使用混合梯度的语言模型的原则性 Top-$k$ 选择

摘要

从 $m$ 候选中选择最佳的 $k$ 项目是现代大语言模型系统的关键组成部分，例如检索增强生成 (RAG) 中的文档选择和专家混合 (MoE) 中的专家路由。然而，由于梯度信号较弱和探索-利用权衡不理想，训练这些选择模块仍然具有挑战性。此外，先前的工作通常依赖于启发式方法，缺乏明确建模和解决 top-$k$ 选择问题的原则性目标和方法。在这项工作中，我们提出了训练选择模块的原则目标，其梯度自然地以混合形式提供更丰富的训练信号——包含监督梯度组件和策略-梯度组件。我们表明，随着 $m$ 的增加，选择问题变得更加困难，并且我们的算法以 $O(1/\sqrt{T})$ 的速率收敛，通过平衡偏差和方差实现最佳上限。实际上，我们将我们的方法应用于涉及 top-$k$ 选择的一组任务，包括综合回归问题、RAG 和 MoE 系统，表明我们的方法在 next-token预测困惑度和 QA 准确性方面优于基线。

使用混合梯度的语言模型的原则性 Top-$k$ 选择的原论文方法或结果图 — 图 2：我们的框架（左）及其在 RAG 和 MoE 中的应用。中：在 RAG 中，给定输入查询 $x$，重排序器选择最相关文档的子集来增强 $x$，以预测目标 $y$。 RAG 总共有一个选择和预测模块，预测模块（即 LM）在我们的设置中被冻结。右：在每个 MoE 层中，给定输入隐藏状态 $x$，路由器选择专家模型的子集来计算输出隐藏状态 $y$。完整的 MoE 模型包含多个选择和聚合的 MoE 层，以及选择和预测模块可训练。