Looped Transformers 中共享内存的惊人有效性

arXiv:2610.02383v1 发布类型：新 摘要：Looped Transformers 在每个 token 上多次应用相同的层，通过增加计算量来提升质量，而无需增加参数量。然而，每次递归都会写入自己的键值缓存，因此内存仍会随着计算量增加。推理时技术可以缩小这一缓存，但会牺牲质量。我们对循环语言模型进行预训练，使其共享内存：只有第一次递归写入缓存，后续递归读取该缓存，同时保留一个属于自身的短窗口。出人意料的是，我们发现共享内存并不会牺牲质量，反而能够提升质量。在 150M–1B 参数规模下，我们的 Looped Prediction Transformer（LPT）及其混合变体为循环模型设定了新的质量-内存前沿：在进行五次递归时，该混合模型在 FineWeb-Edu 上的验证集困惑度比同规模的标准 Transformer 低 1.12–1.82，同时使用的上下文内存减少了 76%–79%。通过广泛分析，我们探究了共享内存发挥作用的原因。共享内存和局部内存会形成不同的表示，后续递归主要关注共享内存，而共享内存还充当了通向第一次递归的梯度高速公路。