arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2606.03768cs.CL

HybridThinker: 通过压缩记忆和瞬态思考步骤实现高效的思维链推理

HybridThinker: Efficient Chain-of-Thought Reasoning via Compressed Memory and Transient Thought Steps

School of Computer Science and Engineering, Northeastern University(东北大学计算机科学与工程学院);China Unicom Cloud-Link(中国联通云链);NiuTrans Research(牛传研究院)

Xin Liu, Runsong Zhao, Xinyu Liu, Junhao Ruan, Pengcheng Huang, Shichao Dong, Chunyang Xiao, Chenglong Wang, Changliang Li, Jingbo Zhu, Tong Xiao

AI总结:

提出HybridThinker方法,通过保留压缩记忆和临时保留思考步骤,并采用混合训练方案,在保持推理准确性的同时显著压缩思维链长度。

摘要:

扩展的思维链(CoT)轨迹提升了LLM的推理能力,但带来了大量的计算和内存成本。现有的CoT压缩方法通过记忆令牌将思考步骤压缩为紧凑表示,并在推理时仅保留这些表示,但细粒度信息的丢失使得后续步骤更容易出错。为了缓解这一问题,我们提出了 extbf{HybridThinker},除了保留这些表示外,思考步骤也会被临时保留以提供细粒度细节。然而,我们观察到在训练期间天真地让后续步骤可以访问思考步骤,会使模型绕过记忆令牌直接从这些步骤中检索信息,导致模型通过记忆令牌压缩和检索信息的能力训练不足。因此,我们引入了一种混合训练方案,其中只有部分思考步骤可以通过注意力直接访问后续步骤,而其他思考步骤被屏蔽,迫使模型使用记忆令牌进行压缩和检索。在4个推理基准测试中,HybridThinker与未压缩的基线性能相当,在相似的推理时间下,平均准确率比现有的CoT压缩方法提高了5.8个百分点。消融研究证实,临时思考步骤保留和混合训练方案都对这一提升有贡献。

英文摘要:

Extended chain-of-thought (CoT) traces improve LLM reasoning but incur substantial computational and memory costs. While existing CoT compression methods mitigate this by condensing thought steps into compact representations via memory tokens and retaining only these representations at inference time, the loss of fine-grained information makes subsequent steps more error-prone. To alleviate this, we propose \textbf{HybridThinker}, where in addition to preserved these representations, thought steps are also temporarily retained to provide fine-grained details. However, we observe that naively keeping thought steps accessible to subsequent steps \emph{during training} lets the model bypass memory tokens by retrieving information directly from these steps, leaving the model's ability to compress and retrieve information through memory tokens insufficiently trained. We therefore introduce a hybrid training scheme, in which only some thought steps are directly accessible through attention to subsequent steps, while the other thought steps are masked, forcing the model to use memory tokens for compression and retrieval. Across 4 reasoning benchmarks, HybridThinker matches the uncompressed baseline, advancing the state of the art in CoT compression by 5.8 points on average accuracy with similar inference time. Ablation studies confirm that both temporary thought-step retention and the hybrid training scheme contribute to these gains.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑