Principled Top-$k$ Selection for Language Models with Hybrid Gradients

Xuchen Gong; Junfei Sun; Tian Li

cs.LGOct 3, 2026

Principled Top- $k$ Selection for Language Models with Hybrid Gradients

Authors: Xuchen Gong, Junfei Sun, Tian Li

Organizations: University of Chicago

Abstract

Selecting the best $k$ items out of $m$ candidates is a critical component of modern large language model systems, such as document selection in Retrieval-Augmented Generation (RAG) and expert routing in Mixture-of-Experts (MoEs). However, training these selection modules remains challenging due to weak gradient signals and suboptimal exploration-exploitation tradeoffs. Furthermore, prior works often rely on heuristics, lacking principled objectives and approaches that explicitly model and solve the top- $k$ selection problem. In this work, we propose a principled objective for training selection modules, whose gradient naturally provides richer training signals in a hybrid form---containing a supervised-gradient component and a policy-gradient component. We show that the selection problem becomes harder as $m$ increases, and our algorithm converges at rate $O(1/\sqrt{T})$ , with the optimal upper bound achieved by balancing between bias and variance. Practically, we apply our method to a set of tasks involving top- $k$ selection, including synthetic regression problems, RAG, and MoE systems, showing that our method outperforms the baselines in next-token prediction perplexity and QA accuracy.

Figures & tables

Figure 1: Our method outperforms prior ones that abandon or approximate the policy gradient.

Figure 2: Our framework (left) and its applications to RAG and MoE. Middle: In RAG, given input query $x$ , the reranker selects a subset of the most relevant documents to augment $x$ to predict the target $y$ . RAG has one selection and prediction module in total, with the prediction module (i.e., LM) frozen in our setting. Right: In each MoE layer, given the input hidden state $x$ , the router selects a subset of expert models to compute the output hidden state $y$ . The full MoE model contains multiple MoE layers of selection and aggregation, with both the selection and prediction modules trainable.

Figure 5: Performance of retrieval-augmented Qwen-2.5-1.5B and Llama-3.2-3B with the reranker trained via different methods. We observe improved performance as

m

increases, and our method outperforms all the baselines. The detailed metric values are reported in Tables 2 - 6 .

Figure 6: Test perplexity of different methods after training from scratch on Wikitext-103. Left: Our method outperforms baselines under different numbers of total experts. Right: The best bias/variance tradeoff is achieved by adjusting

\lambda

. The exact numerical values are reported in Table 7 .

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Gradient norm and loss standard deviation of the supervised component and the policy component. Under appropriate

\lambda

, the two metrics for the two components are comparable.

Method	MSE Loss ( $\times$ 1e-3, $\downarrow$ )				Top-1 Selection Accuracy ( $\%,\uparrow$ )
Method	$m=64$	$m=128$	$m=256$	$m=512$	$m=64$	$m=128$	$m=256$	$m=512$
Deterministic	1.01	2.02	5.73	5.99	89.0	66.7	4.0	0.6
Gaussian	1.39	2.15	5.14	6.22	80.8	64.5	14.3	1.0
DefaultMoE	1.20	1.44	1.87	3.29	80.1	68.6	53.5	41.4
SparseMixer	1.39	1.92	3.11	4.30	86.6	73.9	46.8	27.0
Supervised-only	0.99	1.43	1.91	4.67	89.9	80.0	61.2	11.9

Appendix

Table 1: Performance on the regression task.

Figure 8: Decrease of test loss and the evolution of selection decisions (on two test samples) during optimization from

\theta_{0}

\theta_{T}

under different methods. The line chart represents loss decrease, and each background color represents one selection decision on two test samples. We observe that the selection module explores more selection decisions with our method.

Figure 9: Bias/variance tradeoff adjusted by

\lambda

, showing that the optimal tradeoff is obtained at

\lambda\in(0,1)

under different

m

’s.

Figure 10: Document recall of rerankers trained via different methods.

Method	Qwen-2.5-1.5B			Llama-3.2-3B
Method	$m=8$	$m=16$	$m=32$	$m=8$	$m=16$	$m=32$
Frozen	9.41	10.52	11.14	7.70	8.60	9.11
Policy-only	7.45	7.59	7.89	6.70	7.12	7.46
ICRALM	6.61	6.99	7.64	5.37	5.56	6.09
REPLUG	6.54	6.84	7.61	5.33	5.59	5.98
Supervised-only	6.52	6.79	7.26	5.31	5.56	5.89

Appendix

Table 2: Perplexity on Wikitext-103 (

\downarrow

Method	Qwen-2.5-1.5B			Llama-3.2-3B
Method	$m=8$	$m=16$	$m=32$	$m=8$	$m=16$	$m=32$
Frozen	23.7	26.2	26.5	28.3	30.4	31.4
Policy-only	25.7	29.7	30.8	29.8	31.9	33.2
ICRALM	27.5	29.7	30.8	30.8	34.1	35.0
REPLUG	26.9	29.7	31.1	31.2	33.5	34.5
Supervised-only	27.7	31.3	32.4	31.8	34.6	35.4

Appendix

Table 3: Exact match accuracy on NQ (

\uparrow

Method	Qwen-2.5-1.5B			Llama-3.2-3B
Method	$m=8$	$m=16$	$m=32$	$m=8$	$m=16$	$m=32$
Frozen	63.1	68.6	74.4	63.1	68.6	74.4
Policy-only	56.2	62.7	67.5	55.3	63.1	67.3
ICRALM	65.5	74.6	81.5	65.3	73.4	79.0
REPLUG	64.9	73.3	79.2	65.0	73.3	78.1
Supervised-only	65.9	75.3	82.8	65.8	74.2	81.8

Appendix

Table 4: Document recall on NQ (

\uparrow

Method	Qwen-2.5-1.5B			Llama-3.2-3B
Method	$m=8$	$m=16$	$m=32$	$m=8$	$m=16$	$m=32$
Frozen	55.7	58.3	59.3	71.4	72.0	73.0
Policy-only	56.8	58.3	61.0	70.4	71.2	72.8
ICRALM	60.3	61.1	61.9	72.1	74.5	74.6
REPLUG	60.3	61.6	63.3	72.7	74.6	75.2
Supervised-only	61.2	62.6	65.2	73.4	74.9	75.7

Appendix

Table 5: Exact match accuracy on TriviaQ (

\uparrow

Method	Qwen-2.5-1.5B			Llama-3.2-3B
Method	$m=8$	$m=16$	$m=32$	$m=8$	$m=16$	$m=32$
Frozen	84.9	88.5	90.0	84.9	88.5	90.0
Policy-only	80.5	83.9	86.6	80.1	83.9	85.9
ICRALM	86.3	88.9	90.6	86.7	90.1	91.5
REPLUG	86.5	89.5	91.1	86.4	90.1	91.6
Supervised-only	86.8	90.3	92.6	86.8	90.8	91.8

Appendix

Table 6: Document recall on TriviaQ (

\uparrow

Figure 11: Bias/variance tradeoff of Qwen-2.5-1.5B adjusted by

\lambda

, whose optimal tradeoff is obtained at

\lambda\in(0,1)

Method	One Layer					Two Layers
Method	$m=8$	$m=16$	$m=32$	$m=64$	$m=128$	$m=8$	$m=16$	$m=32$	$m=64$
Deterministic	86.33	84.44	86.66	86.06	85.82	71.07	69.45	68.56	69.21
Gaussian	85.29	83.11	82.75	81.90	83.26	69.47	67.17	68.06	66.14
DefaultMoE	87.52	89.42	87.37	86.99	84.94	70.80	70.59	69.65	67.76
SparseMixer	87.87	86.00	86.89	86.19	86.05	71.37	68.74	68.70	69.95
Supervised-only	84.80	80.99	77.97	74.73	74.39	69.02	64.94	62.72	59.67

Appendix

Table 7: Perplexity on Wikitext-103 by MoE (

\downarrow

Figure 12: Bias/variance tradeoff adjusted by

\lambda

, showing that the optimal tradeoff is obtained at

\lambda\in(0,1)

under different

m

’s. The upper row shows the results for 1 MoE layer, and the lower row shows the results for 2 MoE layers.

Figure 13: Convergence of difference methods for regression, MoE, and RAG applications.

LESSER: Post-Training Data Selection with Output-Layer Gradients
Oct 2, 2026Lyuxin David Zhang, Eric Wong, Surbhi Goel +1Training Data Selection Data Selection for Language Model Fine-Tuning
University of Pennsylvania · University of Texas at Austin
Prof-K: Probabilistic One-Pass Filtering for Efficient Top-k Selection
Aug 12, 2026Tadeusz Dziarmaga, Witold Sikora, Łukasz Struski +2Streaming Algorithms
Jagiellonian University
HARP: Efficient Data Selection for Finetuning Large Language Models
Jun 5, 2026Ning Wang, Zhengxin Zhang, Maosen Tang +3LLM Fine-Tuning Active Learning
Cornell University · The Hong Kong University of Science and Technology

Oct 2, 2026J/K move · Enter open · S save

LESSER: Post-Training Data Selection with Output-Layer Gradients

Lyuxin David Zhang, Eric Wong, Surbhi Goel, Anton Xue

University of Pennsylvania · University of Texas at Austin

Training Data Selection Data Selection for Language Model Fine-Tuning LLM Post-Training Language Model Post-Training Instruction-Tuning Data Selection

Open paper AskPDF arXiv

The choice of post-training data for large language models substantially affects downstream performance. Gradient-based data selection is a popular approach that ranks training data by how well their gradients align with those of a small validation set. However, ranking with full-parameter gradients requires an expensive backward pass on every sample, making computation intractable for large candidate pools. This raises a natural question: can we approximate full-gradient features at a fraction of the cost? Conveniently, we find that output-layer gradients suffice for effective data selection, yet require only the cheaper forward pass. We implement this as LESSER, a drop-in wrapper for selection methods that reduces the feature-extraction FLOP cost by

9.7\times

for SFT and

3.0\times

for RL benchmarks, while tracking full-gradient performance on downstream tasks. Empirically, we find that even when output-layer and full gradients rank individual samples differently, they select batches with aligned gradients.

Abstract

Explore similar work

LESSER: Post-Training Data Selection with Output-Layer Gradients

Prof-K: Probabilistic One-Pass Filtering for Efficient Top-k Selection

HARP: Efficient Data Selection for Finetuning Large Language Models