R2R：使用小型-大型模型令牌路由高效导航分歧推理路径

在人工智能快速发展的今天，大型语言模型（LLMs）展现出令人印象深刻的推理能力，但这些能力往往伴随着巨大的推理开销，给部署带来重大挑战。近日，来自清华大学、无限极AI和上海交通大学的研究团队发表了一篇题为《R2R: Efficiently Navigating Divergent Reasoning Paths with Small-Large Model Token Routing》的研究论文，提出了一种名为"通向罗马之路"（Roads to Rome，简称R2R）的创新方法，通过在小型和大型语言模型之间进行智能令牌路由，显著提升了推理效率。该研究已提交至审稿阶段，代码已在GitHub上开源（https://github.com/thu-nics/R2R）。

虽然经过蒸馏的小型语言模型（SLMs）可以大幅提高生成效率，但它们在尝试跟随大型语言模型的推理路径时常常会出现性能下降。研究团队发现了一个有趣的现象：只有一小部分令牌真正导致LLMs和SLMs之间的推理路径发生分歧。大多数生成的令牌要么完全相同，要么只有中性差异，如缩写或表达方式的细微变化。

想象一下，当你要解一道复杂的数学题时，你可以请一位经验丰富的数学教授（大模型）或一位热心的高中生（小模型）来帮忙。大部分时候，高中生的解题思路与教授相似，只是偶尔在某些关键步骤上会走偏。如果我们能在高中生即将犯错的关键时刻让教授临时指导一下，然后再让高中生继续，这不就既保证了解题质量，又节省了教授的时间吗？

这正是R2R方法的核心思想。研究团队基于这一发现，设计了一种神经令牌路由方法，只在那些关键的、会导致路径分歧的令牌上选择性地使用大型模型，而将大部分令牌生成任务交给小型模型。此外，团队还开发了一个自动数据生成管道，用于识别分歧令牌并生成令牌级路由标签，以训练轻量级路由器。

让我们深入了解这项研究的细节和成果。

一、研究背景与挑战

大型语言模型展示出强大的能力，但它们的推理过程需要大量计算资源。通常，这些模型会生成详细的思维链（Chain-of-Thought，CoT）推理，然后才给出最终答案。然而，这种方法需要拥有数百亿参数的大型模型为每个查询生成数千个令牌，导致推理开销巨大。

经过蒸馏的小型语言模型（SLMs）只有几十亿参数，生成效率高得多。通过对大模型响应进行监督微调，这些小模型可以模仿大模型的推理行为。但问题是，SLMs在推理过程中仍可能产生与大模型不同的推理路径，导致性能严重下降。

研究团队举例说明：与32B参数的R1-32B大模型相比，1.5B参数的R1-1.5B小模型在AIME基准测试中45%的问题上给出了不同的最终答案，准确率下降了4.8倍。

然而，研究团队发现，SLMs和LLMs在给定相同上下文的情况下，往往在下一个令牌预测上达成一致。它们之间的性能差距主要来自于累积错误：在部分响应中出现了一些关键差异后，它们的推理路径逐渐偏离。

通过详细分析，团队发现在32B模型生成的总计760万个令牌中，1.5B模型仅在11%的令牌上产生了不同的预测结果。更重要的是，这些差异中的大部分只是中性变化，例如缩写或表达方式的不同（如"let's"与"let us"），并不影响推理结果。真正导致推理路径分歧的是一小部分特殊令牌，团队将其称为"分歧令牌"（divergent tokens）。这些令牌真正改变了当前句子的含义、逻辑或结论，从而导致后续推理路径的偏离。

这一发现引发了一个关键问题：SLMs是否可以通过仅替换分歧令牌来跟随LLM的推理路径？如果能解决这个问题，我们就能在保持LLM高质量推理的同时，充分利用SLM的效率优势。

二、R2R方法：通向罗马之路

解决SLM-LLM混合推理的主要挑战包括两个相互关联的部分：在特定目标下标记首选模型，以及设计路由方案在推理过程中执行这种选择。

以往的方法通常在查询级别进行路由，为整个响应选择SLM或LLM，以在成本预算内最大化人类偏好的胜率。然而，这些方法依赖于人工标注和复杂的路由器设计，其数据标注和路由方案对于细粒度的令牌级路由来说都太昂贵了。

另一方面，推测性解码方法（speculative decoding）旨在令牌级别上实现SLM和LLM之间的完全相同输出。它们使用SLM草拟输出，并定期用LLM验证。虽然这种方法精确，但严格的验证导致接受率低。此外，中间的差异会使所有后续令牌无效，严重限制了接受长度。

为了解决这些挑战，研究团队提出了"通向罗马之路"（R2R）方法，这是一种令牌级路由方法，仅在SLM生成过程中针对路径分歧令牌选择性地使用LLM。

### 1. 令牌级路由公式化

在自回归语言模型中，推理可以表示为一系列下一个令牌预测的序列。研究团队专注于贪婪采样，通过以下方式定义：

yi = argmax_y P_mi(y|x0, ..., xn-1, y0, ..., yi-1) = argmax_y P_mi(y|S<i)

其中，xi和yi分别表示输入和输出令牌，S<i表示在步骤i的令牌序列。下一个令牌概率P_mi由步骤i处的模型mi预测，mi可以是SLM（θs）或LLM（θl）。

路由策略的本质是定义一个路由函数R，为每个解码步骤选择模型：

mi = R(S<i, θs, θl)

团队的目标是在确保输出序列与仅使用LLM的输出质量相匹配的同时，最小化总生成成本。

### 2. 路径跟随路由策略

最优地解决令牌级路由问题在计算上是难以实现的，特别是对于大规模数据生成。虽然可能存在更好的路由序列（可能与LLM的推理路径不同），但找到它们需要彻底搜索O(2^n)的空间，并为每次搜索生成数千个输出令牌。

为了克服这一实际限制，研究团队提出了一种贪婪的、句子级路径跟随路由策略，将搜索复杂度降低到O(n)。该策略不是探索所有可能的模型选择，而是逐步将混合模型生成与LLM建立的推理路径对齐。

具体来说，在每个生成步骤，策略都会优先选择高效的SLM，除非这会导致与LLM预期推理路径的有意义偏离。为了确定差异是中性的还是分歧的，团队使用了一种延续验证机制：

1. 首先比较SLM和LLM的下一个令牌预测 2. 如果预测相同，则选择SLM 3. 如果预测不同，则构建两个候选序列，分别从SLM和LLM的预测开始 4. 使用LLM继续生成这两个序列，直到满足停止标准（如生成EOS令牌） 5. 如果SLM的延续序列在验证函数V下与LLM的延续序列质量相当，则差异被视为中性的，否则被视为分歧的 6. 根据这一判断选择对应的模型

虽然这种策略可以保证混合生成的令牌序列与其仅LLM的对应物达到相同的质量，但全面延续在计算上对大规模数据生成来说太昂贵了。此外，当前差异对数千个令牌之后的最终输出质量的影响对于要训练的神经路由器来说太难学习了。

因此，在实践中，团队使用句子级路径跟随路由，其中延续在当前句子结束时结束。他们在延续过程中监控句子结束符号（如句号），并使用现有的语义句子分隔器在句子真正结束时结束生成。为了验证这种局部延续，一个强大的LLM作为句子级验证器，用于比较延续并确定初始令牌差异是否引入了与LLM预期推理路径的有意义偏离。

团队验证了句子级路径跟随路由的有效性，使用Qwen2.5-72B作为验证模型。在R1-32B在8K令牌限制内正确解决的17个AIME-24问题中，他们的路径跟随策略实现了可比的准确率（16个问题正确回答），同时仅在3%的生成令牌上依赖较小的R1-1.5B模型。

然而，这种策略直接用于实时推理是不切实际的，因为它依赖于昂贵的LLM延续和验证。相反，团队设计并训练了一个轻量级神经路由器，它可以有效地近似这种策略，仅依靠SLM输出来确定在推理过程中何时使用LLM。

三、令牌级神经路由器

### 1. 训练数据生成

研究团队使用句子级路径跟随路由为神经路由器生成训练标签，并采用了几种优化来控制数据标注开销。

数据生成管道包括以下步骤：

1. 从现有数据集获取查询，获取完整的LLM响应（直接从数据集获取或通过批处理LLM推理） 2. 使用高度并行的SLM预填充，高效地识别SLM预测与LLM相同的令牌，允许排除约90%无需进一步处理的令牌 3. 对剩余10%的不同令牌，从每个SLM预测执行批处理LLM延续 4. 为提高效率，应用前缀缓存重用多个延续之间共享上下文前缀的KV-Cache计算 5. 验证模型比较延续并标记路由偏好

使用这个管道，团队在约2.3天内使用8个A800 GPU高效生成了760万个路由标签，涵盖了Bespoke-Stratos数据集中的数学、编码和问答主题。

### 2. 分歧的预测指标

团队探索了可以帮助识别分歧令牌的预测指标，专注于仅在SLM下一个令牌预测期间可以获取的指标。

SLM的输出logits：数据显示，分歧令牌在SLM的输出logits中表现出明显更高的熵，平均值是非分歧令牌的3.8倍。团队观察到使用其他不确定性度量（如随机不确定性）也有类似趋势。这些实证结果表明，SLM预测中的不确定性增加与令牌分歧强烈相关。

令牌频率：数据集中的低频令牌更可能是分歧的。这可能是由于训练数据中的长尾令牌分布，使得罕见令牌由于SLM的有限容量而更难有效建模。

基于这些见解，路由器将SLM的top-100 logit值作为输入特征之一，并通过使用令牌嵌入作为路由器输入来明确合并令牌频率偏差。

### 3. 路由器设计和路由方案

模型架构：团队将神经路由器设计为一个轻量级的五层前馈网络（FFN），参数为56M。它接收SLM的输出logits和标记化嵌入，以及最后一层隐藏状态以获取额外的语义上下文。所有输入都经过线性投影，连接起来，并送入FFN主干。路由器输出一个二元分类概率，指示当前令牌是否偏离LLM的推理路径。

训练方案：使用交叉熵损失和第4.1节描述的标记数据训练路由器。为了解决由低分歧率引起的类别不平衡，团队对损失进行了反比于类别频率的重新加权。训练后，使用验证集选择满足用户定义的LLM使用率的路由阈值。

路由方案：与定期验证SLM输出的推测性解码方法不同，研究团队的路由方案旨在立即决定是否接受每个SLM令牌，从而消除回滚的需要。如图2所示，这种方法减少了不必要的草稿和验证计算，这在计算密集型批处理服务场景中特别有益。具体来说，神经路由器在每个生成步骤使用SLM输出估计分歧概率。当这个概率超过预定义的阈值pth时，就会调用LLM来纠正当前的输出令牌。

四、实验结果与分析

### 1. 实验设置

基线：研究团队使用DeepSeek-R1-Distill-Qwen模型作为基线，用R1-MB表示，其中M表示模型大小（单位：十亿）。他们将R1-1.5B和R1-32B分别指定为SLM和LLM，中间大小（7B，14B）捕获蒸馏缩放行为。团队比较了RouteLLM框架中的各种查询级路由（QR）方法，包括相似度加权排名（QR-SW）、矩阵分解（QR-MF）、基于BERT的分类器（QR-BERT）和基于Llama3-8B的分类器（QR-LLM）。对于推测性解码，他们采用了EAGLE2和HASS，LLM为R1-32B。

R2R设置：R2R在R1-1.5B和R1-32B之间路由，使用轻量级56M参数FFN路由器，该路由器在第4.1节描述的760万个令牌级路由标签上训练。路由阈值pth在验证集上选择，平均参数使用为6B。性能效率权衡仅通过调整pth控制，无需重新训练路由器。

基准测试：评估方法涵盖具有挑战性的推理基准，包括数学（AIME 2024-2025），研究生级问答（GPQA-Diamond）和编码任务（LiveCodeBench 2024-08-2025-01）。所有实验使用32K令牌的最大输出长度和零生成温度以确保可重复性。

效率指标：研究团队遵循之前的工作，使用每令牌平均激活参数作为与硬件无关的效率指标，简称为平均参数（M）。对于查询级路由，M根据SLM和LLM在所有输出中的激活比率计算其参数的加权平均值。对于R2R，M包括SLM和路由器参数，以及根据LLM激活比率加权的LLM参数。团队还报告了总成本（C），定义为平均激活参数乘以每个查询的平均输出令牌数。平均参数和总成本分别反映了平均解码速度和总延迟。此外，团队还使用SGLang框架在NVIDIA A800-80GB GPU上报告了特定硬件的解码速度。

### 2. 性能结果

缩放行为：图5显示了平均激活参数与准确率的缩放关系。查询级路由（QR）方法从1.5B到32B参数显示出近线性的准确率缩放。经过广泛训练的蒸馏模型（R1-7B，R1-14B）实现了超线性增益，在14B时达到R1-32B准确率的88%，参数大小仅为50%。通过仅将分歧令牌路由到LLM，R2R实现了92%的平均准确率，平均参数仅为17%，达到了更好的缩放效果，形成了新的Pareto前沿。此外，由于输出长度减少，R2R在准确率与总测试时间成本C的权衡方面提供了更好的效果。R2R中的路由阈值还能够在训练后灵活控制这种权衡。

数值比较：表2显示了7B和14B左右平均参数大小的模型性能的数值细节。平均参数大小为5.6B的R2R分别比最佳查询级路由方法（在7B和14B中）提高了1.4-2.4倍和1.2-1.4倍。与蒸馏模型相比，R2R比R1-7B提高了1.4-1.7倍的准确率，甚至在平均准确率上超过了R1-14B 1.1倍。相对于极端情况，R2R实现了比R1-1.5B高4.6倍的准确率，同时保持了R1-32B 92%的准确率，而LLM的使用仅为11-15%。

### 3. 实际速度

表3报告了AIME基准上所有方法的实际延迟和速度。所有基线都使用官方的、高效的SGLang框架，并在两个NVIDIA A800-80GB GPU上使用张量并行性进行评估。R2R使用与表2相同的阈值；查询级路由方法使用14B版本以获得可比的性能。R2R分别比R1-14B和R1-32B实现了1.62倍和2.76倍的生成速度。与查询级路由相比，R2R提供了1.48-1.52倍的加速。它还优于具有树状草稿的高度优化推测性解码方法，这些方法主要在当前的单批处理设置中加速。进一步的系统级优化可以为R2R带来更大的收益。

### 4. 消融研究

从表4第一行的R2R开始，团队通过重新训练具有替代目标或减少输入的路由器来评估设计的有效性，保持所有其他设置不变。所有实验都在AIME基准上进行。

路由目标：如第3节所述，团队将不同的下一个令牌预测分类为中性或分歧。R2R通过容忍中性差异并仅将真正分歧的令牌路由到LLM来提高效率。当路由器被训练为将所有不同的令牌都使用LLM时，它无法在相同数量的LLM使用下达到原始准确率，面临1.4倍的准确率下降，如表4第二行所示。这证实了将LLM使用限制在仅分歧令牌对于在保持高准确率的同时降低成本至关重要。

路由器输入：如第4节所述，SLM logits和令牌嵌入都是分歧的强指标，用作路由器输入。当这些特征逐渐被移除时，路由准确率下降了高达1.3倍，凸显了它们的重要性。值得注意的是，虽然SLM logits理论上可以从路由器内的最后一层隐藏状态计算，但这需要234M参数嵌入层的容量，超过了56M参数神经路由器的容量。

### 5. 路由结果观察

研究团队分析了R2R在AIME基准上的路由行为，考虑了在32K令牌限制内完成的响应。图6(a)显示了跨响应位置的LLM使用率。每个响应被分为思考过程和后续回复，位置归一化为[0, 1]。子图宽度反映了它们各自的平均长度。结果显示，R2R在回复阶段路由到LLM的令牌明显减少。这反映了一种直觉，即在内部思考之后，回复本身是直接的，要求较低。

遵循先前的工作，团队进一步将思考过程基于诸如"Wait"和"Alternatively"等令牌分割为顺序思考。图6(b)检查了每个思考内的LLM使用率。它显示R2R在每个思考的开始和结束时更多地依赖LLM。这与预期一致，即初始令牌为思考设定方向，而结束令牌决定是否结束思考、分支到替代方案或继续更深入的推理。值得注意的是，这些路由模式不是手工制作的，而是自然地从路由器训练中出现的。它帮助R2R有效地分配LLM以获得更好的测试时间缩放。

五、结论与展望

通过这项研究，研究团队提出了R2R，一种令牌级路由方法，通过纠正仅分歧令牌使SLM能够跟随LLM推理路径。他们设计了一种路径跟随路由策略进行有效的数据标注，并研究了预测指标，使神经路由器能够准确选择令牌。在具有挑战性的基准测试中的实验表明，R2R以不到7B的平均参数超越了R1-14B模型，将SLM性能提高了4.6倍，LLM使用率不到15%，并在可比准确率下实现了比LLM快2.8倍的实际速度。

然而，当前的路由策略假设贪婪采样。探索替代采样策略可能进一步增强多功能性。此外，需要更专门的系统级优化来充分实现R2R的理论成本优势。

总的来说，R2R展示了一种有前景的方法，通过智能令牌路由有效地导航分歧推理路径，为大型语言模型的高效部署铺平了道路。