R1-Compress通过块级压缩和块间搜索机制有效压缩长链式推理(Long-CoT),在减少约20% token使用量的同时保持了与基线接近的推理准确率(92.4% vs 93.0%)。
Large Language Model, Reasoning, Efficiency, Supervised Learning, Prompt Engineering
Yibo Wang, Li Shen, Huanjin Yao, Tiansheng Huang, Rui Liu, Naiqiang Tan, Jiaxing Huang, Kai Zhang, Dacheng Tao
Tsinghua University, Sun Yat-sen University, Didichuxing Co. Ltd, Nanyang Technological University
Generated by grok-3
Background Problem
长链式推理(Long-CoT)是大型语言模型(LLMs)在复杂推理任务(如数学问题求解)中通过逐步分解问题提升性能的重要技术。然而,Long-CoT 显著增加了token长度,导致推理速度变慢和KV缓存内存使用量激增,限制了其实际部署能力。现有压缩方法(如实例级和token级压缩)要么牺牲了关键的局部推理信号(如反思能力),要么生成不连贯的输出。因此,亟需一种既能减少计算开销又能保留推理质量的压缩方法。
Method
R1-Compress 提出了一种两阶段的块级压缩框架,旨在通过分块处理和搜索机制平衡效率与连贯性:
- 核心思想:将长CoT分割成多个语义单元(chunks),对每个chunk进行独立压缩,同时通过块间搜索确保整体连贯性。
- 具体步骤:
- 块分割:基于预定义长度和格式约束(如双换行符和最小长度要求)将原始CoT分割成多个chunk,确保每个chunk包含完整的推理单元。
- 块内压缩:利用一个独立的LLM(参数化为πϕ)对每个chunk生成多个压缩候选版本,通过精心设计的提示(prompt)引导模型生成简洁且保留原意的压缩结果。
- 块间搜索:为每个chunk生成多个候选压缩版本后,先通过长度过滤(丢弃最长的α比例候选)减少搜索空间,然后基于原始LLM(参数化为πθ)的条件概率,贪婪地选择与前文最连贯的候选chunk。
- 监督微调(SFT):使用压缩后的CoT数据集对基础模型进行微调,优化其生成简洁推理链的能力。
- 关键点与潜在问题:方法试图在块级尺度上保留局部信息(如反思),并通过搜索机制解决块间连贯性问题。然而,块间搜索依赖于概率选择,可能对模型的预测偏差敏感,且生成多个候选chunk增加了计算负担。此外,块分割策略可能对不同任务或CoT风格的适应性有限,缺乏对分割语义完整性的动态调整机制。
Experiment
实验基于Qwen2.5-Instruct系列模型(14B和32B)进行,数据集包括OpenR1-Math-220k(训练)和MATH500、AIME24、GPQA-Diamond(评估)。
- 设置:对比基线包括Long-CoT(未压缩)、CoT-Valve(实例级压缩)、TokenSkip(token级压缩)以及R1-Compress的随机变体。评估指标为准确率、平均token长度和有效token长度(正确回答的token长度)。
- 结果:在MATH500上,R1-Compress在Qwen2.5-32B模型上实现了92.4%的准确率,仅比Long-CoT基线(93.0%)低0.6%,同时将有效token长度从2406减少到1949(约20%减少)。在其他数据集(如GPQA-Diamond)上也表现出较好的泛化性。相比其他基线,R1-Compress在准确率和token效率上均有优势。
- 反思与连贯性分析:R1-Compress保留了约78%的反思步骤(相比Long-CoT),远高于CoT-Valve和C3oT,且token级损失(衡量连贯性)低于TokenSkip,表明其输出更符合LLM的自然语言模式。
- 消融研究:较小的chunk大小有助于保留局部信息,提升压缩质量;不同的搜索模型(如DeepSeek-Distill)对准确率和token长度有不同影响。
- 评价与问题:实验设置较为全面,涵盖了不同模型规模和任务类型,结果显示方法在压缩与性能间取得了较好平衡。然而,实验缺乏对统计显著性的分析,token减少的20%是否在实际部署中带来显著速度提升未明确。此外,搜索机制的额外计算开销未被充分量化,可能抵消部分效率收益。GPQA-Diamond作为分布外任务的表现虽好,但样本量和任务多样性可能不足以证明泛化性。
Further Thoughts
R1-Compress 的块级压缩思路为高效推理提供了一个有趣的视角,但其搜索机制的计算开销和对固定分割策略的依赖可能限制其在动态任务中的应用。未来可以探索自适应的块分割方法,例如基于语义相似性或推理步骤重要性动态调整chunk大小。此外,结合潜在空间推理(latent space reasoning)或离线计算(如Sleep-time Compute)可能进一步减少在线推理负担。另一个值得思考的方向是,是否可以通过强化学习(如RLHF)直接优化压缩后的CoT质量,而不仅仅依赖概率搜索,以减少对额外候选生成的依赖。这也让我联想到,类似方法是否可以应用于多模态推理任务中,通过压缩视觉或文本推理步骤提升多模态系统的效率。