RelayLLM:通过协作解码实现高效推理 RelayLLM: Efficient Reasoning via Collaborative Decoding
让小模型主动调用大模型,用1.07%的token成本恢复60%性能差距
前置知识
GRPO (Group Relative Policy Optimization)
GRPO是一种强化学习算法,用于训练大语言模型。它通过采样一组输出并相对于组平均值计算优势来优化策略。具体来说,对于每个查询采样G个输出,计算每个输出的优势为$A_i = (r_i - \text{mean}(\{r_j\})) / (\text{std}(\{r_j\}) + \epsilon)$,其中$r_i$是奖励。目标函数包含KL散度正则化项防止策略偏离参考模型太远。
本文使用GRPO来训练小模型学会何时调用大模型,理解GRPO的组采样机制和优势计算对于把握训练策略至关重要。
Token-level Collaborative Decoding
令牌级协作解码是指在生成过程中,不同模型在token粒度上进行协作,而不是在查询级别选择模型。小模型生成大部分token,只在关键时刻生成特殊命令n来暂停自身生成并请求大模型生成n个token,然后继续自身生成。这个过程可以多次迭代。
这是本文的核心创新点,与传统query-level路由方法的本质区别。理解这种细粒度协作机制对于把握方法的有效性至关重要。
RLVR (Reinforcement Learning with Verifiable Reward)
RLVR是一种强化学习范式,适用于生成结果可以确定性验证的领域。它使用基于规则的验证器$v: X \to \{0, 1\}$来分配二元奖励。对于数学推理任务,可以通过验证最终答案是否正确来给出确定性的奖励信号,而不是依赖主观评分。
本文在GRPO训练中采用RLVR范式设计奖励机制,理解可验证奖励的特点有助于理解训练的有效性和局限性。
Call Ratio
调用比率是指大模型生成的token占总生成token的比例,计算公式为$\rho(y) = \frac{\text{tokens from ML}}{\text{total tokens}} \in [0, 1]$。这个指标衡量协作推理的计算成本,调用比率越低表示使用的大模型资源越少。在本文中,RelayLLM实现了仅1.07%的调用比率。
这是衡量方法效率的核心指标,理解其计算方式和意义有助于评估本文方法在实际应用中的成本效益。
研究动机
现有的级联或路由方法操作粒度太粗,会在查询级别做出"全或无"的决策,即将整个问题卸载给大模型。这导致严重的计算浪费,因为小模型通常有能力处理大部分推理步骤,只需要在特定关键位置得到专家干预。例如,在数学推理任务中,小模型可能能够处理90%的推理步骤,但传统路由方法会为整个问题调用大模型,造成10倍以上的资源浪费。文章指出,这种粗粒度的"all-or-nothing"策略无法有效平衡性能和成本。
本文的目标是本文的目标是设计一个token-level的协作解码框架,让小模型能够自主决定何时以及在何处调用大模型,只在关键的token位置请求帮助。通过这种方式,既能利用大模型的强大推理能力,又能最大程度地减少计算成本。具体来说,目标是让小模型充当主动控制器,在推理过程中动态地"传递"生成过程给大模型,在获得必要指导后继续自己的推理。
与已有工作不同的是,本文的独特切入角度在于不依赖额外的路由器或控制器,而是让小模型本身具备主动调用大模型的能力。这与传统的token-level路由方法(如CITER需要额外的MLP控制器)有本质区别。更重要的是,本文设计了难度感知的奖励机制,根据查询的难度动态调整奖励策略,这是现有方法所未考虑的。此外,本文深入研究了在无教师设置下的内在推理能力,发现模型能够从协作过程中学习有效的推理模式,这一观察也是独特的。
核心方法
RelayLLM的核心思想是让小模型担任双重角色:既是问题解决者又是中央控制器。在推理过程中,小模型可以生成特殊的命令令牌n来暂停自己的生成,并请求大模型生成n个token。大模型生成完成后,控制权返回小模型,小模型消化专家指导后继续推理。这个过程可以迭代多次。为了训练小模型具备这种策略性委托能力,本文提出了两阶段训练框架:首先是监督预热阶段,教模型学习调用命令的语法结构;然后是强化学习阶段,使用GRPO算法设计难度感知的奖励机制,引导模型平衡独立性和寻求帮助。
核心创新点在于让小模型自身作为主动控制器,而不是依赖外部的路由器。小模型通过生成特殊命令n来触发大模型的干预,其中n是动态预测的调用长度。这与需要额外MLP控制器的CITER方法有本质区别。另一个关键创新是难度感知的奖励设计,将查询分为三种场景:学生可解决、教师依赖、教师不可解决,每种场景采用不同的奖励策略。这种设计避免了简单奖励可能导致的过度依赖或不足依赖问题。
方法步骤详情
RelayLLM的推理过程分为三个步骤:第一步,小模型MS基于当前上下文自回归地生成token,直到生成特殊命令$C_{cmd}(n) = \langle call \rangle \oplus n \oplus \langle /call \rangle$,其中$n \in \mathbb{Z}^+$表示请求的大模型token数量。第二步,检测到命令后,MS的生成暂停,当前上下文(包括查询和MS生成的token)被转发给大模型ML,但会剥离命令token以保持兼容性。ML生成接下来的n个token(或在遇到EOS时提前停止)。第三步,控制权返回MS,上下文用ML生成的新token更新,MS保留完整历史包括自己的命令token,然后继续生成。训练过程分为两个阶段:监督预热阶段从MS自身采样的序列中插入随机位置的命令token,合成训练数据使用交叉熵损失;强化学习阶段使用GRPO算法,采样一组输出计算组平均优势,设计简单奖励和难度感知奖励两种信号。
技术新颖性
技术新颖性体现在三个方面:一是无控制器的token-level协作,小模型自身具备调用能力,避免了额外控制器的计算开销和分布偏移;二是动态长度预测,模型可以根据当前情况决定需要多少token的帮助,相比固定长度策略更加灵活高效;三是难度感知奖励机制,根据查询难度场景分配不同奖励,引导学生模型学习最优策略。实验表明,相比固定100-token策略,RelayLLM在相似精度下将调用比率从2.87%降至1.07%,相比固定20-token策略,精度从39.71%提升至43.75%,证明了动态预测的优势。
实验结果
在六个数学推理基准上的实验表明,RelayLLM实现了卓越的性能-效率权衡。使用Qwen3-1.7B作为小模型,Qwen3-8B作为教师模型时,RelayLLM(难度感知奖励)的平均准确率达到49.52%,相比基线模型42.50%提升了7.02个百分点,接近教师模型54.12%的性能,恢复了约60%的性能差距。更重要的是,这一提升仅用1.07%的调用比率实现,相比性能匹配的随机路由器减少了98.2%的token成本。相比资源等价的随机路由器,RelayLLM还获得了6.9%的准确率提升。在不同模型规模上,RelayLLM都保持优势:Qwen3-0.6B从27.17%提升至33.04%,调用比率仅0.77%。在未见过的推理域(BBEH、MMLU-Pro、SuperGPQA)上,RelayLLM也一致超越基线,如在MMLU-Pro上达到59.03%,大幅超过GRPO基线49.76%。无教师设置下的评估显示,RelayLLM在简单数据集上即使没有教师访问也能超越GRPO基线(61.12% vs 59.51%),表明模型从协作过程中学到了内在的推理能力。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Minerva数学推理 | 准确率 | 43.75% | Qwen3-1.7B Base: 33.82% | +9.93个百分点(相对提升29.4%) |
| MATH-500 | 准确率 | 81.40% | Qwen3-1.7B Base: 74.60% | +6.8个百分点 |
| GSM8K | 准确率 | 86.28% | Qwen3-1.7B Base: 82.64% | +3.64个百分点 |
| 六个基准平均 | 平均准确率 | 49.52% | Qwen3-1.7B Base: 42.50% | +7.02个百分点 |
| 调用效率 | 调用比率 | 1.07% | 性能匹配随机路由器: 约100% | 降低98.93% |
局限与改进
作者承认的方法局限性包括:依赖可验证奖励机制,限制了在非数学领域的应用;训练和推理教师模型需要保持一致,否则性能会下降,这限制了实际部署的灵活性;方法主要在数学推理任务上验证,在开放域对话、长文本生成等任务上的效果未知。此外,我观察到一些额外的局限性:当前实现假设大模型API调用是可靠的,但实际部署中需要处理网络延迟和API故障;命令tokenn可能被误触发或需要多次调用才能获得充分帮助;方法在需要多轮对话或长期上下文推理的任务上可能需要调整;在小模型能力严重不足的情况下,频繁调用可能导致效率优势丧失。
独立分析的弱点
方法的第一个弱点是对数学领域的依赖性太强,因为奖励机制依赖于可验证答案。改进方向可以是设计更通用的语义相似度奖励或使用LLM-as-Judge机制。第二个弱点是训练-推理教师一致性要求,限制了实际部署的灵活性。改进方向可以包括领域自适应技术或教师模型选择策略。第三个弱点是单次调用的局限性,有时需要多次往返才能获得充分帮助。改进方向可以设计多轮调用的优化策略或上下文感知的调用长度预测。第四个弱点是在开放域任务上的泛化性未知,可以在更多样化的任务上进行验证和优化。
未来方向
作者提出的未来方向包括扩展到更多推理域和非数学任务,探索更复杂的协作模式(如多个小模型和一个大模型的协作),以及优化训练效率。基于本文成果,可以延伸的研究方向包括:实时成本优化,根据当前资源约束动态调整调用策略;个性化协作,根据用户偏好或任务特性定制协作模式;多模态扩展,将token-level协作应用到视觉-语言任务;理论分析,建立协作推理的理论框架,分析最优调用策略的性质;以及实际部署优化,研究如何在实际生产环境中处理延迟、容错和并发调用等问题。
复现评估
本文的可复现性良好。代码已在GitHub开源(https://github.com/Chengsong-Huang/RelayLLM),使用了公开的DAPO数据集和标准基准(Minerva、MATH-500、GSM8K、Olympiad-Bench、AIME-2024、AIME-2025)。训练基于EasyR1框架,超参数配置详细:使用AdamW优化器,学习率$1 \times 10^{-6}$,权重衰减$1 \times 10^{-2}$,全局批大小32,最大prompt长度4096,最大响应长度8192,GRPO组大小G=8,KL系数$\beta=0.01$,采样温度1.0。推理使用vLLM引擎。算力需求相对合理,单epoch训练。但实际复现仍需要访问Qwen3系列模型和足够的GPU资源来支持大模型推理。
论文图表