HMPO: Hybrid Median-length Policy Optimization for Chain-of-Thought Compression

📄 arXiv: 2606.01934v1 📥 PDF

作者: Minghui Zheng, Hongxu Chen, Huimin Ren, Hongsheng Xin, Xiaoyang Qu, Ze Wang, Shuling Yang, Ziyu Peng, Kaike Zhang, Pan Zhou, Kun Zhan

分类: cs.LG, cs.CL

发布日期: 2026-06-01


💡 一句话要点

提出HMPO,通过混合中值策略优化实现CoT压缩,降低推理开销。

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 思维链压缩 强化学习 策略优化 语言模型 推理效率

📋 核心要点

  1. 现有CoT压缩方法存在手动调整长度预算、多阶段训练成本高昂以及模型泛化能力不足等问题。
  2. HMPO采用单阶段强化学习框架,通过自适应中值预算、余弦衰减奖励和平滑奖励函数来压缩CoT。
  3. 实验结果表明,HMPO在压缩CoT的同时,保持了较高的准确率,并降低了训练成本,可扩展到大模型。

📝 摘要(中文)

大型语言模型通过扩展的思维链(CoT)推理实现了卓越的性能,但这种冗长的过程导致了巨大的推理开销。现有的CoT压缩方法受限于不灵活的手动长度预算、计算成本高的多阶段训练流程以及脆弱的可扩展性(仅限于小型模型)。我们提出了HMPO(混合中值长度策略优化),这是一种经济高效的单阶段强化学习框架。HMPO通过三个协同组件有效地压缩CoT:从成功rollout中导出的自适应中值预算,以消除手动调整;用于平滑长度惩罚的余弦衰减token奖励;以及通过严格优先考虑答案正确性来显著减轻琐碎奖励利用的乘法奖励公式。HMPO仅在数学数据上进行训练,便可无缝地推广到数学、代码、科学和指令遵循任务。从9B到122B参数的密集和混合专家(MoE)架构的广泛实验表明,HMPO实现了19%--46%的token压缩,而精度下降可忽略不计,同时与现有的多阶段基线相比,大大降低了训练成本。

🔬 方法详解

问题定义:论文旨在解决大型语言模型中,思维链(CoT)推理过程过长,导致推理开销巨大的问题。现有CoT压缩方法的痛点在于需要手动调整长度预算,训练流程复杂(多阶段),且难以扩展到大型模型上,同时容易出现奖励利用问题,导致模型只关注长度压缩而忽略了答案的正确性。

核心思路:论文的核心思路是利用强化学习,通过策略优化来学习如何压缩CoT,同时保证答案的正确性。通过设计合适的奖励函数,鼓励模型生成更短且正确的CoT。关键在于如何平衡长度压缩和答案正确性,以及如何避免模型利用奖励函数的漏洞。

技术框架:HMPO是一个单阶段的强化学习框架,主要包含以下几个模块:1) 环境:语言模型本身,用于生成CoT;2) 策略网络:用于生成压缩后的CoT;3) 奖励函数:用于评估生成的CoT的质量,包括长度和正确性;4) 优化器:用于更新策略网络的参数。整个流程是:给定一个输入,策略网络生成CoT,环境执行CoT并给出答案,奖励函数根据答案的正确性和CoT的长度给出奖励,优化器根据奖励更新策略网络。

关键创新:HMPO的关键创新在于以下几个方面:1) 自适应中值预算:通过分析成功rollout的长度分布,自动确定CoT的长度预算,避免了手动调整的麻烦;2) 余弦衰减token奖励:使用余弦函数来平滑长度惩罚,避免了奖励的突变,使得训练更加稳定;3) 乘法奖励公式:将答案正确性和长度惩罚相乘,而不是相加,从而严格优先考虑答案的正确性,避免了奖励利用问题。

关键设计:HMPO的关键设计包括:1) 奖励函数的设计:奖励函数由答案正确性奖励和长度惩罚组成,采用乘法形式,确保答案正确性是首要目标。长度惩罚采用余弦衰减函数,避免了奖励的突变。2) 自适应中值预算的计算:通过统计成功rollout的长度分布,计算中值长度作为长度预算。3) 策略网络的选择:可以使用各种语言模型作为策略网络,例如Transformer模型。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,HMPO在9B到122B参数的密集和混合专家(MoE)架构上,实现了19%--46%的token压缩,而精度下降可忽略不计。与现有的多阶段基线相比,HMPO大大降低了训练成本。例如,在数学数据集上训练的HMPO,可以泛化到代码、科学和指令遵循任务上,表现出良好的泛化能力。

🎯 应用场景

HMPO可应用于各种需要CoT推理的场景,例如数学问题求解、代码生成、科学推理和指令遵循等。通过压缩CoT,可以显著降低推理成本,提高推理效率,使得大型语言模型能够更广泛地应用于资源受限的设备上。该方法在教育、科研、智能客服等领域具有潜在的应用价值。

📄 摘要(原文)

Large language models achieve remarkable performance via extended chain-of-thought (CoT) reasoning, yet this lengthy process incurs substantial inference overhead. Existing CoT compression methods struggle with inflexible manual length budgets, computationally expensive multi-stage training pipelines, and fragile scalability restricted to small models. We propose HMPO (Hybrid Median-length Policy Optimization), a cost-effective, single-stage reinforcement learning framework. HMPO efficiently compresses CoT via three synergistic components: an adaptive median-based budget derived from successful rollouts to eliminate manual tuning, a cosine-decay token reward for smooth length penalization, and a multiplicative reward formulation that substantially mitigates trivial reward hacking by strictly prioritizing answer correctness. Trained exclusively on mathematical data, HMPO generalizes seamlessly across math, code, science, and instruction-following tasks. Extensive experiments scaling from 9B to 122B parameters across dense and Mixture-of-Experts (MoE) architectures demonstrate that HMPO achieves 19%--46% token compression with negligible accuracy degradation, all while drastically reducing training costs compared to existing multi-stage baselines.