Loop Transformer:让同一组层反复思考
不再靠"堆更多层"变深,而是把一个 Transformer 块当作函数,在深度方向上循环调用。本页用 8 个可交互动画拆解它与标准多层 Transformer 的区别、关键技术细节,以及它的好处与代价。
标准 LLM(GPT / LLaMA / Qwen …)
L 层依次堆叠,每层一套独立参数 θ₁…θ_L。深度在训练时就固定了,推理时每个 token 走一遍、只走一遍。
Loop Transformer
只有 k 层组成的一个块 θ,重复执行 T 次:h_T = f_θ(f_θ(…f_θ(h₀)))。有效深度 = k × T,而参数量只与 k 有关,T 在推理时还能调。
一句话总结好处
把"参数量"和"计算深度"解耦:用更少参数达到同等深度;推理时多循环几次就能"多想一会儿",而且是在潜空间里想,不必多写 token。
01一张图看懂:堆叠 vs 循环
同样执行 8 次"层计算"(计算量相同),左边需要 8 套参数,右边只需要 2 套。
标准 Transformer 里,"第 3 层"和"第 7 层"是两个完全不同的函数;而 Loop Transformer 把"层"变成了可以重复调用的函数:隐藏状态从块的顶部出来,再回到块的底部,被同一组参数再加工一次。拖动右侧的循环次数 T,观察有效深度变了、参数量却不变。
每个方块代表一个 Transformer 层(Attention + MLP),颜色代表它使用的参数。圆点是某个 token 的隐藏状态,每经过一层就被更新一次。左右两边每层耗时相同,所以当 T = 4 时两边同时到达输出层——计算量(FLOPs)一样,参数量差 4 倍。
| 标准多层 Transformer | Loop Transformer | |
|---|---|---|
| 参数组织 | 每层独立 θ₁, θ₂, …, θ_L | 一个块 θ(k 层)被反复调用 |
| 深度 | 训练后固定 = L | 有效深度 = k × T,推理时可调 |
| 参数量 | ∝ L | ∝ k(与循环次数 T 无关) |
| 每 token 计算量 | ∝ L | ∝ k × T |
| "多想一会儿"的方式 | 只能多生成 token(思维链) | 多循环几次,在隐藏状态里迭代精炼 |
| 代表工作 | GPT、LLaMA、Qwen、DeepSeek… | Universal Transformer、Huginn、Ouro、MoR、TRM… |
02数学形式与"循环的方向"
循环发生在深度方向,而不是 RNN 那样的时间方向。
实际可用的大模型通常采用 "前奏 — 循环核心 — 尾声"(Prelude – Recurrent Block – Coda) 三段式(Huginn-3.5B 即如此):
以 Huginn-3.5B 为例:前奏 2 层、循环核心 4 层、尾声 2 层,一共只有 8 个"真实"层;当循环 r = 32 次时,有效深度为 2 + 4×32 + 2 = 132 层。标准 Transformer 可以看成 T = 1 的特例;ALBERT 是"所有层共享参数但次数固定";Universal Transformer 则在共享之外加入了自适应停止(ACT)。
03关键细节①:参数量与计算量解耦
循环把三样通常绑在一起的东西拆开了:参数量(显存、知识容量)、计算量(FLOPs、延迟)、有效深度(推理步数)。
在标准 Transformer 里,想要更深就得加层,参数和计算同步增长。Loop Transformer 中,参数量只由块的层数 k 决定,而计算量和有效深度由 k × T 决定。下面用常见的近似估算(每层约 12·d² 个参数,每 token 每层前向约 24·d² FLOPs,忽略注意力分数项与词表)来比较 "k 层 × 循环 T 次" 与 "同等有效深度 L = k·T 的标准模型":
KV Cache 按 fp16、多头注意力(无 GQA)估算:每个 token、每次层计算需存 K 和 V 两个 d 维向量。朴素实现中每一次循环都会产生新的 KV,所以 Loop 模型的 KV Cache 与等深标准模型一样大;Mixture-of-Recursions 等工作提出跨循环共享 KV 来降低这部分开销(可能以少量精度为代价)。
04关键细节②:输入注入(Input Injection)
只把块首尾相接是不够的——循环越多,输入信息反而会被"洗掉"。
为了让深度循环稳定,迭代函数通常要接近收缩映射(每一步把状态之间的距离拉近)。但收缩映射有一个数学性质:无论从哪里出发,最终都会收敛到同一个不动点。如果输入只作为初始状态 s₀ = e 出现,迭代足够多次后,不同输入的状态会坍缩到同一点——模型"忘了题目"。
解决方法是每一步都重新注入输入:st = f(st−1, e)。这时不动点 s*(e) 由输入决定,而与初始状态无关。下面用一个二维的玩具系统 s ← tanh(W·s + U·e + b)(‖W‖ = 0.75 < 1)实时演示:
彩色方块是 6 个输入 e 的位置,曲线是对应状态 s 的迭代轨迹,空心圆是最终的不动点。无注入时 6 条轨迹汇聚到同一点,输入差异被抹平;有注入时各自收敛到不同的不动点。再点"随机初始状态 s₀":有注入时无论从哪里出发,都会回到同一个 s*(e)——这就是路径无关性。
工程实现
Huginn 把上一步状态与输入嵌入拼接 [st−1; e],经线性适配器(2h → h)后送入 4 层循环核心;也可以简单相加 s + e。初始状态 s₀ 从高斯噪声采样,迫使模型学到不依赖初始化的迭代。
为什么路径无关很重要
Anil 等(2022)发现:收敛点与初始化无关的"平衡模型",在测试时增加迭代次数更能稳定提升表现;反之,状态可能在训练范围外发散。Bansal 等(2022)称注入为 recall,并用它缓解"想太多"(overthinking)。
05关键细节③:自适应计算——难的 token 多想几轮
循环结构天然支持"按需深度":每个 token 在状态收敛后就可以提前退出。
标点、虚词往往几轮就"想清楚了",而算术结果这类 token 需要更多轮。每一轮之后计算变化量 Δt(例如 ‖st − st−1‖,或相邻两轮输出分布的 KL 散度),当 Δt < ε 时该 token 停止循环。调节阈值 ε,观察平均循环次数与节省的计算:
上方每一列是一个 token,每一行是一轮循环;颜色越亮表示该轮状态变化越大,虚线格是提前退出后省下的计算。下方曲线是所选 token 的 Δt(对数坐标),虚线为阈值 ε,点击上方 token 可切换。曲线由简化模型生成,仅作示意。
四种退出机制
ACT(Universal Transformer):每步输出停止概率,累计到阈值即停;收敛判据(Huginn):相邻两步 KL 小于阈值即停,无需额外训练;学习的退出门(Ouro):训练时用熵正则化学习每步的退出分布;路由器(Mixture-of-Recursions):轻量 router 直接为每个 token 指定递归深度。
工程上的麻烦:KV Cache
某个 token 提前退出后,后面的 token 在更深的循环里做注意力时仍需要它的 K/V。常见做法是把它最后一轮的 KV 复制给后续各轮使用,或干脆让所有循环共享同一份 KV。另外每个 token 深度不同,批处理(batching)也更难做满。
06关键细节④:怎么训练——随机循环次数 + 截断反向传播
循环 32 次就等于训练一个 132 层的网络,显存和稳定性都是问题。
问题 1:如果训练时永远循环固定的 T 次,模型只会在 T 附近工作良好,推理时改变 T 反而变差。做法:每个训练步随机采样循环次数 r。Huginn 使用对数正态-泊松分布(均值 r̄ = 32,σ = 0.5),多数落在十几到六十几次之间,并带一条长尾:
问题 2:对 r 次循环完整反向传播,激活显存随 r 线性增长。做法:截断反向传播——前面 r − k 次循环在 no_grad 下只做前向,梯度只穿过最后 k 次(Huginn 取 k = 8)。由于参数是共享的,最后 k 次得到的梯度同样更新了"每一次循环"所用的参数。
上图:P 为前奏层、C 为尾声层,中间每个方块是一次循环核心。蓝色波纹是前向传播,橙色是反向梯度;灰色方块只做前向、不保存激活。下图:累计采样到的 r 的分布(重尾,偶尔超过 60)。
归一化
深度循环会放大状态范数漂移。Huginn 发现在每个子层前后都加归一化("三明治"式 norm)是大规模稳定训练的关键,单纯 Pre-Norm 不够。
逐步监督
Bansal 等的 progressive loss:从随机步数的中间状态继续迭代再算 loss,逼模型学会"可续接的增量改进"。Ouro 则对每个循环步都计算语言模型 loss,并按学到的退出分布加权、配合熵正则。
从预训练模型改造
Relaxed Recursive Transformers 用已训练好的 LLM 中被绑定在一起的各层权重(取平均或按步长选取)初始化共享块,再给每次循环加一个小 LoRA,"放松"严格共享,以少量继续训练恢复大部分性能。
07关键细节⑤:潜空间推理 vs 显式思维链
同样是"多想一会儿":思维链在 token 空间里写草稿,循环模型在隐藏状态里打草稿。
让标准 LLM 推理更强的主流办法是思维链(CoT):生成中间步骤 token。每个中间 token 都要跑完整的 L 层,序列和 KV Cache 不断变长,而且每一步的中间结果都被压缩成一个离散 token(词表约 13 万时只有 log₂V ≈ 17 bit)。循环模型则在同一个位置上把隐藏状态反复精炼,每一轮传递的是整个高维连续向量(例如 4096 维),序列长度不变。Saunshi 等(2025)从理论上证明:循环 T 次的模型可以模拟 T 步思维链。
每生成一个中间 token = 完整跑一遍 24 层,并新增一条 KV Cache。步与步之间只能通过离散 token 传递信息。
色条表示最后一个位置的隐藏状态向量(只画 32 维),每轮循环都被整体更新,收敛后直接输出答案。
08为什么擅长"算法型"推理:每一轮 = 算法的一步
很多推理问题的本质是"把同一个步骤重复很多次"。循环结构与此天然对齐。
图搜索、动态规划、多跳推理、迭代求解……都可以写成 for t in range(T): state = step(state, input)。循环模型只需学会一个 step,然后按问题难度重复;而固定层数的模型每层可能学到不同的"特化"功能,遇到比训练时更难的实例就无法再"加层"。下面这个玩具例子假设模型学到的一步是"把可达区域向外扩一跳":
下图在 400 个随机迷宫上实时统计:横轴为推理时的循环次数,纵轴为解对的比例。假设两种模型都在"最多 6 跳"的题目上训练:标准模型深度固定,准确率是一条水平线;Loop 模型可以在推理时加循环,越过训练范围继续提升(测试时扩展)。真实模型要做到这一点,依赖 §04 的输入注入与 §06 的随机深度训练,否则可能"想过头"而发散。
学习"学习算法"
Yang 等(2023):在上下文学习线性回归等任务上,循环 Transformer 用不到 10% 的参数达到标准 Transformer 的效果,相当于学会了迭代式优化器。
长度与难度外推
Fan 等(2024)在加法、奇偶校验等任务上,让循环次数随输入长度增长,实现了明显更好的长度泛化;Schwarzschild、Bansal 等在迷宫、棋题上实现了"小题训练、大题求解"。
"推理的归纳偏置"
Saunshi 等(2025):k 层块循环 L 次,在推理任务上接近 kL 层的非共享模型;在语言建模上困惑度更差,但下游推理表现相当甚至更好。
09好处与代价
Loop 不是免费的午餐:它用计算换参数、用潜空间换可读性。
✓ 好处
参数效率
同等有效深度下参数量降为约 1/T。权重显存小,适合端侧与显存受限场景,权重可常驻高速缓存。
测试时计算可调
推理时增大 T 即可在难题上"多想",不需要输出更长文本;简单问题可以少循环,形成一个可调的"计算旋钮"。
自适应计算
按 token 难度分配循环次数(ACT、收敛退出、路由器),把算力花在真正需要的地方。
推理的归纳偏置
重复同一步骤的结构与迭代算法对齐,在多跳推理、组合泛化、长度外推上表现更好。
高带宽潜空间推理
中间状态是高维连续向量而非离散 token,不受语言表达限制,序列长度与 KV Cache 不会因推理变长。
更强的"知识运用"
Ouro 的对照实验指出:循环带来的提升并非来自更大的知识容量,而是来自对已有知识更好的组合与操控。
✗ 代价与挑战
计算量不减
每 token 的 FLOPs 与等深标准模型相当,延迟甚至更高(多次循环串行执行)。参数少 ≠ 跑得快。
知识容量受参数限制
事实性知识主要存储在参数里。在记忆型任务(如闭卷问答)上,循环模型通常不如同计算量、参数更多的标准模型。
训练稳定性
深度递归容易出现梯度爆炸/消失与状态范数漂移,需要专门的归一化、初始化、学习率和随机深度策略。
外推不保证
超出训练范围的循环次数可能让状态发散或振荡(overthinking),性能不升反降,依赖路径无关性设计。
系统工程复杂
KV Cache 随循环次数倍增;每个 token 深度不同导致难以批处理;现有推理框架大多为固定层数优化。
可解释性更差
推理过程发生在隐藏状态里,无法像思维链那样直接阅读、检查和监督。
10发展脉络
从 RNN 的自适应计算,到今天的循环语言模型。
参考文献
- Dehghani et al. Universal Transformers. ICLR 2019.
- Graves. Adaptive Computation Time for Recurrent Neural Networks. 2016.
- Lan et al. ALBERT: A Lite BERT for Self-supervised Learning of Language Representations. ICLR 2020.
- Bai, Kolter, Koltun. Deep Equilibrium Models. NeurIPS 2019.
- Schwarzschild et al. Can You Learn an Algorithm? Generalizing from Easy to Hard Problems with Recurrent Networks. NeurIPS 2021.
- Bansal et al. End-to-end Algorithm Synthesis with Recurrent Networks: Logical Extrapolation Without Overthinking. NeurIPS 2022.
- Anil et al. Path Independent Equilibrium Models Can Better Exploit Test-Time Computation. NeurIPS 2022.
- Giannou et al. Looped Transformers as Programmable Computers. ICML 2023.
- Yang et al. Looped Transformers are Better at Learning Learning Algorithms. ICLR 2024.
- Fan et al. Looped Transformers for Length Generalization. ICLR 2025.
- Bae et al. Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA. ICLR 2025.
- Saunshi et al. Reasoning with Latent Thoughts: On the Power of Looped Transformers. ICLR 2025.
- Geiping et al. Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach. 2025.
- Bae et al. Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation. 2025.
- Wang et al. Hierarchical Reasoning Model. 2025.
- Jolicoeur-Martineau. Less is More: Recursive Reasoning with Tiny Networks. 2025.
- Zhu et al. Scaling Latent Reasoning via Looped Language Models. 2025.
点击标题可在 arXiv 按标题检索原文。