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

arXiv: https://arxiv.org/abs/2610.04162 · HTML: https://mlsanity.com/paper/2610.04162/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Generation and context integration, one of 15 directions in Retrieval-augmented generation.
- Builds on 43 of its 56 references that are in the map.
- First posted on arXiv 3 Oct 2026.

## Placed in

- Retrieval-augmented generation / [Generation and context integration](https://mlsanity.com/rag/generation-and-context-integration/)
- Mixture-of-experts / [MoE large language models](https://mlsanity.com/moe/moe-large-language-models/)

## Builds on (most cited first, 30 of 43)

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020)
- [Mixtral of Experts](https://mlsanity.com/paper/2401.04088/) (2024)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [Leveraging Passage Retrieval with Generative Models for Open Domain Question Answering](https://mlsanity.com/paper/2007.01282/) (2020)
- [Atlas: Few-shot Learning with Retrieval Augmented Language Models](https://mlsanity.com/paper/2208.03299/) (2022)
- [Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learning](https://mlsanity.com/paper/1702.03118/) (2017)
- [DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models](https://mlsanity.com/paper/2401.06066/) (2024)
- [REPLUG: Retrieval-Augmented Black-Box Language Models](https://mlsanity.com/paper/2301.12652/) (2023)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [In-Context Retrieval-Augmented Language Models](https://mlsanity.com/paper/2302.00083/) (2023)
- [Passage Re-ranking with BERT](https://mlsanity.com/paper/1901.04085/) (2019)
- [Mixture-of-Experts with Expert Choice Routing](https://mlsanity.com/paper/2202.09368/) (2022)
- [BASE Layers: Simplifying Training of Large, Sparse Models](https://mlsanity.com/paper/2103.16716/) (2021)
- [Large Language Models are Effective Text Rankers with Pairwise Ranking Prompting](https://mlsanity.com/paper/2306.17563/) (2023)
- [Unified Scaling Laws for Routed Language Models](https://mlsanity.com/paper/2202.01169/) (2022)
- [Query Rewriting for Retrieval-Augmented Large Language Models](https://mlsanity.com/paper/2305.14283/) (2023)
- [Stochastic Beams and Where to Find Them: The Gumbel-Top-k Trick for Sampling Sequences Without Replacement](https://mlsanity.com/paper/1903.06059/) (2019)
- [Multi-Stage Document Ranking with BERT](https://mlsanity.com/paper/1910.14424/) (2019)
- [Conditional Computation in Neural Networks for faster models](https://mlsanity.com/paper/1511.06297/) (2015)
- [Routing Networks: Adaptive Selection of Non-linear Functions for Multi-Task Learning](https://mlsanity.com/paper/1711.01239/) (2017)
- [DSelect-k: Differentiable Selection in the Mixture of Experts with Applications to Multi-Task Learning](https://mlsanity.com/paper/2106.03760/) (2021)
- [RankT5: Fine-Tuning T5 for Text Ranking with Ranking Losses](https://mlsanity.com/paper/2210.10634/) (2022)

## Abstract (arXiv)

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/√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.