The Surprising Effectiveness of Shared Memory in Looped Transformers

论文解决了什么问题

循环Transformer通过多次应用相同网络层来增加计算量，从而在不增加参数的前提下提升模型质量。然而，标准设定下每次递归都会写入独立的KV缓存，导致推理内存随递归深度线性增长，在长上下文中甚至超过权重内存。本文旨在探究若将内存共享作为架构归纳偏置从头预训练，能否在打破内存膨胀瓶颈的同时进一步提升模型表现。

核心方法

作者提出循环预测Transformer（LPT），规定仅第一次递归写入完整的全上下文KV缓存作为共享内存。后续的预测递归在注意力计算中直接读取该共享缓存，并将自身的局部缓存截断为一个极短的滑动窗口。这种设计将内存构建与预测任务解耦，并应用于全注意力层及混合线性注意力架构中。

与未经训练直接复用缓存的Huginn或Ouro相比，本文从头训练内存共享，避免了仅保留首次递归缓存导致的质量锐减。与并行循环Transformer（PLT）相比，LPT在保持序列递归特性的同时，在相近缓存大小下达到了更低的困惑度。与标准循环模型LT和LT2相比，LPT和H-LPT在相同递归深度下显著降低了KV内存占用并提升了下游准确率。

实验结论

• 在1B参数、5次递归设定下，H-LPT的验证困惑度比同规模标准Transformer低1.12至1.82，上下文内存大幅减少。
• LPT比同深度基线LT的困惑度最多降低1.09 nats，下游任务准确率提升最高4.8个百分点，KV内存减少达4.7倍。
• 后续递归在注意力分配中高度依赖共享内存，其注意力占比在早期预测递归中为57%至71%，在最终递归中上升至73%至80%。
• 共享内存与局部内存在表示空间（PCA聚类）中形成明显分离，表明两者在模型中发展出了不同的功能表征。
• 共享内存机制构建了从最终损失直达首次递归的梯度高速公路，恢复了因深度增加而衰减的梯度范数，使用final loss训练时表现最稳定。

局限性与证据边界

适合谁阅读