小米大模型团队 2025-06-20 10:01 广西
- 论文地址:https://arxiv.org/pdf/2505.16552
训练代码:https://github.com/xiaomi-research/CoLaR模型权重:https://huggingface.co/AlbertTan/CoLaR
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