本文提出Compressed Latent Reasoning (CoLaR)框架，通过潜在空间动态压缩和强化学习优化大型语言模型的推理过程，在数学推理任务中显著提升效率并保持较高准确率。
Large Language Model, Reasoning, Reinforcement Learning, Efficiency, Representation Learning
Wenhui Tan, Jiaze Li, Jianzhong Ju, Zhenbo Luo, Jian Luan, Ruihua Song
中国人民大学, 小米公司
Generated by grok-3
Background Problem
大型语言模型（LLMs）通过链式推理（Chain-of-Thought, CoT）在数学推理任务中表现出色，但冗长的推理链带来了高昂的计算成本和低效率问题，特别是在高并发场景下对服务器负载造成显著压力。本文提出了一种新的框架Compressed Latent Reasoning (CoLaR)，旨在通过在潜在空间中动态压缩推理过程，解决推理链过长的问题，同时保持模型性能和探索-利用能力。
Method
CoLaR框架通过两阶段训练实现推理过程的动态压缩：
- 监督微调（SFT）阶段：引入辅助任务‘下一压缩嵌入预测’，在每个训练步骤中随机采样压缩因子，通过Embedding Compress模块将连续个推理token的嵌入合并为一个压缩嵌入，并训练一个专门的Latent Head预测后续压缩嵌入的分布（均值和标准差），从而实现潜在空间中的自回归推理。嵌入压缩避免了简单均值池化导致的分布扭曲，通过缩放调整保持原始分布特性。此外，语言头通过多标签近似预测压缩推理token，提供密集监督信号。
- 强化学习（RL）阶段：利用Latent Head的概率特性，通过Group Relative Policy Optimization (GRPO)算法探索多样化的推理路径并利用更短的推理链。具体而言，对每个问题采样一组输出，计算相对奖励并对每个token平均化奖励，鼓励模型在探索正确推理路径和利用高效路径之间取得平衡。
- 核心创新：动态压缩因子允许推理链长度灵活调整，概率性Latent Head支持探索多样化路径，强化学习进一步优化效率和准确性。
- 批评性思考：虽然方法设计新颖，但潜在空间压缩可能丢失关键推理信息，尤其是在高压缩因子下。此外，随机采样压缩因子的方式可能导致训练不稳定，难以适应非整数或超出训练范围的压缩因子，限制了方法的泛化能力。
Experiment
实验在四个小学级数学推理数据集（GSM8k-Aug, GSM-Hard, SVAMP, MultiArith）和一个更具挑战性的MATH数据集上进行评估：
- 数据集与指标：GSM8k-Aug用于主要训练和测试，其他数据集用于域外泛化测试，MATH用于评估复杂任务表现。指标包括准确率（Acc.）和推理链长度（# L）。
- 基线对比：与CoT、iCoT、Coconut和Distill等方法对比，CoLaR在压缩比相似的情况下平均准确率提升14.1%，推理链长度减少53.3%，仅损失4.8%的性能；在MATH数据集上，强化学习增强的CoLaR准确率提升5.36%，推理链长度减少82.8%。
- 消融研究：验证了概率性Latent Head和密集监督信号的重要性，确定性Latent Head和NLL损失在简单数据集上表现不佳，均值池化导致分布偏移和性能下降。
- 实验设计分析：实验设置较为全面，涵盖了不同难度的数据集和多种压缩因子，但对简单数据集上强化学习效果不显著的解释不够深入，可能存在方法对任务复杂度的依赖性。此外，测试中未充分探讨非整数压缩因子的影响，实验结果可能高估了方法在实际应用中的泛化能力。
- 批评性思考：虽然结果显示效率和准确性有显著提升，但与显式CoT相比仍有性能差距，压缩过程可能导致信息丢失，尤其在复杂任务中可能影响推理质量。实验中对基线方法的实现细节（如Distill的再现）缺乏充分说明，可能影响对比的公平性。
Further Thoughts
CoLaR框架在潜在空间压缩推理链方面的创新值得关注，但其对任务复杂度的依赖性提示我们需要在不同领域进一步测试其适用性。例如，在自然语言推理或多模态任务中，潜在空间压缩是否仍能有效保留关键信息？此外，强化学习在简单任务上的效果不显著，可能意味着需要设计更自适应的奖励机制，以根据任务难度动态调整探索和利用的平衡。另一个有趣的方向是探索与知识蒸馏技术的结合，通过将显式CoT的推理知识更精细地转移到潜在空间中，可能进一步提升性能并减少信息丢失。结合近期关于LLM推理路径多样性的研究，CoLaR的概率性Latent Head或许可以与其他探索策略（如蒙特卡洛树搜索）结合，以在复杂任务中实现更高效的推理路径搜索。