#Unlocking the Working Memory of Large Language Models for Latent Reasoning:让大模型在“工作记忆”里推理

论文:Unlocking the Working Memory of Large Language Models for Latent Reasoning

作者:Lukas Aichberger, Sepp Hochreiter

时间:2026-05-28

arXiv:<https://arxiv.org/abs/2605.30343>

#0. 一句话先讲清楚

这篇论文提出 Reasoning in Memory(RiM):不要再让模型把中间推理一步一步生成出来,无论是生成自然语言 CoT,还是像 Coconut 那样自回归地产生 continuous thoughts;而是在输入后面放一串固定的特殊 token 组成 memory blocks,训练模型把这些 block 的上下文表示当作内部工作记忆,在一次 forward pass 里完成潜在推理。

换句话说,RiM 想做的事情是:

让 LLM 不必“边想边说”,而是学会在一段固定的内部工作区里“默想”。

它的核心判断很直接:

  1. CoT 的问题不是“语言没用”,而是把计算和交流绑定太死。模型为了推理,必须先生成一长串可读文本。
  2. 现有 latent reasoning 虽然把文字换成连续表示,但很多方法仍然要一步一步自回归地产生 latent state,本质上还是“连续空间里的边想边说”。
  3. 如果中间计算能放在固定 memory tokens 的 hidden states 里,那么这些 tokens 可以并行处理,推理延迟接近 direct answer。
  4. 关键难点不在架构,而在训练信号:普通 filler tokens 很容易被模型忽略,必须用合适 curriculum 把它们“激活”为工作记忆。
Figure 1:RiM 总览。Stage 1 让 memory blocks 学会预测中间推理步骤;Stage 2 去掉中间步骤监督,让每个 block 后的 readout 直接预测最终答案。
Figure 1:Reasoning in Memory 总览

这篇文章和 wenjun 最近关心的 latent-space reasoning、长轨迹 agent RL、model-based RL for LLM Agent 很相关。因为它不是单纯在问“怎么多采样几个答案”,而是在问一个更基础的问题:

Transformer 能不能把一段 token position 训练成内部可读写的 latent workspace?如果可以,它能不能成为未来 agent 长程推理、规划、世界模型 roll-out 的低成本计算基底?

下面按“问题背景 → 方法 → 实验 → 机制证据 → 局限与启发”来讲。


#1. 这篇论文想解决什么问题?

#1.1 CoT 的核心矛盾:推理计算被迫以语言形式外化

当前提升 LLM 推理能力最常见的办法是增加 test-time compute:

  • Chain-of-Thought:让模型写出中间推理步骤;
  • self-consistency / repeated sampling:生成多条 reasoning traces;
  • search / verifier / PRM:围绕中间步骤做搜索和评估;
  • RL reasoning models:训练模型生成更长、更有效的思维链。

这些方法的共同点是:模型要想更多,就得说更多。

但语言本来是为了交流,不是为了内部计算。自然语言 CoT 有几个代价:

代价直白解释
生成成本高每个中间 token 都要自回归生成,不能并行
语法负担模型要花算力生成通顺、可读、格式正确的文字
表达瓶颈有些内部计算未必适合压成自然语言步骤
延迟高first token / full answer latency 都会显著增加

所以论文提出一个区分:

  • external communication:把思考过程写给人看;
  • internal computation:模型为了得到答案而进行的内部状态变换。

CoT 把这两者混在一起了。RiM 想把它们拆开。

#1.2 现有 latent reasoning 仍然没有完全摆脱自回归瓶颈

近两年有不少 latent reasoning 工作,典型包括 Coconut、compressed CoT、pause/filler tokens、DART 等。

论文把它们大致分成几类:

  1. 显式 latent reasoning:把自然语言 reasoning token 替换成连续表示。例如 Coconut 生成 continuous thoughts(CTs),再把 CT feed back 到后续解码。
  2. 隐式 latent reasoning / filler-token reasoning:给模型插入没有语义的 token,让它们承载内部计算。
  3. recurrent / vertical latent reasoning:通过循环模块或 hidden-state refinement 做多步计算。

RiM 对 Coconut 的批评尤其关键:Coconut 虽然不写自然语言思维链,但它还是要一个 CT 一个 CT 地生成。也就是说,它把“文字 token 的自回归”换成了“连续表示的自回归”。

这当然比完整 CoT 短,但中间计算仍然被迫 externalize 成一串 sequential latent states。后一个 latent state 必须等前一个产生后才能计算。

论文的观点是:

latent reasoning 真正要解耦的是“中间计算”和“自回归生成”。如果 latent state 仍要一个个生成,那么只是把语言空间换成连续空间,还没有彻底释放并行计算能力。


#2. RiM 的核心想法:固定 memory blocks 作为工作记忆

RiM 的设计非常简单:在问题后面追加若干个固定的 memory blocks。每个 memory block 是一串特殊 token:

<b> <m> <m> ... <m> </b>

其中:

  • <m> 是真正的 memory token;
  • <b></b> 标记 block 边界;
  • 每个 block 由固定数量 M 个 memory token 组成;
  • 实验默认 M = 2;
  • 多个 block 串起来形成 K 个 latent computation slots。

形式上,输入从:

question -> answer

变成:

question -> memory block 1 -> memory block 2 -> ... -> memory block K -> answer readout

关键在于:这些 memory tokens 的 token identity 和位置是固定的,但它们经过 Transformer 后的 contextual representations 是输入相关的。也就是说,同样的 <m> token,在不同问题里会形成不同 hidden states。

这就是论文所谓的 working memory:

固定 token 提供位置和容器,hidden states 承载任务相关的中间计算。

#2.1 为什么固定 memory blocks 能降低延迟?

因为它们不是模型逐步生成出来的,而是预先放进输入里的 token。Transformer 可以对整段 augmented sequence 做一次 forward pass。

对比一下:

方法中间计算形式是否需要自回归地产生中间状态推理延迟
CoT文本 reasoning tokens需要
Coconutcontinuous thoughts需要中高
RiM固定 memory token 的 hidden states不需要接近 direct answer

这点是 RiM 的最大卖点:它不是减少中间步骤长度,而是改变中间步骤的计算拓扑

#2.2 只插入 filler tokens 为什么不够?

这篇论文并不是第一个想到 filler/memory tokens 的。之前很多工作已经发现:简单插入 pause tokens 或 filler tokens,效果通常不稳定,甚至会变差。

原因很简单:Transformer 没有天然知道这些 token 是“草稿纸”还是“噪音”。如果训练目标没有强迫模型把信息写进去、读出来,它很可能直接忽略它们。

所以 RiM 的核心不只是 memory block 本身,而是两阶段训练:

  1. Stage 1:用显式 reasoning steps 监督 memory blocks,给它们赋予计算角色。
  2. Stage 2:去掉中间步骤监督,只监督最终答案,让 memory blocks 从“复述中间步骤”转向“服务最终答案”。

#3. Stage 1:先教 memory blocks 复原中间推理步骤

Stage 1 的目标是:让 memory blocks 不再是空 token,而是被迫存储和变换与推理相关的信息。

训练样本包含:

问题 x
显式推理轨迹 r1, r2, ..., rT
最终答案 y

RiM 为每个 reasoning step 分配一个 memory block。第 t 个 memory block 后面接一个 readout,让模型预测下一个 reasoning step:

x + m1 -> predict r2
x + m1 + m2 -> predict r3
...
x + m1 + ... + mT -> predict y

论文里的目标函数大意是:

Stage 1 loss = 对每个 memory block 后的 readout 做 next reasoning step NLL

更直白地说:

你不能直接看到前面的文字推理步骤,只能看问题和已经经过的 memory blocks;如果你要预测下一步推理,就必须把有用信息压到 memory blocks 里。

Figure 3:RiM Stage 1 和 Coconut curriculum 的差异。Coconut 逐阶段增加 continuous thoughts;RiM 一开始就用所有 memory blocks 做密集监督。
Figure 3:RiM Stage 1

#3.1 RiM 为什么不直接照搬 Coconut 的 curriculum?

Coconut 的思路是逐步把显式 reasoning steps 替换成 continuous thoughts。RiM 认为这种做法有两个问题:

  1. 监督可能绕过 latent workspace。 早期阶段里,目标 reasoning step 可以 attend 到前面的 written steps,于是模型可以直接从文字步骤预测下一个步骤,不一定需要把信息写进 latent tokens。
  2. 监督不够密集。 Coconut-style staging 需要预设最多多少阶段,可能只有部分 reasoning steps 被 latent 替换,很多中间步骤不会直接从 memory block 读出来监督。

RiM 的解决方案是:一开始就把所有 reasoning steps 对齐到 memory blocks,并用 attention mask 防止信息泄漏。


#4. Attention mask:防止模型偷看答案,强迫信息经过 memory blocks

RiM 能做 dense supervision 的关键是自定义 attention mask。

Figure 2:RiM attention mask。memory blocks 能看问题和之前 memory blocks;target reasoning steps 不能看其他 target steps,从而避免信息泄漏。
Figure 2:RiM Attention Mask

它的基本规则是:

  • memory blocks 可以 attend 到问题和之前的 memory blocks;
  • target reasoning step readout 可以 attend 到之前的 memory blocks,也可以选择是否 attend 到问题;
  • target reasoning steps 之间不能互相 attend;
  • 因此所有 target 可以在一个 forward pass 里训练,但不会互相抄答案。

这点很重要。否则 Stage 1 会退化成普通 teacher forcing:模型看到前面的推理文字,自然能预测后面的步骤,memory blocks 学不到东西。

可以把 Stage 1 想象成一种“带遮挡的中间状态预测”:

你要预测下一步推理,但我把可直接利用的文字步骤遮住,只留给你 latent memory。于是你必须在 memory 里组织信息。

这和 JEPA / predictive representation learning 的精神有点像:通过预测缺失结构,让表示学会承载任务相关信息。


#5. Stage 2:去掉中间步骤监督,训练最终答案 refinement

Stage 1 的问题是:它仍然依赖显式 reasoning steps 作为 teacher。这样训练出来的 memory blocks 可能只是“中间步骤压缩器”,不一定最适合最终答案。

所以 Stage 2 做了一个 hard switch:

  • 不再使用中间 reasoning trace;
  • 固定使用 K 个 memory blocks;
  • 每个 memory block 后面都接 answer readout;
  • 每个 readout 都直接预测最终答案 y;
  • 后面的 readout 权重更大,因为它能看到更多 memory blocks。

形式上类似:

x + m1 -> answer
x + m1 + m2 -> answer
...
x + m1 + ... + mK -> answer

这一步的直觉是:

Stage 1 先教模型如何在 memory 里放中间计算;Stage 2 再教模型如何把这些 latent computations 用来逐步修正最终答案。

这和人类学习也有点像:一开始可能需要把步骤写出来,后来内化以后,不必再逐字复述过程,而是直接在脑中操作并给出答案。

#5.1 Stage 1 和 Stage 2 的分工

阶段训练目标主要作用风险
Stage 1从 memory blocks 预测下一步 reasoning step赋予 memory blocks 中间计算角色过度绑定显式步骤
Stage 2从不同 memory budget 预测最终答案把 latent workspace 转成答案 refinement若无 Stage 1,memory tokens 可能没学会计算

论文附录里的 ablation 显示:只做 Stage 2 虽然能快速提高 final answer,但 plateau 更低;只做 Stage 1 any-block 可能不错,但 final-block 不够好。两者都需要。


#6. 实验设置:在 GSM8K-Aug 上训练,在 GSM8K/GSM-Hard 上评估

论文实验主要关注三个问题:

  1. memory blocks 是否真的学到了非平凡 latent computation?
  2. RiM 相比 SFT、CoT、Coconut 等方法,在准确率和延迟上如何?
  3. RiM 对不同 inference-time memory budget 是否鲁棒?

#6.1 数据与模型

训练数据:GSM8K-Aug

  • 约 385,620 个小学数学问题;
  • 每个样本有最多 13 个 reasoning steps;
  • Stage 1 用这些显式 arithmetic reasoning steps 做监督。

评估数据:

  • GSM8K:in-distribution test;
  • GSM-Hard:out-of-distribution,更难的数学问题。

模型:

  • GPT-2;
  • Llama-3.2-1B;
  • Llama-3.2-3B。

#6.2 Baselines

论文比较了:

baseline说明
SFT w/o CoT直接问题到答案,和 RiM 一样不生成显式推理
SFT w/ CoT生成完整 CoT,强 explicit reasoning baseline,但延迟高
Coconut最相关的 latent reasoning baseline,自回归生成 continuous thoughts
DART作为文献官方结果比较,非完全同协议

一个值得注意的评估细节:论文没有直接在 test set 上挑最好 checkpoint,而是用 GSM8K 的 held-out split 做 cross-validation checkpoint selection,再报告结果。这是为了减少 selection overfitting。


#7. 结果一:RiM 在低延迟下超过 Coconut 和 direct SFT

主结果如下:

Table 1:主结果。RiM final block 在 GPT-2、Llama-3.2-1B、Llama-3.2-3B 上都超过 SFT w/o CoT 和最强 Coconut 变体,同时 TTFT 接近 direct-answer SFT。
Table 1:Main Results

最重要的数字:

  • GPT-2:RiM final block 在 GSM8K 上 33.6%,Coconut 31.1%,SFT w/o CoT 15.4%。
  • Llama-3.2-1B:RiM 42.1%,Coconut 36.9%,SFT w/o CoT 23.9%。
  • Llama-3.2-3B:RiM 48.8%,Coconut 41.3%,SFT w/o CoT 36.2%。

论文总结为:RiM final-block readout 相比最强 Coconut,在 GSM8K 上提升约 2.5 到 7.5 个百分点;相比 direct-answer SFT,提升更大。

注意这里的 final block 是可部署设置:只看最后一个 memory block 后的答案,不假设 oracle 能在多个 block 中挑正确答案。

#7.1 Any-block 结果:latent trajectory 里常常已经出现正确答案

论文还报告了 any-block readout:如果任意一个 memory block 后的 readout 答对,就算正确。这不是直接可部署指标,但能反映 latent trajectory 的潜力。

Table 2:附加结果。RiM any-block 在部分设置中接近甚至超过 SFT w/ CoT,但延迟远低于显式 CoT。
Table 2:Additional Results

例如:

  • GPT-2:RiM any-block GSM8K greedy 39.5%,接近 SFT w/ CoT 的 39.8%;GSM-Hard greedy 9.4%,超过 SFT w/ CoT 的 8.4%。
  • Llama-3.2-1B:RiM any-block GSM8K 51.4%,超过 SFT w/ CoT 的 49.1%;GSM-Hard 13.0%,超过 11.2%。
  • Llama-3.2-3B:SFT w/ CoT 仍明显更强,RiM any-block 不及 CoT。

这说明一个有趣现象:

memory blocks 的中间 readouts 里经常已经包含正确答案,只是最终 block 未必总能选中或保持它。

论文进一步用 linear probes 发现,memory representations 中包含可预测 correctness 的信息;在至少一个 block 答对的 recoverable subset 上,简单 probe-based selection 能 90% 选中正确答案。这对未来做 learned verifier / controller 很有启发。


#8. 结果二:RiM 的延迟几乎等于 direct answer

RiM 最大优势不是绝对准确率超过 CoT,而是 accuracy-latency tradeoff

论文报告 TTFT:RiM 几乎和 SFT w/o CoT 一样,因为 memory blocks 是固定输入 token,一次 forward pass 处理。Coconut 因为要自回归生成 continuous thoughts,约慢 7 倍;SFT w/ CoT 因为要生成长文本 reasoning,约慢 27 倍。

附录 Table 7 报告 full answer wall-clock latency:

Table 7:Llama-3.2-1B 上的完整答案生成延迟。RiM 几乎等于 SFT w/o CoT,而 CoT 和 Coconut 明显更慢。
Table 7:Inference latency

Llama-3.2-1B 上:

方法生成 token 数wall-clock 时间
SFT w/o CoT3.1126.0 ms
SFT w/ CoT36.71108.7 ms
Coconut3.1304.7 ms
RiM3.1126.5 ms

这个结果很关键:RiM 不是靠“多生成一些中间 token”换准确率,而是试图在 forward computation 内部重分配计算。

当然要小心:memory tokens 仍然会增加输入长度和 attention FLOPs,只是避免了自回归循环。因此当 K、M 很大、上下文很长时,成本仍会增长。但在论文设置下,它几乎不增加 measured latency。


#9. 结果三:memory blocks 真的形成了输入相关的 latent workspace 吗?

如果 RiM 只是靠特殊 token trick 提升,那意义有限。论文因此做了 representation analysis。

它们收集 Llama-3.2-1B 在 GSM8K 测试问题上的 memory block representations,然后投影到 PCA 空间,观察训练前后变化。

Figure 4:memory block representations。训练后,不同 memory blocks 的表示沿着平滑、block-specific、sample-dependent 的轨迹移动。
Figure 4:Memory block representations

论文观察到:

  1. 训练轨迹是平滑的。 这说明训练不是随机扰动 hidden states,而是在系统组织表示空间。
  2. 不同 memory block 有不同轨迹。 说明 block 不是同质 filler token,而是形成了位置相关的计算角色。
  3. base model 表示更坍缩,final model 表示更分散。 说明训练后不同问题诱导出不同 latent workspace。
  4. first-to-final block delta 有结构。 说明沿 memory blocks 的方向确实发生了状态变化,而不是每个 block 都给同一个答案。

这部分是论文最有价值的机制证据之一:它尝试证明 memory blocks 不只是 prompt padding,而是变成了 task-dependent latent state。


#10. 训练曲线:Stage 2 是 final answer 能力跃迁的关键

Figure 5 展示 Llama-3.2-1B 在 GSM8K 上的训练曲线。

Figure 5:训练曲线。RiM 在进入 Stage 2 后 final-block accuracy 明显提升,并在相同训练预算下超过多个 Coconut 变体。
Figure 5:Llama-3.2-1B training curves

左图对比 RiM 和 SFT:

  • Stage 1 主要让 latent workspace 成形;
  • Stage 2 后 final memory block 准确率迅速提升;
  • any-block accuracy 变得接近 SFT w/ CoT。

右图对比 RiM 和 Coconut:

  • Coconut 加更多 CT stages 并没有稳定改善;
  • RiM 在同训练预算下超过多个 Coconut 变体;
  • 说明固定 memory blocks + dense supervision 的 curriculum,比逐步替换 reasoning steps 更有效。

这里我觉得最重要的是:RiM 的收益不是仅来自“多给几个 latent tokens”,而是来自训练方式让这些 token 承担计算角色。


#11. Memory budget 鲁棒性:给多少 memory block 都还能工作吗?

如果一种方法只能在训练时固定的 block 数量上工作,部署会很麻烦。论文因此测试不同 inference-time memory budget。

Figure 6:memory budget 鲁棒性。Stage 2 后,RiM 在不同 memory token/block 数量下准确率更稳定;额外 memory block 仍会改变答案,且总体净效果为正。
Figure 6:Robustness across memory budgets

Figure 6a 显示:

  • Stage 1 后,增加 block 数量反而可能使准确率下降,因为 Stage 1 的 readout 仍绑定到 reasoning-step positions;
  • Stage 2 后,准确率在不同 memory budgets 下稳定在约 40% 左右,最高约 43%;
  • 说明 Stage 2 把 memory blocks 从“对应第几步推理”转成了“固定 latent computation sequence”。

Figure 6b 显示:

  • 随着 memory block 增加,答案仍会变化;
  • incorrect-to-correct 的转变总体多于 correct-to-incorrect;
  • cumulative net accuracy 是正的。

这说明 latent workspace 没有 collapse 成每个 block 都输出同一个答案,而是在逐步 refinement。


#12. 训练诊断:latent state 的统计性质在 stage switch 后变化

附录 Figure 13 给了训练过程中的诊断指标:latent similarity、latent norm、effective rank、total memory tokens、training perplexity。

Figure 13:RiM 训练诊断。stage switch 附近,latent representation 的相似度、范数、有效秩等指标出现系统变化。
Figure 13:Training diagnostics for RiM

这些图不是主结论,但提供了一个信号:Stage 1 到 Stage 2 的切换不是简单继续训练,而是目标函数变化导致 latent representation 的组织方式也变化。

对做 mechanistic / representation analysis 的人来说,这里还有继续挖的空间:

  • 哪些层承担 memory computation?
  • memory tokens 之间是否形成类似 scratchpad slots 的分工?
  • final answer readout 是否在读取某些低维方向?
  • correctness probe 能否变成可部署 selector?

#13. 这篇论文真正贡献在哪里?

我认为它的贡献可以分成三层。

#13.1 概念层:把 latent reasoning 从“连续 CoT”推进到“工作记忆”

很多 latent reasoning 方法只是把中间文字 token 换成 hidden states,但仍保留 sequence generation paradigm。RiM 强调:

真正的 latent reasoning 不应该只是“看不见的思维链”,而应该是一个内部 workspace。

这点很重要。因为如果 latent reasoning 仍然自回归地产生中间状态,它在计算图上仍然像 CoT;而 memory blocks 更像 Transformer 内部的一段可训练 scratchpad。

#13.2 方法层:用两阶段 curriculum 激活 filler tokens

RiM 没有引入复杂新模块,也没有多路径 self-distillation。它主要靠:

  • 固定特殊 tokens;
  • attention mask;
  • Stage 1 reasoning-step supervision;
  • Stage 2 final-answer refinement。

这说明一个可能的方向:与其设计更复杂的 latent module,不如设计更强的 training signal topology。也就是说,关键是让梯度如何穿过 latent workspace。

#13.3 实验层:在低延迟下超过 Coconut

在 controlled evaluation 下,RiM final-block 超过 Coconut;延迟又接近 SFT w/o CoT。这给了它很强的 practical appeal。

当然,它还没有证明能替代显式 CoT。在 Llama-3.2-3B 上,SFT w/ CoT 仍然明显更强。但 RiM 提供了一个有趣折中:

如果你不能接受长 CoT 的延迟,又希望比 direct answer 更会推理,那么固定 memory workspace 可能是一个方向。


#14. 局限:不要把这篇论文过度解读

这篇论文很有启发,但也有明显边界。

#14.1 任务主要是 GSM 系列数学题

实验集中在 GSM8K-Aug → GSM8K/GSM-Hard。虽然有 ID/OOD 区分,但仍然属于短程数学 word problems。

还没有证明:

  • 多步代码推理;
  • 长上下文规划;
  • agent tool-use;
  • theorem proving;
  • open-ended research reasoning;
  • 需要搜索和回溯的任务;

也能同样受益。

#14.2 Stage 1 依赖显式 reasoning traces

RiM 需要中间 reasoning steps 来 ground memory blocks。这在 GSM8K-Aug 上可行,但在很多真实任务中,高质量 step supervision 很贵,甚至不可得。

未来可能需要:

  • synthetic traces;
  • teacher model distillation;
  • process supervision;
  • RL final reward;
  • self-improvement data;
  • environment interaction traces。

论文结论里也提到,未来可研究 Stage 2 用 final-answer reward 做 RL。

#14.3 固定 memory blocks 的容量和可解释性仍不清楚

memory blocks 的 hidden states 到底编码了什么?

PCA 和 linear probe 提供了证据,但还远远不够。我们还不知道:

  • 每个 block 是否对应某类推理操作;
  • block 间是否有可复用算法结构;
  • 是否只是学会了任务分布上的 shortcut;
  • memory token 数量扩大后是否仍稳定;
  • 在更大模型上是否会出现不同 scaling law。

#14.4 和 explicit CoT 的关系还不是替代,而是互补

RiM 的 final-block 低延迟很强,但显式 CoT 在较大模型上仍然有更高上限。更现实的方向可能是 hybrid:

  • 简单题:RiM 直接 silent reasoning;
  • 不确定题:RiM + verifier 选择 block;
  • 困难题:触发显式 CoT / search / tool use;
  • agent 任务:RiM 作为每一步 action 前的内部 planning workspace。

#15. 和 wenjun 关心方向的关系

#15.1 对 latent-space reasoning 的启发

RiM 说明 latent reasoning 的关键问题不是“能不能不用文字”,而是:

如何训练一个 latent state,使它真的承担中间计算,而不是被模型忽略?

这可以抽象为三个问题:

  1. workspace design:latent computation 放在哪里?token positions、hidden states、recurrent state、external memory?
  2. supervision topology:梯度如何强迫模型把信息写入 workspace?
  3. readout/control:如何从 workspace 中读出答案、判断正确性、决定是否继续计算?

RiM 对第 2 点给了一个漂亮范例:用 attention mask + dense step prediction 防止监督绕过 workspace。

#15.2 对 model-based RL / Dreamer for LLM Agent 的启发

如果把 LLM Agent 的长轨迹看作环境交互序列,那么一个核心瓶颈是:直接在超长 token/action 轨迹上做 RL 成本很高,信用分配也困难。

RiM 提供了一个可能的类比:

  • Stage 1 像是用显式轨迹/teacher traces 学一个 latent workspace;
  • Stage 2 像是把 workspace 转向 final outcome;
  • 未来可用 RL reward 替代 final-answer NLL,使 workspace 为最终任务成功服务。

这和 Dreamer-style model-based RL 有相通点:

不一定要在外部 action/token 空间里展开所有推理,可以在 latent state 里做内部 roll-out / refinement,再读出 action。

当然,RiM 现在还不是 world model。它没有学习环境 transition,也没有多步 imagined rollout。但它证明了一个小尺度事实:固定 token workspace 可以被训练成任务相关内部计算空间。

这对 agent 方向的启发是:

  • 可以在每个 observation 后插入 latent memory blocks,让模型先做 silent planning;
  • 用 expert trajectories 或 process traces ground blocks;
  • 再用 outcome reward / verifier reward 调整 final action readout;
  • correctness/progress probe 可能成为是否继续思考或触发 tool use 的 controller。

#15.3 对代码 Agent 的启发

代码任务里,很多中间状态不一定适合自然语言表达,例如:

  • 当前 bug hypothesis;
  • symbol dependency;
  • execution trace summary;
  • patch plan;
  • test failure localization;
  • hidden constraints。

显式 CoT 会很长,而且未必忠实。RiM-style workspace 可能用于:

  1. 在读取代码上下文后插入 memory blocks;
  2. 用 repo-level traces / debugging traces 做 Stage 1;
  3. 让 block readouts 预测下一步诊断、patch intent、test expectation;
  4. Stage 2 用最终 patch correctness 或 test pass 监督。

这比“让模型写更长分析”更接近训练一个内部调试状态机。


#16. 我对这篇论文的判断

这篇论文不是那种“结果已经碾压一切”的工作,而是一个很有方向感的范式探索。它的重要性在于提出了一个清晰问题:

LLM 的推理能力是否必须通过自回归 token 序列展开?还是可以训练出一段并行可处理的内部工作记忆?

RiM 的答案是:至少在 GSM 类数学题上,可以。

我会把它放在 latent reasoning 发展脉络里的这个位置:

阶段代表思路核心限制RiM 推进点
CoT写出自然语言步骤高延迟、语言格式负担把 reasoning 从外部文字移到内部状态
Coconutcontinuous thoughts仍需自回归生成 latent states固定 memory blocks,一次 forward pass
filler/pause tokens插入无语义 tokens很难训练出有用计算dense supervision + attention mask 激活 workspace
RiM工作记忆式 latent workspace仍依赖 step traces,任务范围有限提供低延迟 latent reasoning 训练范式

最值得后续追的问题有三个:

  1. 能否用 RL 替代 Stage 2 的 final-answer NLL? 这直接连接 long-horizon agent / outcome reward。
  2. 能否把 memory blocks 变成可控 compute budget? 类似 adaptive computation:不确定就多用 blocks,确定就少用。
  3. 能否学一个 selector/verifier 从多个 block readouts 里选答案? any-block 和 linear probe 结果说明这里潜力很大。

如果未来有人把 RiM-style memory workspace 和 agentic RL、tool-use traces、world-model imagination 结合起来,它可能会从“数学题低延迟推理技巧”变成更通用的 latent planning substrate。

#17. 最后用人话总结

这篇论文可以用一个比喻讲完:

以前我们让模型做题,是让它把草稿纸一行一行念出来。CoT 是念中文草稿,Coconut 是念一种人看不见的连续向量草稿,但还是一行一行念。RiM 说:为什么一定要念?我们能不能直接在模型输入里放几格“脑内草稿纸”,训练模型把中间计算写到这些格子的 hidden states 里,最后只说答案?

它的训练方法是:

  1. 先让每格草稿纸能复原下一步推理,确保模型真的往里面写东西;
  2. 再取消逐步推理监督,让每格草稿纸都服务最终答案;
  3. 用 attention mask 防止模型偷看文字步骤;
  4. 最后得到一个几乎不增加推理延迟、但比 direct answer 和 Coconut 更强的 latent reasoning 方法。

对 latent reasoning 来说,这篇论文的价值不是“GSM8K 提了几个点”,而是把问题重新表述成:

如何训练 LLM 使用内部工作记忆,而不是把所有思考都外化成 token 序列?

这正是后续做长程 agent、潜空间规划、model-based RL for LLM 时绕不开的问题。