发表机构
University of Chicago（芝加哥大学）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
本文提出一种用于语言模型top-k选择模块训练的原则性目标，其混合梯度提供丰富信号，算法以O(1/√T)收敛，并在RAG和MoE等任务上优于基线。
AI 中文摘要
从$m$个候选项中选择最佳的$k$项是现代大型语言模型系统的关键组成部分，例如检索增强生成（RAG）中的文档选择和混合专家（MoEs）中的专家路由。然而，由于梯度信号弱和探索-利用权衡次优，训练这些选择模块仍然具有挑战性。此外，先前的工作往往依赖启发式方法，缺乏明确建模和解决top-$k$选择问题的原则性目标和途径。在这项工作中，我们提出了一个训练选择模块的原则性目标，其梯度自然地以混合形式提供更丰富的训练信号——包含监督梯度分量和策略梯度分量。我们表明，随着$m$的增加，选择问题变得更加困难，我们的算法以$O(1/\sqrt{T})$的速率收敛，通过平衡偏差和方差达到最优上界。实际上，我们将我们的方法应用于涉及top-$k$选择的一组任务，包括合成回归问题、RAG和MoE系统，表明我们的方法在下一个词元预测困惑度和问答准确性方面优于基线。
英文摘要
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.
