Adaptive Loops and Memory in Transformers: Think Harder or Know More?
作者: Markus Frey, Behzad Shomali, Ali Hamza Bashir, David Berghaus, Mehdi Ali
分类: cs.CL
发布日期: 2026-03-09
备注: Published at Latent & Implicit Thinking Workshop @ ICLR 2026
💡 一句话要点
提出自适应循环与记忆Transformer，提升数学推理与常识任务性能
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: Transformer 自适应循环 门控记忆库 数学推理 常识推理 链式思考 语言模型
📋 核心要点
- 链式思考提示依赖中间步骤的显式表达，循环Transformer虽然参数效率高，但存储能力受限。
- 论文提出自适应循环与门控记忆库，使模型能够学习迭代隐藏状态并利用额外存储。
- 实验表明，循环提升数学推理，记忆库改善常识任务，二者结合优于iso-FLOP基线模型。
📝 摘要（中文）
链式思考（CoT）提示能够增强语言模型的推理能力，但需要显式地表达中间步骤。循环Transformer提供了一种替代方案，通过在隐藏状态中迭代地细化表示。这种参数效率是有代价的，因为循环模型缺乏更深层模型的存储容量，后者在每层使用独特的权重。本文研究了具有自适应逐层循环（其中每个Transformer块学习通过学习到的停止机制来迭代其隐藏状态）和门控记忆库（提供额外的学习到的存储）的Transformer模型。研究发现，循环主要有利于数学推理，而记忆库有助于恢复常识任务的性能，与参数和FLOP匹配的模型相比。结合这两种机制的模型在数学基准测试中优于iso-FLOP基线——具有三倍的层数。对模型内部结构的分析表明，层存在专业化分工：早期层学习以最小的循环和少量地访问存储器，而后期层则更多地执行这两项操作。
🔬 方法详解
问题定义：现有的大语言模型在进行复杂推理时，通常依赖于链式思考（Chain-of-Thought, CoT）提示，这需要模型显式地生成中间步骤，增加了计算负担。循环Transformer虽然通过在隐藏状态上迭代来提高参数效率，但其存储容量不如深度Transformer，限制了其在需要大量知识的任务上的表现。因此，如何提高Transformer模型的推理能力，同时保持参数效率和存储容量，是一个重要的研究问题。
核心思路：论文的核心思路是将自适应循环机制和门控记忆库结合到Transformer模型中。自适应循环允许模型根据任务的复杂程度，动态地决定每个Transformer块的迭代次数，从而提高计算效率。门控记忆库则为模型提供了额外的存储空间，使其能够更好地处理需要大量知识的任务。通过结合这两种机制，模型可以在保持参数效率的同时，提高推理能力和知识存储能力。
技术框架：该模型基于标准的Transformer架构，并在每个Transformer块中引入了自适应循环机制和门控记忆库。自适应循环机制通过一个学习到的停止机制来控制每个Transformer块的迭代次数。门控记忆库则由一组可学习的向量组成，模型可以通过门控机制来选择性地读取和写入记忆库中的内容。整体流程如下：输入序列首先经过嵌入层，然后进入一系列Transformer块。每个Transformer块首先执行自适应循环，然后通过门控机制访问记忆库。最后，模型的输出经过线性层和softmax函数，得到最终的预测结果。
关键创新：论文的关键创新在于将自适应循环机制和门控记忆库结合到Transformer模型中。自适应循环机制允许模型动态地调整计算量，从而提高计算效率。门控记忆库则为模型提供了额外的存储空间，使其能够更好地处理需要大量知识的任务。此外，论文还对模型内部结构进行了分析，发现不同层学习到了不同的功能，例如早期层主要负责特征提取，而后期层则主要负责推理和知识检索。
关键设计：自适应循环机制通过一个sigmoid函数来计算停止概率，当停止概率大于一个阈值时，循环停止。门控记忆库使用一个门控机制来控制对记忆库的读取和写入，门控值由当前Transformer块的隐藏状态计算得到。损失函数包括交叉熵损失和正则化项，用于防止过拟合。具体的参数设置和网络结构细节在论文中有详细描述。
🖼️ 关键图片
📊 实验亮点
实验结果表明，结合自适应循环和门控记忆库的模型在数学推理和常识任务上都取得了显著的提升。在数学基准测试中，该模型优于iso-FLOP基线模型，后者具有三倍的层数。在常识任务上，该模型也恢复了与参数和FLOP匹配的模型相当的性能。这些结果表明，自适应循环和门控记忆库是提高Transformer模型推理能力和知识存储能力的有效方法。
🎯 应用场景
该研究成果可应用于需要复杂推理和大量知识的自然语言处理任务，例如问答系统、机器翻译、文本摘要等。通过提高模型的推理能力和知识存储能力，可以显著提升这些任务的性能。此外，该研究还可以为设计更高效、更强大的语言模型提供新的思路。
📄 摘要（原文）
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 -- on 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.