[论文解读] Adaptive Loops and Memory in Transformers: Think Harder or Know More?

Markus Frey, Behzad Shomali+4 more|arXiv (Cornell University)|Mar 9, 2026

Topic Modeling被引用 0

一句话总结

该论文提出一种自适应循环 Transformer，同时具备每层迭代和门控局部/全局记忆，以研究循环（更深思考）和记忆（知道更多）如何影响数学与常识推理。循环提升数学性能；记忆有助于常识，二者结合的效果优于等 FLOP 基线。

ABSTRACT

Chain-of-thought (CoT) prompting enables reasoning in language models but requires explicit verbalization of intermediate steps. Looped transformers offer an alternative by iteratively refining representations within hidden states. This parameter efficiency comes at a cost, as looped models lack the storage capacity of deeper models which use unique weights per layer. In this work, we investigate transformer models that feature both adaptive per-layer looping, where each transformer block learns to iterate its hidden state via a learned halting mechanism, and gated memory banks, that provide additional learned storage. We find that looping primarily benefits mathematical reasoning, while memory banks help recover performance on commonsense tasks compared to parameter and FLOP matched models. Combining both mechanisms yields a model that outperforms an iso-FLOP baseline, with three times the number of layers, across math benchmarks. Analysis of model internals reveals layer specialization: early layers learn to loop minimally and access memory sparingly, while later layers do both more heavily.

研究动机与目标

• 调查自适应的每层循环是否能改善 Transformer 的推理能力。
• 检验门控记忆库是否能弥补循环模型丢失的存储容量。
• 分析循环与记忆如何相互作用影响数学与常识任务的表现。
• 刻画在使用循环和记忆时内部层级的专门化特征。

提出的方法

• 在解码器单向 Transformer 中增加自适应循环与局部/全局记忆库。
• 使用学习到的终止机制对多达 Nmax 次迭代的中间循环状态进行加权。
• 引入局部（逐层）和全局（共享）记忆 KV 库并进行门控记忆集成。
• 在一个 12 层、约 2e8 参数、训练 14B tokens 的模型上进行训练，并使用 BPB 和准确率对常识与数学基准进行评估。
• 与等参数与等 FLOP 基线进行比较，以评估循环相对于加深深度的效率。

实验结果

研究问题

• RQ1 自适应的每层循环是否相对于标准 Transformer 在数学推理任务上有提升？
• RQ2 记忆库是否能弥补循环模型容量下降带来的影响，特别是在常识任务上？
• RQ3 循环深度与记忆增长在性能与内部动态方面的相互作用如何？
• RQ4 早层与晚层在循环和记忆的使用上有何差异？

主要发现

Model	CS Acc	CS BPB	Math BPB
IsoPar	0.477	0.859	2.163
Loop-3	0.501	0.813	1.687
Loop-5	0.503	0.823	1.737
Loop-7	0.498	0.832	1.659
IsoFLOP	0.523	0.780	1.801
IsoPar-M	0.459	0.823	2.108
Mem (g0=-3)	0.472	0.810	1.619
Mem (g0=0)	0.481	0.810	1.662
Mem (g0=3)	0.511	0.794	1.616
IsoFLOP-M	0.535	0.749	1.761

• 自适应循环（Nmax=3）将数学 BPB 提升约 0.476 点（从 2.163 到 1.687），常识准确率也有温和提升（0.477→0.501）。
• 更多循环对数学的收益递减，且增加迭代可能略微降低常识表现。
• 增加局部/全局记忆进一步提升了数学与常识，说明循环与记忆具有互补效应。
• 记忆增强的循环在数学任务上常常优于等参数基线，且相对于通过加深深度增加来弥补的差距，降低了常识差距。
• 层级分析显示早层循环较少且对记忆依赖较低，晚层循环更多且更依赖记忆，表明在“更努力思考”和“知道更多”之间存在专门化。
• 循环与记忆的组合可以在拥有三分之一层数的情况下实现优于等 FLOP 基线的性能，同时在数学基准上达到具有竞争力的水平。