The Surprising Effectiveness of Shared Memory in Looped Transformers
论文地址：https://arxiv.org/abs/2610.02383
学术交流：922230617
目录
0. 一句话总结
1. 背景：循环 Transformer 的内存困境
2. 方法：Looped Prediction Transformer（LPT）
2.1 第一次递归写共享记忆，后续递归读共享记忆
2.2 混合版本：H-LPT
2.3 关键消融设置
3. 计算效率
4. 主要结果：共享记忆不仅省，还更强
4.1 LPT vs LT
4.2 H-LPT vs LT2
5. 为什么共享记忆有效？
5.1 共享记忆让深度递归更稳定
5.2 预测递归不一定需要很长的局部记忆
5.3 后续递归主要看共享记忆
5.4 共享记忆和局部记忆学出了不同表示
5.5 共享记忆是一条 “梯度高速公路”
6. 相关工作定位
7. 局限与未来方向
8. 谈谈我的理解
8.1 循环 Transformer 是为了重复思考一件事
8.2 共享记忆是第一次思考的完整内容
8.3 后续递归通过短窗口补充局部信息，再结合共享记忆来精炼预测
0. 一句话总结
论文发现：在 Looped Transformer 中，让第一次递归写入一份共享 KV Cache，后续递归只读取这份共享记忆 + 自己的短滑动窗口，不仅大幅省内存，而且不降质量，反而提升质量。在 150M～1B 参数规模上，提出的方法（LPT / H-LPT）刷新了循环模型的 “质量-内存前沿”。
共享内存不仅节省内存，还提升质量。
- 左图：在 Looped Prediction Transformer（LPT）中，后续递归会读取第一次递归写入的 KV Cache，并且只保留自己的一个短窗口。
- 右图：1B 参数模型在 2～5 次递归下的困惑度与上下文内存，与参数量匹配的标准 Transformer（STD）对比。每个模型都在其上方对应的循环基线上加入了共享记忆：LPT 对应循环 Transformer LT，H-LPT 对应门控 DeltaNet 混合模型 LT2。
1. 背景：循环 Transformer 的内存困境
Looped Transformer（LT）会多次复用同一组 Transformer 层，每个 token 多跑几次递归。这样可以在不增加参数量的情况下增加有效计算量，理论上是用时间换质量。
但问题在于：标准 LT 的每一次递归都会写入自己的 KV Cache。于是：
参数没变多；
但上下文内存随递归次数线性增长；
长上下文时，KV Cache 甚至可能比模型权重还大；
这直接违背了循环模型 “省内存” 的初衷。
已有的推理时技巧可以复用或压缩 Cache，但通常会损失质量。于是作者提出一个问题：
如果 “不训练就能共享 Cache” 都有一定效果，那如果从预训练开始就让模型学会共享记忆，会发生什么？
2. 方法：Looped Prediction Transformer（LPT）
论文提出的核心设计非常简洁：
2.1 第一次递归写共享记忆，后续递归读共享记忆
左图在说：
1）State 行：第 1 次递归正常运行，为所有 token 写入完整的 KV Cache。这些 KV 变成蓝色圆点，也就是 “共享记忆/Shared KV”。
2）Predict 1 行：第 2 次递归不再写完整 KV。它做两件事：
- 通过虚线箭头，去读第 1 次递归留下的蓝色共享 KV；
- 自己只写最近几个 token 的 KV，也就是蓝色方块（变红），更早的变成白色，被擦除。
3）Predict 2 行：第 3 次递归也一样，继续读共享 KV，自己只保留短窗口。
最上面的 x_{t+1} 是最终预测的下一个 token，由最后一次预测递归输出。
第 1 次递归负责 “做完整笔记”，后面几次递归只 “翻完整笔记 + 记最近几行草稿”。
形式化地，对于第 t 次递归、位置 i 的 query，其注意集合为：
其中，C^(1) 表示第一次递归写入的上下文 cache，ω 是后续递归保留的局部窗口大小，默认取 64。
后续递归会把自己的局部 KV 和共享 KV 放在同一个 softmax中注意力计算。
也就是说：
第一次递归负责 “建记忆”；
后续递归负责 “用记忆做预测”；
后续递归自己的 Cache 被截断成短窗口，因此上下文内存几乎不随递归深度增长。
2.2 混合版本：H-LPT
同样的思想也可以用于混合架构。比如 LT2 使用 Gated DeltaNet + 每 5 层一个 full attention。
在 H-LPT 中：
线性注意力层保留自己的 recurrent state，大小固定；
full attention 层使用和 LPT 一样的共享内存 + 短窗口策略。
这样，full attention 的 KV Cache 不再随递归深度增长。
2.3 关键消融设置
如果设 ω=T，也就是后续递归也保留全上下文 Cache，那么 LT 和 LPT 的唯一区别就只剩：后续递归能不能读第一次递归的共享记忆。
论文用这个设置来分离 “共享记忆” 和 “短窗口” 的贡献。
3. 计算效率
训练：用 FlexAttention 块掩码实现共享 KV + 局部窗口注意力，跳过完全遮蔽块（mask block）。
解码：每个预测递归只保留最多 ω 个前序 KV，当前条目先参与注意力再进入窗口。
缓存大小：
具体开销：T=4096、ω=64、5 次递归时，注意力条目仅增加1.25%。
内存收益：全注意力 KV Cache 相比 LT 减少4.7×，节省约79%。
结论：LPT 几乎不增加计算（相比于 LT），却显著降低长上下文推理内存。
4. 主要结果：共享记忆不仅省，还更强
4.1 LPT vs LT
在相同规模、相同递归次数下：
验证困惑度降低0.38～1.09；
下游准确率提升0.4～4.8 个百分点；
KV 记忆减少2.0～4.7 倍；
推理 FLOPs 只增加不到 1%。
更重要的是，增益随递归深度增加而变大：
LT 超过 3 次递归后，困惑度不再下降；
LPT 到 5 次递归仍在下降。
这说明共享记忆让循环模型更 “可加深”。
4.2 H-LPT vs LT2
混合版本 H-LPT 也全面优于 LT2：
上下文记忆保持在 STD 的0.18～0.24 倍；
5 次递归时，困惑度比 LT2 低0.09～0.15；
下游准确率高0.8～1.2 个百分点；
上下文记忆少4.1～4.2 倍；
推理 FLOPs 只多约 0.1%。
5. 为什么共享记忆有效？
5.1 共享记忆让深度递归更稳定
没有中间监督时，LT 增加递归次数会持续变差，5 次递归甚至崩溃到 unigram 级别。
而 LPT 保持稳定，并且使用 final loss 比 adaptive loss 更好，困惑度还能再降0.16～0.50。
结论：LPT 更适合 final loss，且深度递归更稳定。
【
这个问题在我的上一篇博文中也提到了。
Looped Transformer 最新进展：从循环 MoE、缩放定律到优化器与扩散生成
Looped-DiT 指出：
- 直接更新会逐渐丢失空间信息
- 有效循环需要中间监督 + 自适应调节注意力更新强度
】
5.2 预测递归不一定需要很长的局部记忆
论文扫描预测窗口 ω 从 0 到 4096：
即使 ω=0，后续递归完全没有自己的历史，只读共享记忆 + 当前 KV，LPT 仍然优于 LT；
即使 ω=4096，LT 和 LPT 唯一区别是共享记忆，LPT 仍低0.24～1.15；
窗口大小整体影响不大，真正关键的是能不能访问共享记忆。
所以默认 ω=64 已经足够。
5.3 后续递归主要看共享记忆
论文统计了后续递归对共享 KV 的注意力占比：
较早的预测递归：57%～71%；
最后一次预测递归：73%～80%。
也就是说，后续递归虽然有自己的局部 Cache，但它们主要依赖第一次递归写下的共享记忆。
5.4 共享记忆和局部记忆学出了不同表示
论文用 PCA 比较 query/key 表示：
在 LT 中，不同递归的 key 重叠较多；
在 LPT 中，共享 key 和局部 key 形成明显不同的簇。
这说明：
共享内存和局部内存承担了不同角色；
共享内存有效并不依赖 “不同递归写出相似表示”；
第一次递归可以专门学 “怎么写记忆”，后续递归专门学 “怎么用记忆”。
这有点像 state-prediction separation：把记忆构建和下一 token 预测分开。
5.5 共享记忆是一条 “梯度高速公路”
这是论文最关键的机制解释之一。
因为后续递归会读取第一次递归写下的共享 KV，所以 final loss 的梯度可以直接回传到第一次递归，绕过中间递归。
实验发现：
在 LT 中，到达第一次递归的梯度范数随深度急剧下降；
在 LPT 中，这个梯度被大幅恢复；
5 次递归时，大部分梯度通过共享 KV Cache 回传；
如果把后续读取共享 KV 的梯度 detach 掉，前向不变，但质量明显下降，深层模型甚至训练失败。
结论：共享记忆不仅是推理时的内存优化，更是训练时帮助第一次递归学会 “该存什么” 的优化路径。
6. 相关工作定位
论文把 LPT 放在两条研究线的交叉处：
分离记忆与预测：
- 如 XLNet 的双流注意力、State-Prediction Separation Transformer、YOCO 等。
- LPT 的不同点在于：它不增加额外 stream 或层，而是在循环递归之间分工——第一次递归建记忆，后续递归读记忆并预测。
Looped Transformer 与内存共享
- 已有工作如 Huginn、Ouro、PLT 等，多把 Cache 共享视为省内存手段。
- 论文则把共享内存当作一种架构归纳偏置，并证明它还能提升质量。
- 与最接近的 PLT 相比，LPT 保持递归顺序执行，并在一个 softmax 中同时读共享和局部 Cache，效果更好。
7. 局限与未来方向
论文也坦诚列出限制：
LPT 不是 compute-optimal：它用时间换质量。5 次递归时，解码 FLOPs 是 STD 的4.3～4.8 倍，和 LT 差不多。
每个 token 都跑相同递归次数，未来可以做 per-token 自适应递归。
比较主要是参数匹配，没有训练解码 FLOPs 匹配的 STD。
只在一个语料、4096 上下文、最高 1B 参数上验证，未测长上下文和非常高级任务。
分析主要集中在 S/M 规模，训练动态分析只覆盖 detach 消融。
未来可以探索：
增加 state recursion 数量；
交替 state recursion（首次递归，保存共享记忆）和 prediction recursion（后续递归，调用共享记忆）；
更复杂的 Memory Sharing 策略；
每 token 动态选择递归次数。
8. 谈谈我的理解
8.1 循环 Transformer 是为了重复思考一件事
循环 Transformer 的设计初衷就是：用同一组参数反复处理同一份输入，相当于增加 “思考时间” 而不增加参数。
论文里也提到，LT 是 “trade time for memory”，用更多递归来换更高质量。
8.2 共享记忆是第一次思考的完整内容
第一次递归（State Recursion）写入了完整的全上下文 KV Cache。这份 Cache 包含：
输入 token 的完整信息；
第一遍前向计算后的中间表示；
可以理解为 “第一次完整读题的笔记”。
后续递归不再重新读完整输入，而是直接读这份笔记。
共享记忆帮助后续递归看到全局内容，也让 final loss 的梯度能回传到第一次递归（Sec. 5.5），有助于稳定训练（Sec. 5.1）。
8.3 后续递归通过短窗口补充局部信息，再结合共享记忆来精炼预测
后续递归：
读取共享记忆（共享记忆本身是只读的）；
同时维护一个自己的短窗口；
两者在同一个 softmax 里一起做注意力。
在同一个 softmax 里不是分别做两次注意力再把结果加起来，而是：
把共享 KV 和局部 KV拼接在一起；
当前 query 对这个拼接后的集合做一次点积；
经过同一个 softmax得到一组注意力权重；
再用这组权重对所有 value 做加权求和。
这意味着：模型是在动态分配注意力——query当前 token 时，它自己决定看多少全局记忆、看多少近期局部信息（Sec. 5.3）。
共享记忆的目的就是：不再重写完整记忆，而是拿第一次递归写好的全局笔记，结合最近一点局部信息，把当前 token 的预测改得更准，这就是 “精炼”。
