[ICLR 2025 投稿/调研] Think Harder or Know More? 自适应循环与记忆库的深度协同

TL;DR

Transformer 是该堆叠层数来“死记硬背”，还是该原地循环来“深度思考”？本文提出的 Adaptive Looped Transformer 给出了一份优雅的答卷。通过**自适应循环（Looping）强化数学推理，配合门控记忆库（Memory Banks）**拟合常识知识，该模型在 200M 规模下，仅用 12 层架构就在数学任务上击败了拥有 36 层的等计算量模型。

痛点深挖：知识操纵 vs 知识容量

在 LLM 领域，CoT（思维链）被认为是显式推理的标配。但人类思考往往包含“内心独白”——即在给出答案前，在大脑内部进行多次迭代。

• 循环 Transformer 的优势：利用参数共享实现极高的逻辑压缩率，适合处理算法和数学问题。
• 致命伤：由于参数不随深度增加，模型缺乏足够的“槽位”来存储海量的常识性事实（World Knowledge）。

作者的核心直觉（Insight）是：知识操纵（Manipulation）和知识存储（Storage）应该是解耦的。 循环负责操纵，而额外的记忆模块负责存储。

核心架构：让模型自主选择“思考深度”

模型由两个核心增强组件构成：

1. 自适应循环（Adaptive Looping）

不同于传统的固定循环次数，本文利用 PonderNet 的逻辑。每一层在计算后都会接一个“停止概率”预测器。

• 动态性：对于简单的语法词汇，模型可以一次通过；对于复杂的数学逻辑，模型可以原地“思考”多次（最多 $N_{max}$ 次）。
• 平滑训练：引入了学习化的缩放参数 ${\alpha }_t$，初始状态近乎恒等映射，保证了训练早期的稳定性。

图 1: 左侧为标准 Transformer，中间为带停止机制的循环块，右侧为挂载的局部/全局记忆库。

2. 双层级记忆库（Memory Banks）

为了补齐循环带来的“容量短板”，作者塞进了两类静态记忆：

• Local Memory：每层独有，负责存储该深度特有的表征。
• Global Memory：全局共享，类似一个公共百科全书。 这些记忆通过**门控单元（Gate）**接入残差流，模型可以根据当前输入 Token 的特征，自主决定从记忆库中提取多少信息。

实验结果：数学推理的降维打击

实验在 FineWeb-Edu 数据集上进行，通过 BPB（单位字节比特数）衡量性能。

• 数学专项：在 Loop-3 配置下，模型在数学任务上的 BPB 相比 Base 模型下降了 22%。最令人吃惊的是，它比深度是其三倍的 Iso-FLOP 基线（36层）还要强。这证明了针对逻辑问题，增加迭代次数比增加参数深度更有效。
• 常识补全：单纯的 Loop 模型在常识任务上表现一般，但加入 Memory Banks 后，性能迅速回升，证明了记忆模块确实充当了“知识外挂”。

表 1: 带有记忆库的循环模型（Mem系列）在数学和常识任务上取得了最佳的平衡。

深度洞察：层级分工的涌现

通过分析内部状态，作者发现模型自发演化出了有趣的“分工”：

1. 浅层（Early Layers）：基本不循环，也很少读取记忆。主要处理基础的 Token 映射。
2. 深层（Later Layers）：循环次数显著增加，门控权重也更高。这意味着复杂的语义理解和逻辑推理主要发生在模型的后端。
3. 学习时机：模型并非一开始就学会“思考”。实验显示，只有当验证集 Loss 降到一定阈值（约 3.27）后，模型才开始通过增加循环次数来榨取性能。

图 2：随训练进行，后期图层（红色）比前期图层（蓝色）更倾向于增加循环次数。

总结与启示

这篇论文有力地反驳了“大力出奇迹”的单一路径。它告诉我们：

• 不是所有 Token 都需要同样的深度。
• 推理性能可以通过“时间换空间”的思想，用循环计算来实现。
• 静态记忆库（Memory Banks）是低参数模型扩容的有效手段。

虽然目前该研究还处于 200M 参数的实验室阶段，但其展示的“功能解耦”思想对未来端侧小模型（如 1B-3B 级模型）在算力受限下提升推理上限具有极高的参考价值。

局限性：在大规模（7B+）模型上是否存在同样的“功能解耦”效应，以及循环机制在高并发推理时的吞吐瓶颈，仍需进一步探讨。