小米大模型团队 2025-06-20 10:01 广西
你是否想过，当人工智能解决数学题时，是怎么一步步思考的？
现在的大型语言模型（比如能做数学推理的 AI），常用一种 “思维链” 方法 —— 就像人类做题时在纸上写步骤一样，AI 会逐字生成中间推理过程，最后给出答案。但这种 “逐字写下来” 的方式就像让 AI 用慢动作解题，效率很低。如何让 AI 高效思考，成为亟待解决的问题。
针对此问题，小米大模型 PLUS 团队与中国人民大学的研究团队合作提出了压缩隐空间推理（Compressed Latent Reasoning, CoLaR），即通过将多个词元（字符）压缩到一个连续的数字向量（隐空间变量）进行推理，类似人类“心算”而非在纸上逐字写下计算过程，进而提升推理效率。
在 GSM 系列数据集（小学难度数学题）上，相较基准模型，CoLaR 在使用更少推理步数的同时将解题成功率提升了 14.1%，相较逐字思考的思维链方法，CoLaR 仅以 4.8% 的正确率为代价，实现了 53.3% 的推理速度提升。
在更具挑战性的 MATH 数据集（大学难度数学题）上，经过强化学习训练后的 CoLaR 模型在正确率提升 5.36% 的同时，将推理链长度降低了 82.8%，首次证明了强化学习在隐变量推理领域的有效性。
01
方法详解
监督学习训练（左）以及强化学习训练（右）示意图
推理过程以及检索分析示意图
CoLaR 在推理过程中预测的隐变量可以有效地在数十万词元中检索出多个最有意义的词元，证明了思考过程预测的隐变量具有浓缩且有意义的推理信息。
▍推理压缩与理解
核心模块：通过 “嵌入压缩模块”，每次训练迭代过程随机选择压缩因子c（如 c=2 或 c=5），将连续 c 个 词元 的嵌入向量压缩成一个隐变量。测试时只需修改提示中的压缩因子，即可控制模型的 “思考速度”（压缩因子越高，速度越快）。
技术细节：采用 1/√c 的缩放因子（而非简单平均）保持向量分布稳定，同时训练语言头预测压缩后的词元和答案，通过损失函数确保模型理解压缩后的语义（公式 1）。
▍动态隐变量预测
概率化设计：区别于传统确定性方法，“隐变量头” 预测高斯分布的均值和方差（公式 2-3，图 2 左），通过重参数化技巧采样下一个隐变量，保留探索不同推理路径的能力。
▍强化学习调优
核心算法：采用 GRPO 算法（公式 4-5，图 2 右）对隐变量推理路径进行优化，对每一个问题生成多组推理路径并根据答案正确性分配奖励。
探索-利用：将组奖励平均分配到每个 “词”，既鼓励探索正确路径，又偏好更简洁的推理过程。在 MATH 数据集上，该方法使准确率提升 5.36%，同时推理长度减少 82.8%。
02
实验结果
▍主要实验
小学数学数据集（GSM8k 等）：相较显式 “思维链” 方法，CoLaR 仅损失 4.8% 的成功率，实现 53.3% 的推理性能提升；
数学竞赛题（MATH）：强化学习版 CoLaR 在推理长度压缩 82.8% 的同时，准确率提升 5.36%，首次证明强化学习在隐变量推理中的有效性。
▍分析性试验
泛化能力：不同训练压缩因子可互相泛化，提升彼此性能，模型甚至能适应训练中未见过的压缩因子；
网络利用：深层网络在高压缩因子下被更充分激活
强化学习中的探索与利用：训练过程呈现 “探索 - 利用” 两阶段：初期尝试不同推理路径（准确率快速提升），后期优化路径长度（准确率稳定，步骤减少）。
03
总结与展望
CoLaR 首次实现了动态调节压缩率的隐变量推理，让 AI 能像人类一样 “灵活心算”，突破了传统方法固定步数和缺乏探索的局限。尽管目前性能略低于传统显式推理，但团队认为通过扩大训练数据和迭代训练策略，有望实现质的突破。
未来研究将聚焦模型根据问题的难度自主选择压缩率，甚至在推理过程中动态调整压缩率，实现效果和效率的完美结合，进而推动大模型在高效推理场景中的实际应用。
END
