Diminishing-Returns Token Selection

Integrate this into a ViT token-pruning module, encoder attention pooling layer, retrieval reranker, or MoE router. For each example, let X={1,...,n} index candidate tokens and let h_i in R^d be their hidden states. Build a set utility f_theta(S) by pooling selected states, for example z_S=sum_{i in S} phi_theta(h_i), f_theta(S)=softplus(rho_theta(z_S, |S|)); subtract a learned or fixed size cost beta|S| if selection must be sparse. During training, sample a base set A0, two or three disjoint increments A1,...,Am, and random subsets A,B from the candidate indices. Evaluate f_theta on all unions required by the inequalities. Add the hinge penalty L_sub above to the task loss. A compact training step is: (1) compute h_1,...,h_n; (2) sample A,B,A0,...,Am; (3) evaluate f_theta on A union B, A intersection B, and every A0 union union_{i in K} Ai; (4) form the two positive-part violations; (5) backpropagate task loss plus lambda_sub L_sub plus lambda_strong L_strong. Estimate no mathematical constants: the inequalities are exact on sampled subsets, while lambda_sub, the maximum order m, and sample count are empirical hyperparameters. At inference, initialize S=emptyset and repeatedly choose i=argmax_{i notin S} Delta_theta(i|S), adding i while the marginal exceeds a threshold or until a budget k is reached. First test on CIFAR-100 or ImageNet-100 with a small DeiT-Tiny: compare standard learned top-k token pruning against this regularized selector at equal kept-token ratios of 25%, 50%, and 75%. Measure validation accuracy, attention FLOPs, selected-token redundancy via pairwise cosine similarity, and the empirical fraction of sampled inequality violations. Success means equal-or-better accuracy at lower token budgets, fewer redundant tokens, and monotone marginal gains during greedy selection.