Unlocking the Working Memory of Large Language Models for Latent Reasoning
解锁大语言模型的工作记忆以进行隐式推理
🧭 隐式推理 · 工作记忆 · 两阶段课程
Johannes Kepler University Linz · NXAI GmbH
用固定记忆块代替自回归中间推理，在单次前向传播中实现高效隐式推理
📌 工作参考价值
- 用固定记忆块实现低延迟推理：将中间思考外化替换为少量固定记忆令牌，直接附加到输入上，单次前向传播即可完成推理，大幅降低生成延迟，适用于对实时性敏感的 Agent 或推荐系统。
- 两阶段课程可复用于训练隐式推理模型：先通过中间步骤监督迫使模型把信息编码进记忆块，再逐步转为仅用最终答案优化，这种方法同样可用于训练推荐模型中的用户意图或状态隐式表示，例如在会话式推荐中用记忆块累积用户偏好。
- 注意力掩码实现高效并行训练：通过分离记忆块流和读出头，一次前向传播就能完成所有中间目标的监督，训练效率接近标准微调。在训练其他需要多步推理或序列决策的模型时，可以用类似掩码减少训练时间。
- 推理时记忆预算可调，方便工程折中：训练好的模型在不同记忆块数量和大小下表现稳定，实际部署时可以根据算力限制灵活调整推理成本而不必重新训练，适合线上动态裁剪。
动机：当前大语言模型 (LLM) 的推理能力高度依赖生成中间文本（如链式思考，CoT），但将内部计算与语言生成耦合会带来两大问题：生成流畅文本消耗额外计算，且每个中间步骤必须自回归生成，造成推理延迟。人类认知依赖于工作记忆在不外化思考的情况下进行内部操作，为此，本文提出 Reasoning in Memory (RiM)，将 LLM 的中间推理转移到固定的记忆块中，通过单次前向传播完成隐式推理，解耦计算与自回归生成。
方法关键点：
- 定义记忆块为一段由特殊令牌（
<b>,<m>,</b>）组成的固定序列，作为隐式工作空间，输入时直接附加在问题之后。 - 采用两阶段课程学习解锁记忆块的推理能力：
- 阶段1：用原问题的显式推理步骤作为监督信号，要求每个记忆块后的读出头预测下一个推理步骤。自定义注意力掩码确保读出头只能看到对应记忆块及之前的内容，不能窥视后续步骤，从而强制将中间计算编码进记忆块。
- 阶段2：去除中间步骤监督，改为每个记忆块后直接预测最终答案，训练模型利用隐式工作空间逐步精炼答案。两个阶段均使用标准的下一个令牌预测目标，避免复杂的多路径蒸馏。
- 训练效率优化：通过块级因果注意力掩码，在一次前向传播中并行计算所有读出的监督损失，显著提升训练效率。
- 冻结原有令牌嵌入，仅新增记忆块的特殊令牌嵌入可学习，减少对预训练模型的破坏。
关键实验：
- 在数学推理基准上验证，训练集为 GSM8K-Aug（38.6万条带分步推理的数学题），评估集为 GSM8K (ID) 和 GSM-Hard (OOD)。
- 对比基线包括 SFT（直接答案和 CoT 变体）和显式隐式推理 Coconut，涵盖 GPT‑2 和 Llama-3.2‑1B/3B 等模型规模。
- RiM 在 Llama-3.2‑1B 上 greedy 准确率达 42.1%，比 Coconut 高出 5.2 个百分点，且延迟仅 16.1ms（与直接答案生成相当），Coconut 为 108.3ms。RiM 的任意块准确率（取最佳记忆块输出）可达 51.4%，逼近显式 CoT 的 49.1% 但延迟降低 26 倍。
- 推理时记忆块数量和每块令牌数可灵活调整，性能保持稳定，第二阶段训练后对推理预算不敏感。
核心结论：固定记忆块配合两阶段课程可以有效赋予 LLM 在隐式空间进行工作记忆推理的能力，精度超越现有隐式推理方法，且推理速度与直接答案生成持平，为高效推理提供了新范式。