通过临时保留思维步骤与混合训练,实现高效推理且不损失准确率。
HybridThinker: Efficient Chain-of-Thought Reasoning via Compressed Memory and Transient Thought Steps

- 保留压缩记忆与临时思维步骤双重机制
- 在4个基准上平均准确率提升5.8点
- 适合需要高效推理的模型部署场景
长链思维(CoT)能提升大模型推理能力,但带来显著计算与内存开销。现有压缩方法通过记忆标记将思维步骤浓缩为紧凑表示,并在推理时仅保留这些表示,但细粒度信息丢失导致后续步骤更易出错。为此,我们提出 HybridThinker,除保留压缩表示外,还临时保留部分思维步骤以提供细节支持。然而,我们发现若在训练中让模型直接访问所有思维步骤,会绕过记忆标记机制,导致压缩与检索能力训练不足。因此,我们设计了一种混合训练策略:仅部分思维步骤可通过注意力机制被后续步骤访问,其余步骤被掩码,强制模型依赖记忆标记进行信息压缩与检索。在4个推理基准上,HybridThinker达到未压缩基线的性能,平均准确率提升5.8点,且推理时间相近。消融实验验证了临时保留与混合训练均贡献显著。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。