Looped Transformers：参数共享与循环推理落地实践

Looped Transformers 这个词最近被问到的次数明显变多，但很多人的第一反应是“是不是又堆了一层 Transformer”。恰好相反。它做的事情是把同一组 Transformer 层反复调用若干次，用几乎不变的参数量去换一条更深的计算路径。如果你正在做小模型、长链推理、或者被“想加深网络但显存和参数扛不住”卡住，这个思路值得花时间看一遍。我下面会按实际落地的顺序拆：它到底改了什么、循环在哪几行代码里发生、最小可运行实现长什么样、训练时会撞上哪几堵墙、推理阶段怎么拿循环次数换精度，以及出问题时先查哪里。默认你有 PyTorch 基础和 Transformer 的常识，但不要求你读过相关论文。

1. 先把概念掰开：Looped Transformers 不是“更深的 Transformer”

很多人把 Looped Transformers 理解成“层数很多的 Transformer”，这个理解会让你在调参时走错方向。它的核心不是层数多，而是层被复用了。一个标准 24 层模型有 24 套独立的注意力权重和前馈权重；一个 24 次循环的 Looped 模型可能只有 1 套或几套权重，被调用 24 次。参数量差了一个数量级，计算图深度却接近。

这个区别决定了它的一系列行为：显存曲线不一样、梯度路径不一样、推理时的可调空间不一样、能表达的函数族也不一样。搞清楚这一点，后面的坑才解释得通。

1.1 参数量不变、计算深度变深，这是它最核心的取舍

Transformer 的表达能力很大程度上和深度相关。层数越多，能组合出的中间表示层级越丰富。但加深有两个现实成本：参数线性增长，以及深层网络的训练难度上升（梯度传播变差、表示容易坍塌、需要更精细的初始化和归一化）。

Looped 结构把这两件事解耦了。它假设：某些计算步骤在形式上高度相似，只是输入状态不同。既然相似，就没有必要为每一步准备一套独立参数，共用一套、反复施加即可。

打个比方，标准深层网络像一条有 24 个不同工位的流水线，每个工位干不同的活；Looped 结构像一个人在同一张桌子上反复处理同一份材料，处理 24 轮。前者能做的事更“多样”，后者能做的事更“迭代”。

这个取舍的直接后果是：

• 参数效率高：同样的参数预算下，计算深度可以拉得很深。
• 表达受限：因为每轮施加的是同一个函数，能表达的变换是“同一个映射的多次复合”，不是任意 24 个不同映射的组合。
• 训练更依赖技巧：同一个函数反复施加，容易出现表示收敛到不动点，多轮之后输出几乎不变，白算。

所以判断一个任务适不适合用 Looped 结构，关键问题不是“我的模型够不够大”，而是“我的任务里有没有大量形式相似、需要反复迭代的计算步骤”。状态跟踪、多跳检索、逐步算术、程序执行这类任务，天然贴合。而那种每一步都需要完全不同特征变换的任务，硬套循环收益不大。

1.2 和 Universal Transformer、递归网络、MoE 的分界线

几个概念容易混，这里划清楚，不然选型会跑偏。

Universal Transformer 可以看作 Looped 思路的一个早期代表：它在深度方向共享权重，并且引入了按位置动态决定计算步数的机制。Looped Transformers 更强调“固定的循环块 + 明确的循环次数”这个骨架，把“循环几次”当成一个可调的显式超参数，工程上更好控制。

递归网络（RNN 系） 也是在时间步上复用同一套参数，但它复用的是沿序列方向的隐藏状态，处理的是序列依赖。Looped Transformer 复用的是整个 Transformer block，沿的是“计算步”这个虚拟维度，序列内部仍然靠注意力做全局交互。两者复用的对象和维度都不一样。

MoE（混合专家） 走的是另一条路：参数量变大，但每次前向只激活其中一部分。它和 Looped 是正交的，甚至可以叠加——循环块内部再放 MoE 层。但两者解决的问题不同，MoE 解决的是“参数容量和计算量的解耦”，Looped 解决的是“计算深度和参数量的解耦”。

下面这张表放在选型时看更直观：

表格里最关键的一行是第三行。Universal Transformer 的“动态步数”听起来很美，但实际落地时会带来变长计算图、批次内负载不均、梯度估计偏差等一堆问题。如果你的目标是把东西跑稳，我更建议先从固定循环次数的 Looped 版本开始，把动态步数当成后续优化项。

2. 循环到底怎么发生：一次 forward 里那几行最关键

概念说清楚了，接下来看代码层面。Looped Transformer 的实现改动其实非常小，小到很多人第一次看完会觉得“就这？”。但正是这几行改动，决定了后面所有的训练现象。

理解循环的关键是三件事：哪部分权重被共享、隐状态怎么在轮次间传递、原始输入要不要每轮注回去。这三点不同组合，会得到行为差异很大的模型。

2.1 权重共享的 block 与隐状态回注

最朴素的写法就是一个 for 循环：

python复制x = embed(tokens)
for _ in range(num_loops):
    x = block(x)
logits = head(norm(x))

block 里包含注意力、前馈、残差和归一化，每轮调用的都是同一份权重。这就是最基础的 Looped 结构。

但这里有个细节值得讨论：每轮要不要把原始输入 x0 注回去。纯 for 循环的写法里，随着轮次增加，x 会逐渐远离初始输入。如果循环次数很多，模型可能“忘掉”自己在处理什么，尤其是当 block 内部有较强的收缩性时。

一种常见做法是每轮做一次输入注入：

python复制x = block(x) + inject(x0)

inject 可以是一个线性层，也可以直接是恒等映射加个缩放系数。这样每轮计算都在“当前状态”和“原始输入”之间做一次重新对齐。

这个设计不是可有可无的装饰。它直接影响模型能不能在多次循环后仍然保留输入信息。如果你发现循环次数加到 8 以上效果反而下降，先检查这里——很可能不是循环没用，而是信息在轮次间被稀释掉了。

另一个细节是残差。标准 Transformer 的残差是层内的，Looped 结构里还可以额外加一个跨轮次的残差，把上一轮的输出直接加到这一轮的结果上。这相当于给循环加了一条“慢通道”，能缓解多轮之后的表示坍塌。

2.2 循环次数是超参数，不是训练完就锁死的常量

这是 Looped 结构最有工程价值的一点：训练时的循环次数和推理时的循环次数可以不同。

原因是权重共享。既然每轮用的是同一套参数，那么训练时跑 4 轮、推理时跑 8 轮，在计算图上是完全合法的——模型没见过 8 轮的状态分布，但它有能力处理。

这个特性带来两个实用玩法：

第一，训练时随机采样循环次数。每个 batch 随机取 2 到 8 轮，让模型见过不同深度的状态分布。这样推理时你就能在一个范围内自由调节，而不会出现“训练 4 轮、推理 6 轮就崩”的情况。

第二，推理时按任务难度调。简单输入跑少几轮省延迟，复杂输入多跑几轮换精度。同一个权重文件，导出多个循环配置，部署时按场景选。

要注意的是，这个“可调”不是无成本的。循环次数拉得太高，超出训练时见过的范围，输出质量通常会先持平再下降。原因是多轮之后隐状态趋于收敛，继续算只是重复施加接近恒等的变换，同时数值误差在累积。所以调循环次数时要有监控手段，不能盲调。

监控方法很直接：记录相邻两轮输出之间的差异（比如余弦相似度或 L2 距离）。如果第 6 轮和第 7 轮的输出差异已经很小，那第 7 轮基本是白算。这个指标比看最终 loss 敏感得多，我在调循环次数时一般先看它。

3. 用 PyTorch 写一个能跑的最小实现

理论讲完，进入能跑的部分。这里给的是教学级实现，目的是让你快速验证想法，不是生产级代码。参数规模、初始化策略、训练流程都需要按你的任务改。

我建议第一次测试严格按三步走：先确认形状对、再确认梯度有、最后确认 loss 会降。三步里任何一步没过，都不要急着加循环次数或者调学习率。

3.1 环境与最小依赖

环境要求不高，能跑普通 Transformer 的机器就能跑这个：

• Python 3.9 以上
• PyTorch 2.0 以上（nn.MultiheadAttention 的 batch_first 参数在较老版本上行为不一致，建议用新版本）
• 单卡即可，CPU 也能验证形状和梯度

先确认版本：

bash复制python -c "import torch; print(torch.__version__, torch.cuda.is_available())"

如果这行报错，后面的代码不用看了，先把环境理干净。很多人卡在“模型不收敛”，最后发现是 PyTorch 版本和 CUDA 对不上，白折腾一天。

3.2 核心代码：共享层 + 循环 + 残差

先写循环块。这里用 Post-LN 的写法（先残差后归一化），因为它在小模型上更稳一些；如果你要接预训练权重，改成 Pre-LN。

python复制import torch
import torch.nn as nn

class LoopedBlock(nn.Module):
    def __init__(self, d_model=256, nhead=8, dim_ff=1024, dropout=0.1):
        super().__init__()
        self.attn = nn.MultiheadAttention(
            d_model, nhead, dropout=dropout, batch_first=True
        )
        self.ff = nn.Sequential(
            nn.Linear(d_model, dim_ff),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(dim_ff, d_model),
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.drop = nn.Dropout(dropout)

    def forward(self, x, attn_mask=None):
        a, _ = self.attn(x, x, x, attn_mask=attn_mask, need_weights=False)
        x = self.norm1(x + self.drop(a))
        x = self.norm2(x + self.drop(self.ff(x)))
        return x

再写外层，把循环和输入注入加上：

python复制class LoopedStack(nn.Module):
    def __init__(self, vocab_size=8000, d_model=256, nhead=8,
                 dim_ff=1024, num_loops=4, dropout=0.1):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.block = LoopedBlock(d_model, nhead, dim_ff, dropout)
        self.inject = nn.Linear(d_model, d_model)
        self.norm_out = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)
        self.num_loops = num_loops

    def forward(self, idx, loops=None, attn_mask=None):
        loops = loops or self.num_loops
        x0 = self.embed(idx)
        x = x0
        for _ in range(loops):
            x = self.block(x, attn_mask) + self.inject(x0)
        return self.head(self.norm_out(x))

整个模型只有一套 block 权重。你可以试着把 num_loops 从 2 改到 12，观察参数量——它一个都不变。这就是 Looped 结构的全部魔法，剩下的都是工程问题。

如果要加跨轮残差，把循环体改成：

python复制prev = x
x = self.block(x, attn_mask) + self.inject(x0) + prev

代价是数值范围会变大，归一化的位置要相应调整。

3.3 单条任务验证：形状、梯度、loss 是否正常

写完不要直接上真实数据。先用随机张量跑一遍，确认三件事。

第一，形状对不对：

python复制model = LoopedStack(vocab_size=1000, num_loops=4)
x = torch.randint(0, 1000, (2, 32))
logits = model(x)
print(logits.shape)  # 期望 torch.Size([2, 32, 1000])

第二，梯度有没有：

python复制y = torch.randint(0, 1000, (2, 32))
loss = nn.functional.cross_entropy(
    logits.view(-1, logits.size(-1)), y.view(-1)
)
loss.backward()
print(model.block.attn.in_proj_weight.grad.norm())
print(model.inject.weight.grad.norm())

这里最容易出问题的是 inject 层。如果它初始化的输出接近零，梯度也会很小，看起来像“没学到东西”。可以先把 inject 的权重初始化为小值但不是零，或者干脆先去掉这一层跑通再加回来。

第三，loss 会不会降。用随机数据当然学不到东西，但 loss 应该能从初始值往下走一点。如果完全不降，多半是学习率、初始化或者归一化位置有问题，不要急着加循环次数。

注意：验证阶段用随机数据只是为了确认管道通，不代表模型能力。有人拿随机数据跑出 loss 下降就以为模型可以了，切到真实数据又崩，问题往往在数据侧，不在模型侧。

跑通之后，再逐步加循环次数。我一般按 2、4、8 这个节奏试，每次只改一个变量。

4. 训练阶段的三个硬骨头：显存、梯度、稳定性

Looped 结构在训练上比标准 Transformer 更容易出问题，原因很直接：同一个函数被反复施加，梯度在轮次间来回传播，数值行为比层间传播更敏感。下面三块是最常见的卡点。

4.1 显存不是随循环数线性涨，但也不是不涨

很多人以为参数不变，显存就不变。错。参数只占显存的一部分，训练时的大头是激活值和中间张量。循环 8 次意味着有 8 份 block 的中间激活需要保留（用于反向传播），这部分显存是随循环数增长的。

粗略的量级关系是：显存占用大致等于“单次 block 激活 × 循环次数 + 固定开销”。所以从 4 轮加到 8 轮，激活部分大概翻倍，总显存增长会小于翻倍，但绝对不是不变。

控制显存有三个办法：

• 梯度检查点（gradient checkpointing）：把每轮 block 当成一个独立 segment，前向不保存中间激活，反向时重算。显存大降，时间增加约三成。Looped 结构非常适合这个，因为每轮结构完全相同，实现起来比异构层简单。
• 减批次：最直接，但会拖慢训练。
• 降低序列长度：激活和序列长度相关，长序列下效果明显。

如果只是想先验证想法，我建议直接开梯度检查点，别在显存上纠结。等确定方向对了，再回来优化吞吐。

4.2 梯度回传与截断：BPTT 的思路怎么迁移过来

循环结构天然带来长梯度路径。循环 16 次，梯度要穿过 16 个 block 才能回到第一轮的输入。这和 RNN 沿时间步反传是同一类问题，所以从 RNN 那边搬过来的经验大多适用。

几个实际做法：

截断反传。只对最后 k 轮做梯度回传，前面的轮次前向但不反传。这能显著降显存和计算，代价是前面的轮次学不到东西。如果循环次数很多，这是个必要手段。

中间监督。每一轮都接一个轻量输出头，都算 loss，加权求和。这样梯度不用穿太多轮就能到达参数，训练会稳很多。代价是额外的前向开销，而且每轮都要对齐目标，实现上要小心。

分层学习率。如果循环层和嵌入层、输出头一起训，循环层的学习率通常要小一些。因为它被反复施加，等效的更新幅度被放大了。

这里没有万能配置。我的经验是先不加任何技巧，用较小的循环次数（2 到 4）跑通，看 loss 曲线和梯度范数，再决定要不要上截断或中间监督。一上来就全加上，出问题根本不知道是哪一块导致的。

4.3 稳定性：LayerNorm、残差缩放、warmup 的顺序

循环结构对归一化的位置比标准 Transformer 敏感。原因是同一个归一化层被反复调用，统计量的分布会随着轮次漂移。

如果你用 Post-LN（残差后再归一化），多轮之后输出幅度可能越来越大，最后爆掉。症状是 loss 突然变成 NaN，或者梯度范数逐轮放大。

如果你用 Pre-LN（归一化后再进子层），稳定性好一些，但深堆叠时表达力会受影响。

实际做法通常是：

• 用 Pre-LN 保证稳定，必要时在残差分支上加缩放系数。
• 加 warmup，前几百步线性升温，避免早期大梯度破坏初始化。
• 梯度裁剪，阈值不用太激进，按梯度范数的实际分布定。

下面是一个最小训练循环的骨架，方便对照：

python复制opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.LambdaLR(
    opt, lambda s: min((s + 1) / 200, 1.0)
)

for step, (x, y) in enumerate(loader):
    loops = int(torch.randint(2, 9, (1,)).item())  # 训练时随机循环次数
    logits = model(x, loops=loops)
    loss = nn.functional.cross_entropy(
        logits.view(-1, logits.size(-1)), y.view(-1)
    )
    opt.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()
    scheduler.step()

随机循环次数这一行是 Looped 训练的关键技巧，别省。它让模型对循环深度不敏感，推理时才有调节空间。

5. 推理与调参：循环次数是最好用的那个旋钮

训练完之后，Looped 模型给你留了一个标准 Transformer 没有的调节维度。用得好能省不少事，用得不好会得出“这结构没用”的错误结论。

5.1 用循环次数换精度：什么任务值得多循环

判断标准不是“循环越多越好”，而是“这个任务是不是需要多步迭代”。

值得多循环的任务特征：

• 输出需要经过若干中间推导步骤，步骤之间形式相似。
• 输入里存在需要多跳关联的信息，单次注意力覆盖不到。
• 任务本身有“迭代求解”的性质，比如状态跟踪、路径搜索、逐步计算。

不值得多循环的任务特征：

• 主要是模式匹配或检索，一次前向就能拿到足够信息。
• 输出对中间计算路径不敏感。
• 任务难度主要体现在词汇覆盖和长尾知识，而不是推理步数。

实际测试时，我建议按这个顺序做：

1. 固定训练好的权重，把循环次数从 2 扫到训练上限，记录任务指标和延迟。
2. 同时记录相邻两轮输出的差异，找到“收益开始变平”的位置。
3. 在这个位置附近选部署配置，不要选指标最高但延迟翻倍的点。

第 2 步很关键。如果差异曲线很早就平了，说明模型在训练时没有学会利用多轮计算，这时候加循环纯属浪费。原因可能是训练时循环次数采样范围太窄，或者中间监督没做好。

5.2 什么时候该停：早停、ACT 和固定次数的取舍

固定循环次数最简单，延迟可预测，工程上最好处理。缺点是所有输入都付出同样的计算代价，简单输入被浪费。

动态停止（类似自适应计算时间）能让每个 token 或每个样本自己决定算几轮。理论上更高效，实际落地要考虑：

• 批次内不同样本循环次数不同，需要 mask 和变长处理，实现复杂度上升。
• 计算图变长，编译优化和显存规划更麻烦。
• 停止判断本身需要一个额外的分类头，也要训练。

如果只是做实验验证想法，固定次数完全够用。如果要做线上服务、对延迟敏感，才值得投入做动态停止。这个取舍我在项目里一般按“先固定、后动态”的顺序推进，避免一开始就陷进实现细节。

还有一个折中方案：分档。离线跑一批数据，按输入长度或复杂度分成两三档，每档配一个固定循环次数。实现简单，效果接近动态，延迟也好预估。

6. 出问题时按这个顺序查

Looped 模型的报错和普通 Transformer 有重叠，但有几类问题是它特有的。下面按我实际排查的顺序写，从最表层往最里层走。

6.1 先看形状和 mask

第一层永远是形状。循环结构里最容易错的是 mask 的维度。

nn.MultiheadAttention 的 attn_mask 形状在不同版本、不同 batch_first 设置下要求不一致。标准 Transformer 你只写一次，错了立刻发现；Looped 结构里 mask 在循环中被反复传入，如果某一轮维度对不上，报错信息可能指向循环内部，看起来像是循环本身的问题。

排查方法：

• 先跑 loops=1，确认单轮没问题。
• 再跑 loops=2，确认循环没引入新问题。
• 打印每一轮输入的 shape，确认没有被意外改变。

如果 loops=1 就错，问题和循环无关，回去查标准 Transformer 的常见错误。如果 loops=1 对、loops=2 错，问题一定在循环体内部的状态传递上。

6.2 再看梯度与 loss 曲线

形状没问题但 loss 不降，看梯度。

先看梯度范数随训练步的变化。如果一直是 0，说明某处断开了计算图，常见原因是把张量转成了 numpy 又转回来，或者用了 torch.no_grad() 包住了不该包的部分。

如果梯度范数爆炸，看是不是循环次数太多、又没有做梯度裁剪。循环结构里梯度是复合的，放大效应比层间传播更明显。

如果梯度正常但 loss 平台期很早，看两件事：学习率是不是太大（循环结构对学习率更敏感），以及 inject 层是不是把信息压没了。

这里有个容易误判的现象：loss 在前期下降很快，然后突然卡住。很多人以为是过拟合，其实往往是模型进入了表示收敛状态——多轮施加同一个函数后，隐状态趋近不动点，继续训练也学不到新东西。验证方法是看不同循环次数下的 loss，如果 4 轮和 8 轮的 loss 几乎一样，说明多出来的轮次没在起作用。

6.3 最后看循环语义是否真的被学到

这一层最隐蔽。模型能跑、loss 在降、指标也还行，但循环结构其实没被利用——它退化成了一个普通的小 Transformer，多出来的轮次只是在做接近恒等的变换。

判断方法有两个：

消融对比。固定其他条件，只改循环次数，看指标变化。如果 2 轮和 8 轮差异极小，说明循环没被学到。

轮间差异监控。记录每一轮输出的变化量。健康的循环应该有明显的逐轮变化，且变化量逐轮递减但不为零。如果从第二轮开始变化量就接近零，说明模型在第一轮就把事情做完了。

出现这种情况，通常不是结构问题，是训练配置问题：循环次数采样范围太窄、缺少中间监督、或者任务本身不需要多步计算。前两个可以调，第三个说明这个任务不适合 Looped 结构，换方向比调参更划算。

7. 什么场景值得上，什么场景别浪费时间

最后说边界，这部分比前面的实现细节更能帮你省时间。

适合的场景：小参数量但要处理多步推理的任务；需要在延迟和精度之间做连续调节的部署环境；想把“计算深度”做成一个可配置项的产品形态；研究性质的结构探索，想验证“深度是否可以用循环替代”。

不太适合的场景：任务本身是单步映射，循环没有语义基础；对训练稳定性要求极高、没有余力调参的项目；已经有成熟的大模型方案、参数预算不是瓶颈的场景；需要极强表达灵活性、每一步都要不同变换的任务。

还有一点值得说清楚：Looped 结构不是“免费加深”。它换来的深度是有结构的深度，不是任意深度。如果你的任务真的需要 24 种不同的特征变换，循环 24 次给不了你这个能力，老老实实堆层更合适。

注意：判断要不要用 Looped，先问一句“我的任务里有没有反复做同一件事的步骤”。有，再往下走；没有，别硬套。

落地时的顺序我一般是这样：先用几百条样本跑通形状和梯度，再用小数据集看 loss 会不会降，然后扫循环次数找收益拐点，最后才考虑批量化、动态停止和部署优化。跳过任何一步，后面都会以更难排查的形式还回来。

这类结构真正落地时，最该盯住的不是它能循环多少次，而是轮次之间的信息有没有被保住、梯度有没有被截断、以及多出来的计算到底换回了什么。这三个问题答不上来，循环次数调得再细也没意义。