让Transformer自己决定何时思考、何时存信息,提升推理能力。
Adaptive Loops and Memory in Transformers: Think Harder or Know More?
- 每层自适应循环+门控记忆库,动态控制思考和存储
- 数学题准确率超同算力基线3倍,常识任务表现更好
- 适合需要深度推理的模型优化,如数学与逻辑任务
链式思维提示(CoT)要求显式写出推理步骤,而循环Transformer通过迭代更新隐藏状态实现隐式推理,但参数效率高却缺乏深层模型的存储能力。本文提出结合自适应逐层循环机制与门控记忆库的模型:每个变换器块通过学习的停止机制决定是否迭代隐藏状态,记忆库则提供额外的可学习存储空间。实验表明,循环主要提升数学推理能力,记忆库有助于恢复常识任务性能;两者结合的模型在算力相当的情况下,相比同规模基线,在数学基准上表现更优,层数达其三倍。内部分析显示层间分工:早期层循环少、调用记忆少,后期层则频繁循环与访问记忆。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。