摘要：大型语言模型（LLMs）以巨大的推理开销为代价，实现了令人印象深刻的推理能力，这带来了巨大的部署挑战。 尽管蒸馏的小语言模型（SLM）显著提高了效率，但由于它们无法遵循LLM的推理路径，因此性能受到影响。 幸运的是，我们发现只有一小部分标记真正在LLM和SLM之间发散推理路径。 大多数生成的代币要么完全相同，要么表现出中性的差异，例如缩写或表达上的微小变化。 基于这一见解，我们引入了罗马之路（R2R），这是一种神经令牌路由方法，有选择地仅针对这些关键的、路径发散的令牌使用LLM，同时将大多数令牌生成留给SLM。 我们还开发了一个自动数据生成管道，该管道可以识别不同的标记并生成标记级路由标签来训练轻量级路由器。 我们应用R2R将DeepSeek家族的R1-1.5B和R1-32B模型结合起来，并在具有挑战性的数学、编码和问答基准上进行评估。 R2R的平均激活参数大小为5.6B，其准确率是R1-7B的1.6倍，甚至超过了R1-14B模型。 与R1-32B相比，它在性能相当的情况下实现了2.8倍的时钟加速，在测试时间缩放效率方面达到了帕累托前沿。 我们的代码可以在Github。Huggingface链接：Paper page，论文链接：2505.21600

研究背景和目的

研究背景

近年来，大型语言模型（LLMs）在各种任务中展现出了强大的能力，特别是在复杂的推理任务上。然而，这些强大的能力是以巨大的推理开销为代价的，这给模型的部署带来了显著的挑战。为了应对这一挑战，研究者们开始探索如何通过蒸馏技术创建小语言模型（SLMs），这些SLMs参数较少，生成效率更高，能够模拟LLMs的推理行为。然而，尽管SLMs在效率上有了显著提升，但它们在推理路径上往往与LLMs存在分歧，导致性能下降。例如，在AIME基准测试中，R1-1.5B SLM与R1-32B LLM相比，在45%的问题上提供了不同的最终答案，准确率下降了4.8倍。

进一步的研究发现，SLMs和LLMs在生成令牌时，大多数情况下预测是相同的，性能差距主要源于累积误差：在部分响应出现关键差异后，它们的推理路径开始逐渐偏离。这表明，只有一小部分令牌真正导致了推理路径的分歧。因此，一个自然的研究问题是：SLMs是否能够通过仅替换发散的令牌来跟随LLMs的推理路径？ 如果能够解决这个问题，我们就可以在大多数生成步骤中利用SLMs的高效性，同时保留LLMs高质量推理的优势，从而在测试时间缩放效率上实现更好的性能-效率权衡。

研究目的

本研究旨在开发一种名为罗马之路（R2R）的神经令牌路由方法，通过选择性地在SLMs生成过程中使用LLMs来纠正那些关键、路径发散的令牌，从而在保持高质量推理的同时提高生成效率。具体来说，我们的研究目的包括：

1. 设计一种自动数据生成管道：用于识别发散令牌，并生成令牌级路由标签，以训练轻量级路由器。
2. 开发一种神经令牌路由器：在SLMs生成过程中，根据路由器的预测，选择性地将发散令牌路由到LLMs进行纠正。
3. 在具有挑战性的数学、编码和QA基准上评估R2R的性能：验证其在提高生成效率的同时，能否保持或提升推理质量。

研究方法

1. 数据收集与预处理

为了训练神经令牌路由器，我们需要大量的路由标签数据。这些数据包括LLMs和SLMs在相同上下文中生成的令牌序列，以及每个令牌是否应该被路由到LLMs的标签。我们通过以下步骤生成这些数据：

• 生成LLMs响应：使用LLMs（如R1-32B）为给定的问题生成响应，作为“黄金”推理路径。
• SLMs预填充：使用SLMs（如R1-1.5B）对相同的问题进行预填充，生成部分响应序列。
• 令牌差异识别：比较LLMs和SLMs生成的令牌序列，识别出不同的令牌。
• 继续生成与验证：对于每个不同的令牌，我们分别从LLMs和SLMs的预测继续生成完整的响应序列，并使用另一个LLMs作为验证器来判断这些继续生成的序列是否在语义上等价。如果不等价，则标记该令牌为发散令牌。

2. 神经令牌路由器设计

我们设计了一个轻量级的五层前馈网络（FFN）作为神经令牌路由器。该路由器接受SLMs的输出日志、令牌嵌入和最后层的隐藏状态作为输入，输出一个二元分类概率，指示当前令牌是否偏离了LLMs的推理路径。路由器的训练目标是最小化交叉熵损失，通过反向传播算法更新路由器的参数。

3. 推理过程中的路由策略

在推理过程中，我们采用了一种即时路由策略。对于每个生成的令牌，SLMs首先进行预测，然后路由器根据SLMs的输出预测该令牌是否应该被路由到LLMs。如果路由器的预测概率超过预设的阈值，则使用LLMs来纠正该令牌；否则，使用SLMs的预测。

研究结果

1. 性能提升

我们在AIME、GPQA-Diamond和LiveCodeBench等具有挑战性的基准上评估了R2R的性能。实验结果表明，R2R在平均激活参数大小为5.6B的情况下，准确率比R1-7B高出1.6倍，甚至超过了R1-14B模型。与R1-32B相比，R2R在性能相当的情况下实现了2.8倍的时钟加速，显著推进了测试时间缩放效率的帕累托前沿。

2. 路由效率

通过分析路由器的行为，我们发现R2R在回复阶段路由到LLMs的令牌明显减少，这符合直觉，因为经过内部思考后，回复本身更加直接，要求较低。此外，R2R在每个思考过程的开始和结束时更依赖于LLMs，这反映了初始令牌设定思考方向，而结束令牌决定是否结束思考、分支到替代方案或继续深入推理的直觉。

3. 消融研究

我们通过消融研究验证了R2R设计的有效性。实验结果表明，将路由目标更改为所有不同的令牌都会导致准确性显著下降（1.4倍），这证实了仅将发散令牌路由到LLMs对于在减少成本的同时保持高准确性至关重要。此外，去除SLMs日志或令牌嵌入等输入特征也会导致路由准确性大幅下降，进一步强调了这些特征在路由器中的重要性。

研究局限

尽管R2R在提高生成效率的同时保持了高质量的推理，但本研究仍存在一些局限性：

1. 采样策略限制：当前的路由策略假设了贪心采样。探索替代采样策略可能会进一步提高R2R的通用性。
2. 系统级优化需求：为了充分实现R2R的理论成本效益，需要进行更多的专用系统级优化。
3. 数据集多样性：尽管我们在多个基准上评估了R2R的性能，但这些基准主要涵盖了数学、编码和QA任务。未来研究可以探索R2R在其他类型任务上的表现。

未来研究方向

基于本研究的结果和局限性，未来的研究可以关注以下几个方向：

1. 探索替代采样策略：除了贪心采样外，还可以探索其他采样策略（如束搜索、top-k采样等）对R2R性能的影响。
2. 系统级优化：开发更高效的系统级实现，以充分利用R2R的潜力。例如，可以通过优化KV缓存更新、减少LLMs和SLMs之间的通信开销等方式来提高推理速度。
3. 扩展到其他任务类型：将R2R应用于更多类型的任务（如自然语言生成、对话系统等），以验证其通用性和有效性。
4. 结合其他模型压缩技术：探索将R2R与其他模型压缩技术（如量化、剪枝等）相结合，以进一步提高模型的效率和性能。
5. 可解释性研究：研究R2R的路由决策背后的原因，提高其可解释性，从而帮助用户更好地理解和信任模型。

总之，本研究通过引入R2R方法，为在保持高质量推理的同时提高大型语言模型的生成效率提供了一种新的思路。未来的研究可以进一步探索R2R的潜力，并推动其在更多实际应用中的部署。