跳到主要内容
返回时间线
arXiv来源发表:

循环 Transformer 共享记忆:五次递归下验证困惑度较同规模标准 Transformer 降低 1.12-1.82,上下文记忆减少 76-79%

相关研究与后续进展

核心概要

该工作预训练循环语言模型使其共享记忆——仅第一次递归写入键值缓存,后续递归读取该缓存并保留自身短窗口——在 150M-1B 参数规模上,其 Looped Prediction Transformer(LPT)及混合变体在五次递归时于 FineWeb-Edu 上相对同规模标准 Transformer 将验证困惑度降低 1.12-1.82,同时减少 76-79% 的上下文记忆,并通过分析发现共享记忆与局部记忆形成不同表征、后续递归主要关注共享记忆且共享记忆充当通往第一次递归的梯度高速公路。

Source-provided article image: The Surprising Effectiveness of Shared Memory in Looped Transformers
Figure 1 ·

Figure 1: Shared memory not only saves memory, it also improves quality. Left: in our Looped Prediction Transformer ( lpt ), later recursions read the KV cache written by the first recursion and keep only a short window of their own. Right: perplexity and context memory of 1 1 B-parameter models with two to five recursions, against a parameter-matched standard Transformer ( std ). Each of our models adds shared memory to the looped baseline above it: lpt to the looped Transformer lt , and h-lpt to the gated DeltaNet hybrid lt2 .

arXiv

深度剖析

提出让循环 Transformer 共享记忆的预训练方案:只有第一次递归写入键值缓存,后续递归读取该缓存并保留一个短窗口的自身记忆。 循环 Transformer 每次递归都会写入自己的键值缓存,因此记忆仍随计算量增长;推理时压缩缓存的技术会带来质量损失。该工作改为在预训练阶段就共享记忆,而非仅在推理时压缩。 摘要描述了该预训练方案与推理时技术的对比,并报告在 150M-1B 参数规模上验证。

共享记忆不仅不损失质量,反而提升质量:五次递归时混合变体在 FineWeb-Edu 上相对同规模标准 Transformer 将验证困惑度降低 1.12-1.82,同时使用少 76-79% 的上下文记忆。 为循环模型建立了新的质量-记忆前沿,说明共享记忆可以同时改善质量与记忆效率,而非以质量换取记忆节省。 摘要报告了 150M-1B 参数规模、五次递归、FineWeb-Edu 验证困惑度降低 1.12-1.82 以及上下文记忆减少 76-79% 的具体数值。

通过分析解释共享记忆为何有帮助:共享记忆与局部记忆发展出不同表征,后续递归主要关注共享记忆,且共享记忆充当通往第一次递归的梯度高速公路。 提供了机制层面的解释,将质量提升归因于表征分化、注意力分配与梯度路径,而不仅是经验性观察。 摘要称进行了广泛分析并给出上述机制性发现,但未在摘要中列出具体实验细节。

启示与展望

该结果面向循环 Transformer 语言模型:在 150M-1B 参数规模、五次递归、FineWeb-Edu 验证集上,共享记忆方案同时改善困惑度与上下文记忆占用。它适用于希望在固定记忆预算下增加递归计算量的预训练场景,以及需要长上下文推理时计算扩展的部署环境。共享记忆与局部记忆的分工、后续递归对共享记忆的注意力集中,以及共享记忆作为梯度高速公路的机制,为后续研究提供了可检验的设计方向。

摘要未列出完整实验设置、消融实验、基线细节与统计显著性,因此质量-记忆前沿的稳健性、共享记忆在不同递归次数与不同数据分布下的表现,以及梯度高速公路机制是否可推广到其他架构,仍是读者需要留意的开放问题。此外,摘要未说明短窗口的具体大小与共享缓存的实现细节,这些会影响复现与迁移。

来源