Loop Transformer
Recurrent Depth · Looped Transformer

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,观察有效深度变了、参数量却不变。

动画 1 · 隐藏状态如何流过网络
标准:独立参数
8 层
Loop:独立参数
2 层
标准:有效深度
8
Loop:有效深度 = 2 × T
8

每个方块代表一个 Transformer 层(Attention + MLP),颜色代表它使用的参数。圆点是某个 token 的隐藏状态,每经过一层就被更新一次。左右两边每层耗时相同,所以当 T = 4 时两边同时到达输出层——计算量(FLOPs)一样,参数量差 4 倍。

标准多层 TransformerLoop 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 那样的时间方向。

标准 Transformer(L 层,每层参数不同)
h(l) = Blockθl( h(l−1) ),   l = 1, …, L参数量 ∝ L
最朴素的 Loop Transformer(同一个块反复用)
h(t) = Blockθ( h(t−1) ),   t = 1, …, T参数量 ∝ 块大小 k;有效深度 = k·T

实际可用的大模型通常采用 "前奏 — 循环核心 — 尾声"(Prelude – Recurrent Block – Coda) 三段式(Huginn-3.5B 即如此):

e = P(x)前奏层:把 token 映射到潜空间(只跑一次)
s0 ~ 𝒩(0, σ2I)随机初始化循环状态
st = R( e, st−1 ),   t = 1, …, r循环核心:每一步都重新注入输入 e(见 §04)
p = C( sr )尾声层:解码出下一个 token 的概率分布(只跑一次)

以 Huginn-3.5B 为例:前奏 2 层、循环核心 4 层、尾声 2 层,一共只有 8 个"真实"层;当循环 r = 32 次时,有效深度为 2 + 4×32 + 2 = 132 层。标准 Transformer 可以看成 T = 1 的特例;ALBERT 是"所有层共享参数但次数固定";Universal Transformer 则在共享之外加入了自适应停止(ACT)。

动画 2 · 计算图展开:参数在哪个方向被共享?

横轴是序列位置(token),纵轴是深度。注意:Transformer 同一层的参数本来就在所有位置之间共享(同一行颜色相同)。Loop Transformer 新增的是纵向共享——颜色沿深度周期性重复;RNN 的循环则是横向的,一个 token 接一个 token 串行传递状态。所以 Loop Transformer 保留了 Transformer 在序列方向上的并行性。

03关键细节①:参数量与计算量解耦

循环把三样通常绑在一起的东西拆开了:参数量(显存、知识容量)、计算量(FLOPs、延迟)、有效深度(推理步数)。

在标准 Transformer 里,想要更深就得加层,参数和计算同步增长。Loop Transformer 中,参数量只由块的层数 k 决定,而计算量和有效深度由 k × T 决定。下面用常见的近似估算(每层约 12·d² 个参数,每 token 每层前向约 24·d² FLOPs,忽略注意力分数项与词表)来比较 "k 层 × 循环 T 次" 与 "同等有效深度 L = k·T 的标准模型":

交互 3 · 参数 / 计算 / KV Cache 估算器

KV Cache 按 fp16、多头注意力(无 GQA)估算:每个 token、每次层计算需存 K 和 V 两个 d 维向量。朴素实现中每一次循环都会产生新的 KV,所以 Loop 模型的 KV Cache 与等深标准模型一样大;Mixture-of-Recursions 等工作提出跨循环共享 KV 来降低这部分开销(可能以少量精度为代价)。

常见误解:Loop 不会让推理更便宜。它的价值是:同样的计算,用更少的参数完成(省显存、利于端侧);或者同样的参数,用更多计算换取更强的推理(测试时扩展)。

04关键细节②:输入注入(Input Injection)

只把块首尾相接是不够的——循环越多,输入信息反而会被"洗掉"。

为了让深度循环稳定,迭代函数通常要接近收缩映射(每一步把状态之间的距离拉近)。但收缩映射有一个数学性质:无论从哪里出发,最终都会收敛到同一个不动点。如果输入只作为初始状态 s₀ = e 出现,迭代足够多次后,不同输入的状态会坍缩到同一点——模型"忘了题目"。

解决方法是每一步都重新注入输入:st = f(st−1, e)。这时不动点 s*(e) 由输入决定,而与初始状态无关。下面用一个二维的玩具系统 s ← tanh(W·s + U·e + b)(‖W‖ = 0.75 < 1)实时演示:

动画 4 · 6 个不同输入的状态轨迹(二维相平面)
已迭代次数 t
0
6 个输入的末状态平均间距
—
结论
—

彩色方块是 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 停止循环。调节阈值 ε,观察平均循环次数与节省的计算:

动画 5 · 逐 token 提前退出(最多 24 轮)
平均循环次数
—
相比固定 24 轮节省计算
—
想得最久的 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),多数落在十几到六十几次之间,并带一条长尾:

τ ~ 𝒩( log r̄ − ½σ2, σ ),    r ~ Poisson( eτ ) + 1

问题 2:对 r 次循环完整反向传播,激活显存随 r 线性增长。做法:截断反向传播——前面 r − k 次循环在 no_grad 下只做前向,梯度只穿过最后 k 次(Huginn 取 k = 8)。由于参数是共享的,最后 k 次得到的梯度同样更新了"每一次循环"所用的参数。

动画 6 · 一个训练步:采样 r → 前向 → 截断反传
本步采样的 r
—
需保存激活的循环数
—
已采样步数 / 平均 r
—

上图: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 步思维链。

动画 7 · 计算 17 × 24:两种"思考"方式
显式思维链 CoT 在 token 空间推理
序列长度
0
层计算次数(L=24)
0

每生成一个中间 token = 完整跑一遍 24 层,并新增一条 KV Cache。步与步之间只能通过离散 token 传递信息。

循环潜空间推理 在隐藏状态里推理
序列长度
0
循环次数(块=4 层)
0

色条表示最后一个位置的隐藏状态向量(只画 32 维),每轮循环都被整体更新,收敛后直接输出答案。

但潜空间推理并不能取代思维链:它不可读,难以检查和监督;能力也受训练时循环深度的限制。实践中二者是互补的——循环模型同样可以再写思维链,相当于"每个 token 都想得更深"。

08为什么擅长"算法型"推理:每一轮 = 算法的一步

很多推理问题的本质是"把同一个步骤重复很多次"。循环结构与此天然对齐。

图搜索、动态规划、多跳推理、迭代求解……都可以写成 for t in range(T): state = step(state, input)。循环模型只需学会一个 step,然后按问题难度重复;而固定层数的模型每层可能学到不同的"特化"功能,遇到比训练时更难的实例就无法再"加层"。下面这个玩具例子假设模型学到的一步是"把可达区域向外扩一跳":

动画 8 · 从起点 S 出发能否到达终点 G?(每轮循环扩展一跳)
S → G 最短距离
—
标准模型(固定 6 层)
—
Loop 模型(T 次循环)
—

下图在 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 深度不同导致难以批处理;现有推理框架大多为固定层数优化。

可解释性更差

推理过程发生在隐藏状态里,无法像思维链那样直接阅读、检查和监督。

什么时候值得考虑 Loop 架构?参数/显存受限但可以接受更多计算(端侧、小模型);任务以多步推理、数学、算法、规划为主;希望推理时能按难度调节计算预算。若主要需求是海量知识记忆或最低延迟,标准的深而宽的模型(或 MoE)通常更合适。

10发展脉络

从 RNN 的自适应计算,到今天的循环语言模型。

2016
Adaptive Computation Time(Graves)
RNN 每步自行决定"思考"几次,是后来自适应深度的源头。
2018
Universal Transformer(Dehghani 等)
首次把 Transformer 改为深度方向共享权重的循环结构,并结合 ACT 做逐位置自适应停止。
2019
ALBERT · Deep Equilibrium Models
ALBERT 用跨层参数共享大幅压缩 BERT;DEQ 直接用求根法求"无限深"循环的不动点。
2021–2022
从易到难的外推(Schwarzschild、Bansal、Anil 等)
循环网络在小迷宫上训练、在大迷宫上求解;提出 recall(输入注入)、progressive loss 与路径无关性。
2023
Looped Transformer 的表达能力
Giannou 等:循环 Transformer 可以模拟可编程计算机;Yang 等:循环结构更擅长学习上下文学习算法。
2024
长度泛化 · Relaxed Recursive Transformers
Fan 等用自适应循环次数实现长度泛化;Bae 等把预训练 LLM 改造成递归模型并用逐层 LoRA 放松共享。
2025.02
Huginn-3.5B · Reasoning with Latent Thoughts
Geiping 等训练 3.5B 参数、8000 亿 token 的循环深度模型,推理时增加循环可达约 50B 参数模型的计算量;Saunshi 等给出循环与思维链的理论联系。
2025.06–07
HRM · Mixture-of-Recursions
HRM 用快慢两级递归模块在数独、迷宫、ARC 上取得突出结果;MoR 用路由器为每个 token 分配递归深度并共享 KV。
2025.10
Ouro(LoopLM)· TRM
字节 Seed 的 Ouro 1.4B / 2.6B 在 7.7T token 上预训练循环语言模型,论文报告可比肩最高约 12B 的标准模型;TRM 仅用约 7M 参数的 2 层网络递归,在 ARC-AGI 上超过许多大模型。

参考文献

  1. Dehghani et al. Universal Transformers. ICLR 2019.
  2. Graves. Adaptive Computation Time for Recurrent Neural Networks. 2016.
  3. Lan et al. ALBERT: A Lite BERT for Self-supervised Learning of Language Representations. ICLR 2020.
  4. Bai, Kolter, Koltun. Deep Equilibrium Models. NeurIPS 2019.
  5. Schwarzschild et al. Can You Learn an Algorithm? Generalizing from Easy to Hard Problems with Recurrent Networks. NeurIPS 2021.
  6. Bansal et al. End-to-end Algorithm Synthesis with Recurrent Networks: Logical Extrapolation Without Overthinking. NeurIPS 2022.
  7. Anil et al. Path Independent Equilibrium Models Can Better Exploit Test-Time Computation. NeurIPS 2022.
  8. Giannou et al. Looped Transformers as Programmable Computers. ICML 2023.
  9. Yang et al. Looped Transformers are Better at Learning Learning Algorithms. ICLR 2024.
  10. Fan et al. Looped Transformers for Length Generalization. ICLR 2025.
  11. Bae et al. Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA. ICLR 2025.
  12. Saunshi et al. Reasoning with Latent Thoughts: On the Power of Looped Transformers. ICLR 2025.
  13. Geiping et al. Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach. 2025.
  14. Bae et al. Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation. 2025.
  15. Wang et al. Hierarchical Reasoning Model. 2025.
  16. Jolicoeur-Martineau. Less is More: Recursive Reasoning with Tiny Networks. 2025.
  17. Zhu et al. Scaling Latent Reasoning via Looped Language Models. 2025.

点击标题可在 arXiv 按标题检索原文。