R-Stitch: Dynamic Trajectory Stitching for Efficient Reasoning
作者: Zhuokun Chen, Zeren Chen, Jiahao He, Lu Sheng, Mingkui Tan, Jianfei Cai, Bohan Zhuang
分类: cs.LG, cs.AI, cs.CL
发布日期: 2025-07-23 (更新: 2025-09-26)
💡 一句话要点
提出R-Stitch以解决大语言模型推理效率问题
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 动态推理 语言模型 熵引导 计算效率 混合解码
📋 核心要点
- 现有方法在推理过程中面临高成本,尤其是在长自回归轨迹下,导致效率低下。
- R-Stitch通过熵引导的路由策略,动态分配计算资源于小模型和大模型之间,从而提高推理效率。
- 实验结果显示,R-Stitch在多个模型上实现了高达4.10倍的加速,同时保持与完整LLM解码相当的准确性。
📝 摘要(中文)
链式思维(CoT)增强了大型语言模型(LLMs)的问题解决能力,但由于长自回归轨迹,推理成本显著增加。现有加速策略通过提前停止或压缩缩短轨迹,或采用小模型的推测解码。然而,当模型一致性较低时,推测解码的收益有限。我们提出R-Stitch,这是一种无训练的混合解码框架,利用令牌级熵作为不确定性代理,在小语言模型(SLM)和LLM之间分配计算。我们的分析表明,高熵令牌更可能导致错误,从而激励了一种熵引导的路由策略,使SLM有效处理低熵令牌,同时将不确定的令牌委托给LLM。通过减少每个令牌的解码复杂度和生成的令牌数量,我们的方法实现了显著加速,且准确性损失微乎其微。
🔬 方法详解
问题定义:本论文旨在解决大型语言模型推理过程中的高成本问题,现有方法通过提前停止或压缩来加速,但效果有限,尤其在模型一致性低时。
核心思路:R-Stitch提出了一种基于熵的动态路由策略,允许小语言模型处理低熵令牌,而将高不确定性的令牌交给大型语言模型,从而提高推理效率。
技术框架:R-Stitch的整体架构包括两个主要模块:小语言模型(SLM)和大型语言模型(LLM)。在推理过程中,系统首先计算令牌的熵值,根据熵值决定将令牌分配给哪个模型进行处理。
关键创新:R-Stitch的主要创新在于引入熵作为不确定性代理,动态调整计算资源的分配,这与现有方法的固定策略形成鲜明对比。
关键设计:在设计中,熵值的计算方式、路由策略的实现以及动态调整的机制都是关键细节,确保了在不同计算预算下的适应性和效率。具体的参数设置和损失函数设计也经过精心调整,以优化模型性能。
🖼️ 关键图片
📊 实验亮点
在多个基准测试中,R-Stitch实现了显著的性能提升:在DeepSeek-R1-Distill-Qwen-7B上达到了3.00倍的加速,在14B模型上为3.85倍,而在QWQ-32B模型上则达到了4.10倍,同时保持了与完整LLM解码相当的准确性。这些结果表明,R-Stitch在推理效率和准确性之间实现了良好的平衡。
🎯 应用场景
R-Stitch的研究成果具有广泛的应用潜力,尤其在需要高效推理的自然语言处理任务中,如对话系统、文本生成和信息检索等领域。其动态资源分配的特性使得在资源受限的环境中仍能保持良好的性能,未来可能推动更多智能应用的发展。
📄 摘要(原文)
Chain-of-thought (CoT) enhances the problem-solving ability of large language models (LLMs) but incurs substantial inference cost due to long autoregressive trajectories. Existing acceleration strategies either shorten traces via early stopping or compression, or adopt speculative decoding with a smaller model. However, speculative decoding provides limited gains when model agreement is low and rigidly enforces token-level consistency, overlooking the observation that some smaller models, when correct, produce significantly more concise reasoning traces that could reduce inference length. We introduce R-Stitch, a training-free hybrid decoding framework that leverages token-level entropy as an uncertainty proxy to delegate computation between a small language model (SLM) and an LLM. Our analysis shows that high-entropy tokens are more likely to induce errors, motivating an entropy-guided routing strategy that lets the SLM efficiently handle low-entropy tokens while delegating uncertain ones to the LLM, thereby avoiding full rollbacks and preserving answer quality. We further extend this design with R-Stitch$^{+}$, which learns an adaptive routing policy to adjust the token budget dynamically beyond fixed thresholds. By jointly reducing per-token decoding complexity and the number of generated tokens, our method achieves substantial acceleration with negligible accuracy loss. Concretely, it attains peak speedups of 3.00$\times$ on DeepSeek-R1-Distill-Qwen-7B, 3.85$\times$ on 14B, and 4.10$\times$ on QWQ-32B while maintaining accuracy comparable to full LLM decoding. Moreover, it naturally enables adaptive efficiency--accuracy trade-offs that can be tailored to diverse computational budgets without retraining.