Skip to content
Go back 2504.19095 arXiv logo

Efficient Reasoning for LLMs through Speculative Chain-of-Thought

Published:  at  01:19 AM
79.97 🤔

本文提出了推测思维链(SCoT)框架,通过轻量级草稿模型并行生成多个思维链草稿,并由微调后的目标大模型选择最佳草稿或决定重新思考,从而在保持接近大模型准确率的同时,显著降低了大型语言模型的推理延迟。

Large Language Model, Reasoning, Efficiency, Fine-tuning, Human-AI Interaction

Jikai Wang, Juntao Li, Lijun Wu, Min Zhang

Soochow University, Shanghai AI Laboratory

Generated by gemini-2.5-flash-preview-04-17

Background Problem

大型推理语言模型(如 OpenAI-o1、Deepseek-R1)在各种推理任务中表现出色,这得益于其庞大的模型规模和冗长的思维链(CoT)生成。然而,这导致了巨大的计算开销和显著的响应延迟。现有的高效推理方法主要集中于开发更小的模型或缩短 CoT 长度。但仅使用小型模型或压缩 CoT 可能会牺牲在复杂任务上的性能,而对简单任务使用大型模型则是一种资源浪费。如何在保证复杂任务性能的同时,提高整体推理效率,尤其是在处理不同难度问题时自适应地利用模型能力,是一个关键问题。

Method

Speculative Chain-of-Thought (SCoT) 是一种通过大小模型协同工作来加速 LLM 推理的框架。其核心方法包括:

  1. 草稿生成 (Generating Drafts of Thoughts): 使用一个轻量级的小模型 (MdM_d) 并行生成 nn 个思维链(CoT)草稿 (Td1,ext...,TdnT_d^1, ext{...}, T_d^n)。为了增加草稿的多样性,这里采用了核采样(nuclear sampling)。由于小模型参数量少,其生成速度远快于大模型。
  2. 思维行为对齐 (Thinking Behavior Alignment): 为了解决小模型生成的 CoT 可能过于冗长或包含冗余的问题,通过 LoRA 微调小模型,使其生成的 CoT 在长度和风格上更接近目标大模型生成的 CoT。训练数据是目标模型在 GSM8K 训练集上生成的问答对和 CoT。微调使用交叉熵损失 LDraf=1li=1llogpMd(yixm,y<i)\mathcal{L}_{Draf} = -\frac{1}{l} \sum_{i=1}^{l} \log p_{M_d}(y_i | x_{\le m}, y_{< i})
  3. 草稿选择与错误纠正 (Draft Selection and Error Correction): 使用一个微调后的目标大模型 (MM) 来评估生成的 nn 个草稿。设计一个特殊的 Prompt 模板 SS,将问题和所有草稿(包括一个表示“所有草稿都错误”的特殊选项 Tn+1T_{n+1})作为输入。目标模型通过一次前向传播,选择概率最高的 CoT 索引:index=argmaxiVPM(iS(q,Tn+1)),  V={1,2,,n+1}index = \arg\max_{i \in \mathcal{V}} \, P_M(i | S(q, T_{\le n+1})), \; \mathcal{V} = \{1, 2, \ldots, n+1\}。目标模型也通过 LoRA 进行微调,以提高其选择正确草稿和检测错误的能力。微调损失函数为 LTarget=min{logpM(yS(q,Tn+1))yY}\mathcal{L}_{Target} = \min \{ -\log p_M(y|S(q, T_{\le n+1})) | y \in \mathcal{Y} \},其中 Y\mathcal{Y} 是所有正确草稿索引的集合。如果选择的索引 indexnindex \le n,则直接使用选定的草稿 TindexT_{index} 和目标模型生成最终答案。如果选择索引为 n+1n+1,则表示草稿模型未能生成可用草稿,此时目标模型将重新进行完整的推理过程以生成答案,确保复杂问题的准确性。这个机制使得系统能够根据问题难度和草稿质量自适应地分配计算资源。

Experiment

实验采用了 Deepseek-R1-Distill-Qwen-32B 作为目标模型,Deepseek-R1-Distill-Qwen-1.5B 作为草稿模型。评估数据集涵盖了不同难度的推理任务,包括 GSM8K、MATH、GaoKao、CollegeMath 和 Olympiad。实验环境为单张 A100 GPU (1.5B) 和四张 A100 GPU (32B)。模型通过 LoRA (rank 8) 在 GSM8K 训练集上进行微调,其中 1500 个样本用于草稿模型的思维行为对齐,500 个样本用于目标模型的草稿选择和错误检测。草稿数量设置为 5 个,最大草稿长度 5000。

实验结果显示,SCoT 在保持接近目标模型准确率的同时,显著降低了推理延迟,速度提升比达到 1.90× 至 2.92×。在 CollegeMath 数据集上甚至实现了无损加速。消融实验证明了多草稿生成、目标模型微调进行选择以及错误纠正机制(即允许目标模型重新思考)对提高准确率的有效性。思维行为对齐显著缩短了草稿模型的 CoT 长度,提高了草稿生成效率。草稿选择准确率分析表明,微调后的目标模型能更准确地识别正确草稿或判断所有草稿均错误的情况,尤其是在草稿模型能生成正确链的情况下准确率较高(85.1%),即使在草稿全错的困难样本上(Class 2),准确率也从原始模型的 0% 提升到 52.0%。此外,SCoT 还减少了最终答案的生成长度,进一步提升了推理效率。

Further Thoughts

SCoT 将推测解码(Speculative Decoding)的思想从 token 级别扩展到 thought(CoT)级别,这是一个非常自然的且富有成效的迁移。这种大模型与小模型协同工作的模式,不仅限于推理加速,还可以推广到其他需要权衡效率和质量的任务中。例如,在内容生成、代码生成等领域,可以先用小模型快速生成多个草稿,再用大模型进行筛选、润色或纠错。此外,论文中提到的“思维行为对齐”非常关键,它解决了小模型和大模型在生成风格和冗余度上的差异,确保了草稿的可用性。未来的工作可以进一步探索更精细的对齐方法,或者研究如何让小模型生成的草稿更接近大模型的“最优”思维路径。动态调整草稿数量或根据问题初步判断难度来决定是否生成草稿,也可能是进一步提升效率的方向。这种自适应、分层处理任务的思路,与人类解决问题的过程有相似之处,为构建更智能、更高效的 AI 系统提供了新的视角。



Previous Post
AdaptMI: Adaptive Skill-based In-context Math Instruction for Small Language Models
Next Post
StreamRL: Scalable, Heterogeneous, and Elastic RL for LLMs with Disaggregated Stream Generation