论文精读

精读：RiM — 用工作记忆替代自回归推理

用固定记忆块替代 CoT 自回归生成，单次前向传播完成潜在推理，TTFT 与直接回答相当

论文： Unlocking the Working Memory of Large Language Models for Latent Reasoning
作者： Lukas Aichberger, Sepp Hochreiter
arXiv： 2605.30343
潜在推理 工作记忆 推理加速 记忆块 test-time compute CoT替代 Hochreiter

1 论文概览

一句话总结： RiM 用固定长度的特殊 token 序列（记忆块）替代自回归推理步骤，让 LLM 在单次前向传播中完成潜在推理——推理速度与直接回答相同，准确率匹配或超越 Coconut/DART。

核心问题： 当前扩展 test-time compute 的主流范式（Chain-of-Thought）将推理与自回归生成耦合，导致：(1) 推理延迟线性增长（CoT 慢 ~27 倍）；(2) 内部计算被迫外化为自然语言，受语言表达瓶颈约束；(3) 推理效率与推理质量存在结构性矛盾。

主要贡献：

  1. 提出 RiM 框架：首次用固定记忆块替代自回归推理步骤，实现”推理即工作记忆”——内部计算与外部通信彻底解耦
  2. 两阶段课程训练：Stage 1 用显式推理步骤为记忆块奠基（grounding），Stage 2 去除步骤监督、迭代精炼最终答案——两个阶段缺一不可
  3. 极致推理效率：TTFT 与直接回答完全一致（Llama-3.2-1B: 16.1ms vs CoT 420.3ms），总延迟仅增加 0.5ms
  4. 线性探针答案选择：从记忆块表征中提取正确性信号的线性探针达 90% 准确率，可将 any-block 性能差距从 ~9pp 缩至接近零
  5. 跨模型家族验证：在 GPT-2 和 Llama-3.2（1B/3B）上一致有效，RiM Final Block 在 GSM8K 上比 Coconut 提升 2.5–7.5pp

2 方法详解

架构/方法概述

RiM 的核心思想来自认知科学中的工作记忆（Working Memory）概念：人类能在不外化中间思维的情况下，在内部保持和操作信息。RiM 将这一原理映射到 LLM 推理中。

记忆块结构： 每个记忆块是固定长度的特殊 token 序列：

mk=[<mb>,<m>,…,<m>⏟M tokens,<me>]m_k = [\text{<mb>}, \underbrace{\text{<m>}, \dots, \text{<m>}}_{M \text{ tokens}}, \text{<me>}]

  • 边界 token <mb> 和 <me> 标记记忆块的开始和结束
  • 内部 token <m> 是专用特殊 token，M=2（仅 2 个内部 token）
  • 记忆块被附加到输入问题之后，作为 LLM 的”内部工作空间”

关键设计选择：

  • 冻结原始词表嵌入，仅更新特殊 token 的嵌入——防止干扰已学到的语言知识
  • 自定义注意力掩码：每个 readout 位置只能关注输入问题和之前的记忆块，不能关注其他推理步骤或未来记忆块——强制中间计算编码在记忆块中

训练范式（两阶段课程）：

Stage 1 — 推理步骤监督（Grounding）：

将推理轨迹分割为 T 个显式步骤，分配 T 个记忆块。在每个记忆块 mtm_t 后，模型被监督预测下一个显式推理步骤 rt+1r_{t+1}：

LS1(w)=−∑t=1Tλt(s)log⁡pw(rt+1∣x,m≤t)L_{S1}(w) = -\sum_{t=1}^T \lambda_t(s) \log p_w(r_{t+1} | x, m_{\le t})

退火权重 λt(s)∈[0,1]\lambda_t(s) \in [0,1] 线性从 1 递减到 0——较早的推理步骤先失去监督，迫使模型将更基础的推理能力内化到记忆块中。

“Unlike approaches that gradually replace explicit steps, RiM directly replaces the entire trace with memory blocks in this stage.” — 论文原文

Stage 2 — 最终答案精炼：

固定 K=8 个记忆块，模型在每个记忆块后直接预测并精炼最终答案 y：

LS2(w)=−∑k=1Kαklog⁡pw(y∣x,m≤k)L_{S2}(w) = -\sum_{k=1}^K \alpha_k \log p_w(y | x, m_{\le k})

权重 αk\alpha_k 随 k 线性递增——鼓励后继记忆块（拥有更多潜在计算积累）产生更强答案。切换阶段时重置优化器状态和学习率调度器，使用更低学习率和更高 dropout。

关键创新点

  1. 计算与通信的解耦：推理步骤不再需要以自然语言形式输出，模型可以在潜在空间中自由”思考”——不受语言表达瓶颈约束
  2. 固定 vs 生成的范式转换：Coconut 的 latent token 仍需自回归生成（每步依赖前一步的输出），RiM 的记忆块是预设结构——因此可在单次前向传播中并行处理
  3. 密集监督的 grounding 策略：Coconut 的渐进课程（逐步替换显式步骤）在训练后期才建立完整潜在推理链；RiM 从 Stage 1 开始就用记忆块替换整个推理轨迹，提供更密集的中间监督信号

与基线方法的区别

维度CoT (SFT w/ CoT)CoconutRiM
推理 token自然语言，自回归生成潜在 token，自回归生成固定记忆块，单次前向
推理延迟~27x 直接回答~7x 直接回答~1x 直接回答
中间监督显式步骤标签渐进替换课程密集两阶段课程
计算位置输出空间潜在空间（顺序）潜在空间（并行）
可解释性高（自然语言 CoT）低低（但线性探针 90% 可提取正确性）

3 实验结果

主要指标

核心比较（Final Block 读出）：

模型方法GSM8K GreedyGSM8K Pass@8GSM-Hard GreedyTTFT (ms)
GPT-2SFT w/o CoT15.4%33.3%3.5%7.6
GPT-2Coconut31.1%45.0%7.1%53.4
GPT-2RiM Final33.6%49.1%7.8%7.6
Llama-3.2-1BSFT w/o CoT23.9%41.7%5.3%16.1
Llama-3.2-1BCoconut36.9%51.1%8.5%108.3
Llama-3.2-1BRiM Final42.1%56.1%10.5%16.1
Llama-3.2-3BSFT w/o CoT36.2%45.9%8.5%27.9
Llama-3.2-3BCoconut41.3%55.5%10.2%188.8
Llama-3.2-3BRiM Final48.8%58.8%12.0%27.9

Any Block 读出（选择所有记忆块中最优答案）：

模型方法GSM8K GreedyGSM8K Pass@8GSM-Hard Pass@8TTFT (ms)
GPT-2SFT w/ CoT39.8%57.0%12.9%213.7
GPT-2RiM Any39.5%78.1%19.0%7.6
Llama-3.2-1BSFT w/ CoT49.1%64.7%15.3%420.3
Llama-3.2-1BRiM Any51.4%76.8%19.6%16.1

与 SOTA 的对比

  • vs Coconut：RiM Final Block 在 GSM8K Greedy 上提升 2.5–7.5pp，在 GSM-Hard 上提升 0.7–1.8pp，同时 TTFT 快 7–7.5 倍
  • vs DART：RiM 在 GPT-2 (35.5% vs 24.7%) 和 Llama-3.2-3B (49.4% vs 46.6%) 上均超越 DART 官方报告数字，且 RiM 仅用 8 epochs + LoRA（DART 用 40 epochs + 全参数微调）
  • vs SFT w/ CoT：在小模型上 RiM Any Block 的 Pass@8 已超越 CoT（GPT-2: 78.1% vs 57.0%，Llama-1B: 76.8% vs 64.7%），且快 ~27 倍；但在 Llama-3.2-3B 上 RiM 的 Greedy 仍低于 CoT（57.3% vs 66.9%）

推理延迟对比

方法生成 Token 数总延迟 (ms)延迟增量
SFT w/o CoT3.1126.0—
SFT w/ CoT36.71108.7+982.7
Coconut3.1304.7+178.7
RiM3.1126.5+0.5

RiM 的推理延迟几乎与直接回答完全一致——仅增加 0.5ms。

消融实验关键发现

1. 两阶段课程缺一不可：

  • 仅 Stage 2：最终答案准确率远低于两阶段组合
  • 仅 Stage 1：Any-block 准确率高但 Final-block 准确率低
  • 过早切换到 Stage 2：效果显著下降

2. 密集监督优于渐进课程：

  • 用 Coconut 式渐进课程替换 RiM 的密集监督，整个训练过程中始终大幅低于 RiM

3. 记忆块表征确实在”工作”：

  • PCA 分析显示训练后记忆块表征变得块特异性（block-specific）和样本依赖（sample-dependent）
  • 最远样本余弦相似度：Layer 12 为 0.09，Layer 16 为 0.01——表征在原始空间中高度分化
  • 基础模型的记忆块表征基本坍缩——说明分化是训练习得的

4. 线性探针可提取正确性：

  • 从记忆块表征训练线性探针预测读出正确性：AUROC 84.2–86.0%，AUPRC 80.7–83.3%
  • 基于探针的答案选择准确率达 90%——Final-block 与 Any-block 之间的 ~9pp 差距可通过简单线性探针大幅缩小

5. 推理时 Memory 预算鲁棒：

  • Stage 2 后，准确率在广泛的 (M, K) 组合下保持稳定
  • 额外记忆块持续精炼预测，累积净效应为正

4 个人分析

亮点

1. 范式级概念转换——从”推理需要说出来”到”推理可以在脑子里完成”。这是对 CoT 范式的根本性挑战：CoT 的核心假设是推理必须以自然语言外化，RiM 证明推理可以在潜在空间中完成，且效果更好、速度更快。这一转换的认知科学动机（工作记忆）简洁而深刻。

2. 固定 vs 生成的关键区分。Coconut 虽然将推理移到潜在空间，但 latent token 仍是自回归生成的——每个 latent token 的嵌入取决于前一步的输出，因此必须顺序计算。RiM 的记忆块是预设的固定结构，所有 token 可以在一次前向传播中并行处理。这个区分看似简单，却带来了 27 倍的延迟差异。

3. 两阶段课程的精巧设计。Stage 1 的退火权重设计（较早步骤先失去监督）迫使模型将更基础的推理能力内化，Stage 2 的递增权重鼓励后继记忆块积累更多计算——这种”先奠基、后精炼”的范式与人类学习过程（先学步骤、后内化为直觉）有深层对应。

4. 线性探针的启示意义。90% 的探针准确率意味着：记忆块表征中已经包含了正确性信息，但模型的 final-block 读出机制未能充分利用这些信息。这暗示了一个重要的优化方向——不是需要更好的记忆块，而是需要更好的读出机制。

5. Sepp Hochreiter 的回归。作为 LSTM 的共同发明人（1997 年那篇引用 12 万+ 的论文），Hochreiter 在序列建模的根本性问题上有独特的洞察力。从”如何让序列模型记住长期依赖”（LSTM）到”如何让序列模型在内部推理而不依赖外部输出”（RiM），这是一条连贯的研究脉络——本质上都是在解决信息如何在时间维度上有效流动的问题。

局限性

1. 仅在小规模数学推理基准上验证。GSM8K 和 GSM-Hard 是小学数学题——推理链短、结构规整。论文未在逻辑推理、代码生成、常识推理等更广泛任务上验证泛化性。GSM-Hard 上的提升幅度（0.7–1.8pp）已显著小于 GSM8K（2.5–7.5pp），暗示更难问题可能需要更长的潜在推理链。

2. 大模型上的 Greedy 性能仍低于 CoT。Llama-3.2-3B 上 RiM Any Block 的 Greedy 准确率（57.3%）显著低于 SFT w/ CoT（66.9%），差距达 9.6pp。这说明 3B 模型的记忆块可能不足以承载 CoT 所能表达的复杂推理——记忆容量是硬约束。

3. 可解释性的根本性丧失。CoT 的最大优势是可解释性——你可以逐步检查推理过程。RiM 的记忆块是完全不可解释的潜在表征。虽然线性探针可以提取正确性信号，但这离真正的可解释推理审计还有很大距离。在对齐和安全关键场景中，这是重大局限。

4. 记忆块容量的固定性。K=8 个记忆块（每个 M=2 个内部 token，共 32 个特殊 token）是硬编码的，无法根据问题复杂度自适应调整。相比之下，CoT 可以根据问题难度自然地生成不同长度的推理链。

5. “Any Block”评估的实际可行性问题。Any Block 读出需要遍历所有 K 个记忆块的预测并选择最优——但在没有验证器的情况下，如何知道哪个是”最优”？线性探针是论文给出的答案，但 90% 的准确率仍意味着 10% 的选择错误率。

对大模型领域的影响

1. 潜在推理的工程可行性论证。RiM 是首个在保持直接回答推理速度的同时，达到与 Coconut 相当或更优性能的潜在推理方法。这为潜在推理从学术探索走向工程部署扫清了关键障碍——延迟不再是取舍因素。

2. 与 Overthinking 系列论文的深层对话。本专栏此前精读的 Overthinking 论文揭示推理 token 过多反而有害，Illusion of Thinking 揭示推理在高复杂度下崩塌，EDRM 揭示推理是动态状态而非静态能力。RiM 从另一个维度回应了这些发现：如果推理不需要外化为 token，那么 Overthinking 的问题可能从根源上就不存在——记忆块不存在”过度生成”的问题，因为它们是固定数量的。

3. 与 GRAM/Equilibrium Reasoners 的对照。GRAM（概率生成递归推理）和 EqR（吸引子推理）分别从概率性和确定性角度探索潜在推理。RiM 提供了第三条路径——不是递归迭代，而是并行处理固定结构。三条路径的对比（递归概率 / 递归确定 / 并行固定）为潜在推理的完整方法论图谱提供了关键拼图。

4. 推理效率的新基线。27 倍的延迟优势意味着：在延迟敏感的部署场景（实时对话、代码补全、Agent 交互），RiM 式的潜在推理可能是唯一可行的推理增强方案。这可能改变 test-time compute scaling 的研究方向——从”生成更多 token”转向”设计更好的潜在计算结构”。

5 实践启示

可以借鉴的技术

1. 固定计算结构 + 潜在推理的范式。不仅限于推理增强——任何需要在输出前进行内部计算的场景（如检索增强生成前的查询重写、多步工具调用前的规划）都可以借鉴”固定记忆块”的设计思想，避免自回归生成带来的延迟开销。

2. 两阶段课程的 Grounding → Refinement 范式。Stage 1 的核心洞见是：潜在工作空间需要显式监督信号来建立计算角色。这一洞见可推广到任何需要训练模型使用新引入的潜在组件的场景——先用显式标签奠基，再去掉标签让模型自主使用。

3. 线性探针作为质量评估工具。90% 准确率的线性探针说明正确性信息高度线性可分——这是轻量级质量评估的直接工程工具。可以在不运行完整推理的情况下，快速估计当前推理状态的质量。

4. 冻结原始嵌入 + 仅训练特殊 token。这一设计防止了引入新组件时对已有知识的干扰，是增量式模型扩展的通用最佳实践。

潜在应用场景

1. 延迟敏感的推理增强。实时对话系统、代码补全、搜索查询理解等场景，CoT 推理的延迟不可接受，RiM 式潜在推理可能是唯一可行的推理增强方案。

2. Agent 内部规划。Agent 在执行多步任务前需要”想一下”——与其生成冗长的规划文本（增加延迟和成本），不如在记忆块中完成快速潜在规划，仅输出最终行动计划。

3. 多候选答案的并行评估。RiM 的 K=8 个记忆块自然产生 K 个候选答案，配合线性探针可实现高效的 Best-of-K 选择——比独立采样 K 条 CoT 快 ~27 倍。

4. 后训练数据质检的效率优化。在数据质检场景中，需要对大量样本进行快速推理验证。RiM 的 0.5ms 延迟增量意味着质检吞吐量可提升 ~27 倍，同时保持推理质量。

5. 端侧推理增强。移动/嵌入式设备上，自回归推理的延迟和能耗是硬约束。RiM 的单次前向推理模式天然适合端侧部署——不需要 KV cache 的逐步扩展，不需要投机解码的复杂调度。

论文信息： Aichberger, L. & Hochreiter, S. (2026). Unlocking the Working Memory of Large Language Models for Latent Reasoning. arXiv:2605.30343.

关联精读： 本论文与此前精读的 Overthinking、Illusion of Thinking、GRAM、Equilibrium Reasoners、EDRM 形成潜在推理与推理效率的完整论证链——从”推理为何过度”→“推理在什么条件下崩塌”→“推理何时是动态状态”→“推理如何在潜在空间中递归”→“推理如何向吸引子收敛”→“推理如何在工作记忆中并行完成”。