Looped Transformer：原理、PyTorch实现与训练排错

这篇笔记的标题来自一档叫 Neural Breakdown with AVB 的拆解式科普节目。"Breakdown"这个词用得很准，因为 Looped Transformer 恰恰是一个只看名字完全猜不出在干什么、但拆开来看结构又异常简单的模型：它把同一组 Transformer 层反复调用多次，用更大的"计算深度"换取更强的推理能力，而参数量几乎不变。

如果你之前被"循环""递归深度""隐式思考"这些词绕晕过，或者想在自己的小项目里试一下这种结构却不知道从哪儿下手，这篇文章会从概念、原理、代码实现一路写到训练排错。读完之后你应该能独立实现一个可运行的 Looped Transformer，理解它为什么难训、训练时要注意哪些开关，并且知道在什么场景下值得用它、什么场景下不值得。

1. 先把概念理清：Looped Transformer 到底"循环"了什么

1.1 标准 Transformer 的深度是静态的

先回忆一下普通 Transformer 的结构。一个 12 层的模型，输入经过 12 个结构相同但参数各异的 Block，最后接一个输出头。这里的"12"在训练开始前就定死了：模型有多少层，前向传播就走多少次，一次推理消耗多少 FLOPs 也是固定的。

这种静态深度带来一个很直接的后果：模型的思考时间不能随问题难度变化。一道简单题和一道需要多步推导的难题，走的是完全相同的计算路径。而在很多任务里，"多想几步"恰恰是性能的关键——人做算术题会反复验算，模型做多步推理其实也需要类似的过程。

静态深度还有第二个问题：想提升能力就得加深，而加深意味着参数量线性增长。175B 参数的模型再想翻倍，显存和训练成本都会直接顶到天花板。所以研究者一直在找一种方式：能不能不增加参数，只增加计算？

1.2 Looped Transformer 的核心：共享权重、重复调用

Looped Transformer 给出的答案非常直接：准备 (k) 个独立的 Transformer Block，然后把这 (k) 个 Block 组成的整段结构重复执行 (T) 次。

这里有两个数字要分清：

• (k)：唯一层数（unique layers），决定有多少份不同的参数。
• (T)：循环次数（loop count / recurrence depth），决定前向传播执行多少遍。

前向流程可以写成下面这样（伪代码）：

code复制x = Embedding(tokens) + PositionEncoding
for t in range(T):
    for block in blocks[0:k]:
        x = block(x)
logits = LMHead(Norm(x))

注意内层循环里的 block 是同一批对象，第二轮循环用的还是第一轮那 (k) 个 Block 的参数。也就是说：

• 参数量 ≈ (k) 层的量。
• 计算图的深度 ≈ (k \times T) 层。

当 (k=2)、(T=8) 时，你得到一个"参数量像 2 层、计算深度像 16 层"的模型。这就是它最吸引人的地方——把参数效率和计算深度解耦了。

1.3 和 RNN、Universal Transformer、MoE 的区别

这几个概念经常被混在一起，这里做个明确区分。

关键区别在"沿什么维度循环"。RNN 是沿着序列长度递归，第 (i) 步依赖第 (i-1) 步的隐藏状态；Looped Transformer 是沿着网络深度递归，序列长度方向仍然是一次性并行处理的。所以它保留了 Transformer 的并行优势，只是在深度方向上加了一个外层循环。

和 Universal Transformer 的差别更微妙：Universal Transformer 经典设定里每个 token 可以有不同步数（动态停机），并且每一步都会重新加一次位置编码；Looped Transformer 更常见的是整层统一循环 (T) 次，把"步数"当成一个可控的超参数甚至推理时可调的旋钮。两者思想同源，工程取舍不同。

1.4 它适合解决什么问题

从已有研究和实践经验看，这类结构比较有价值的场景有三类。

第一类是算法类任务，比如排序、查表、求最值、图上的路径计算。这类任务的特点是"解需要迭代逼近"，一次前向不够，多次迭代更自然。循环结构天然匹配这种计算形态。

第二类是推理密集型任务，特别是需要在隐空间里做多步推导的场景。与其让模型把中间步骤全部吐成 token（显式思维链，序列变长、成本升高），不如让它在隐藏状态里多转几圈（隐式思考，序列长度不变）。

第三类是端侧和参数受限场景。如果你的显存只能放下 2 层模型，但又想要 16 层的表达能力，循环是最省参数的手段。

反过来说，如果你的任务就是普通分类、常规文本生成、检索问答，用标准 Transformer 通常更省心——训练稳定、生态成熟、推理框架支持好。循环结构带来的是训练难度和工程复杂度，得确认收益能覆盖成本再上。

2. 环境准备与项目结构

2.1 软硬件与依赖版本

本文的实现只依赖 PyTorch，不引入额外框架，方便你直接嵌进现有项目。关键点在于 torch.nn.functional.scaled_dot_product_attention 这个算子，它在 PyTorch 2.0 之后才提供，如果你的环境低于这个版本，需要换成手写的注意力实现。

具体版本请按你的实际情况调整，本文示例以常见环境为例，重点演示配置思路：

bash复制# 建议在独立虚拟环境中操作，避免污染全局环境
python -m venv venv
source venv/bin/activate     # Windows 用 venv\Scripts\activate

# 安装 PyTorch，具体命令请按本机 CUDA 版本到官网选择
pip install torch --index-url https://download.pytorch.org/whl/cu121

# 验证版本与 GPU 可用性
python -c "import torch; print(torch.__version__, torch.cuda.is_available())"

如果输出里 torch.cuda.is_available() 是 False，本文的代码依然能跑，只是速度慢一些——玩具任务的规模在 CPU 上也能在几分钟内看到趋势。显存方面，本文的示例模型参数量在百万级以下，2GB 显存的显卡足够。

2.2 项目目录结构

为了避免代码散落，先规划好目录。文件名和路径在后面的代码块里会一一对应，建议保持一致：

text复制looped-transformer-demo/
├── looped_transformer/
│   ├── __init__.py
│   └── model.py          # 模型定义：RMSNorm / Attention / Block / LoopedTransformer
├── train_sort.py         # 训练脚本：排序任务 + 随机循环次数
├── eval_loops.py         # 推理脚本：对比不同循环次数下的准确率
└── requirements.txt

requirements.txt 里只需要写一行：

text复制torch>=2.0

2.3 先算一笔账：参数量和计算量

在动手写代码前，先把这笔账算清楚，后面调参时心里才有数。假设 d_model=128、n_heads=4、n_unique_layers=2、max_loops=8。

参数主要来自：词嵌入 vocab_size × d_model、位置嵌入 max_seq_len × d_model、每个 Block 的注意力和前馈层、以及最后的输出头（如果做权重绑定，输出头不额外占参数）。

计算量则与 n_unique_layers × n_loops 成正比。所以当你说"这个模型有 8 层"时，一定要说清楚是8 个唯一层还是8 次循环——这两个数字对显存和耗时的含义完全不同。我建议在配置文件里用 n_unique_layers 和 max_loops 两个字段分开命名，永远不要合并成一个 n_layers，否则三个月后你自己都分不清。

3. 核心组件拆解：一个循环块怎么写

这一节把模型拆成最小零件，逐个实现并解释。零件都不复杂，但每个都有循环结构下特有的注意点。

3.1 归一化与残差：为什么选 RMSNorm + Pre-LN

循环结构本质上等价于一个很深的网络，所以归一化的位置和形式比普通 Transformer 更敏感。

• Pre-LN（先归一化再进子层）在深层网络里比 Post-LN 更稳定，因为残差路径上是一条干净的恒等映射，梯度可以沿残差直通到底。
• RMSNorm 只做均方根缩放，不减均值、不带 bias，计算量比 LayerNorm 略小，且在很多实现里更稳。

python复制# 文件路径：looped_transformer/model.py
import torch
import torch.nn as nn
import torch.nn.functional as F


class RMSNorm(nn.Module):
    """RMSNorm：按最后一维计算均方根并缩放，不含 bias。"""

    def __init__(self, d_model: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d_model))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x 形状：(B, T, C)
        variance = x.pow(2).mean(dim=-1, keepdim=True)
        x_norm = x * torch.rsqrt(variance + self.eps)
        return self.weight * x_norm

这里 keepdim=True 不能省，否则广播维度会对不上；eps 放在 rsqrt 里面而不是外面，是为了防止方差为 0 时除零。这两个细节在自定义归一化里非常容易写错，写完建议用一个小 tensor 打印形状和数值验证一遍。

3.2 因果自注意力：用 SDPA 简化实现

注意力部分直接用 scaled_dot_product_attention，它内部会根据硬件自动选择高效实现，省掉手写 mask 和 softmax 的麻烦。同时把 causal 做成参数，方便在"语言模型"和"编码器类任务"之间切换。

python复制# 文件路径：looped_transformer/model.py（续）
class SelfAttention(nn.Module):
    """多头自注意力，causal=True 时使用因果掩码。"""

    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.0, causal: bool = True):
        super().__init__()
        assert d_model % n_heads == 0, "d_model 必须能被 n_heads 整除"
        self.causal = causal
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        # 一次性算出 Q/K/V，减少三次矩阵乘的调度开销
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)
        self.dropout = dropout

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape
        qkv = self.qkv(x)                                    # (B, T, 3C)
        qkv = qkv.view(B, T, 3, self.n_heads, self.d_head)
        qkv = qkv.permute(2, 0, 3, 1, 4)                     # (3, B, H, T, Dh)
        q, k, v = qkv[0], qkv[1], qkv[2]

        out = F.scaled_dot_product_attention(
            q, k, v,
            dropout_p=self.dropout if self.training else 0.0,
            is_causal=self.causal,
        )
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(out)

几个容易踩的点。permute 之后一定要接 contiguous() 再 view，否则会报维度不连续的错。is_causal 和显式的 attn_mask 不要同时传，语义上会冲突。dropout 只在训练时生效，所以这里用 self.training 做判断，而不是把 dropout 概率写死。

3.3 前馈网络与 LayerScale

前馈层用标准的"升维—激活—降维"。另外加上 LayerScale：在残差分支上乘一个可学习的小向量，让深层残差不至于一开始就把主干冲垮。这个技巧在很深或循环很多的模型里非常有用。

python复制# 文件路径：looped_transformer/model.py（续）
class MLP(nn.Module):
    def __init__(self, d_model: int, d_hidden: int, dropout: float = 0.0):
        super().__init__()
        self.fc1 = nn.Linear(d_model, d_hidden)
        self.fc2 = nn.Linear(d_hidden, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.dropout(self.fc2(F.gelu(self.fc1(x))))


class Block(nn.Module):
    """一个 Pre-LN 的 Transformer Block，带 LayerScale。"""

    def __init__(self, d_model: int, n_heads: int, mlp_ratio: float = 4.0,
                 dropout: float = 0.0, causal: bool = True,
                 layer_scale_init: float = 1.0):
        super().__init__()
        self.norm1 = RMSNorm(d_model)
        self.attn = SelfAttention(d_model, n_heads, dropout, causal=causal)
        self.norm2 = RMSNorm(d_model)
        self.mlp = MLP(d_model, int(d_model * mlp_ratio), dropout)
        # 残差分支的缩放系数，初始值越小越保守
        self.scale1 = nn.Parameter(torch.full((d_model,), layer_scale_init))
        self.scale2 = nn.Parameter(torch.full((d_model,), layer_scale_init))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.scale1 * self.attn(self.norm1(x))
        x = x + self.scale2 * self.mlp(self.norm2(x))
        return x

注意 layer_scale_init 的取值。小规模玩具任务可以用 1.0（相当于不加缩放），跑得快、收敛明显；但在循环次数很大（比如 32、64）或者模型较宽时，建议降到 0.1 甚至更小，否则残差流会随循环次数累积膨胀得很厉害。这不是玄学，而是循环结构里最实际的一个稳定性旋钮。

3.4 把 Block 组装成可循环单元

到这里零件就齐了。真正体现"循环"的地方在模型主类里，也就是用同一组 self.blocks 反复调用。下一节把主类完整写出来。

4. 完整实战：从零实现一个 Looped Transformer

4.1 模型主类与初始化

先看完整的主类。这个类里除了循环，还包含三个重要开关：输入注入（每轮把输入重新加回残差流）、步嵌入（区分第几轮）、梯度截断（跨轮 detach）。这三个开关后面会分别解释。

python复制# 文件路径：looped_transformer/model.py（续）
class LoopedTransformer(nn.Module):
    """
    循环 Transformer：
    - n_unique_layers 个独立 Block 被重复调用 max_loops 次
    - inject_input: 每轮开始时把输入嵌入重新注入残差流，缓解深层表示漂移
    - use_step_emb: 为每一轮加一个可学习的步嵌入，让模型知道"现在是第几轮"
    """

    def __init__(
        self,
        vocab_size: int,
        d_model: int = 128,
        n_heads: int = 4,
        n_unique_layers: int = 2,
        max_loops: int = 8,
        max_seq_len: int = 64,
        mlp_ratio: float = 4.0,
        dropout: float = 0.0,
        causal: bool = False,
        use_step_emb: bool = True,
        inject_input: bool = True,
        tie_embeddings: bool = True,
        layer_scale_init: float = 1.0,
    ):
        super().__init__()
        self.max_loops = max_loops
        self.use_step_emb = use_step_emb
        self.inject_input = inject_input

        self.tok_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(max_seq_len, d_model)
        self.emb_dropout = nn.Dropout(dropout)

        self.blocks = nn.ModuleList([
            Block(d_model, n_heads, mlp_ratio, dropout, causal, layer_scale_init)
            for _ in range(n_unique_layers)
        ])

        # 步嵌入：告诉网络当前处于第几轮循环
        self.step_emb = nn.Embedding(max_loops, d_model) if use_step_emb else None

        self.norm_f = RMSNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        if tie_embeddings:
            # 输入输出共享词表矩阵，省参数且通常有小幅收益
            self.lm_head.weight = self.tok_emb.weight

        self.apply(self._init_weights)

    @staticmethod
    def _init_weights(module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def num_params(self) -> int:
        return sum(p.numel() for p in self.parameters())

这里有个细节值得单独说：tie_embeddings=True 时，lm_head.weight 和 tok_emb.weight 指向同一个 Parameter，apply(self._init_weights) 会把它初始化两遍，结果仍然是同一份权重，不影响正确性。但如果你后续要对输出头单独设置学习率或者冻结，就要注意它和词嵌入是绑在一起的，改一个等于改两个。

4.2 前向传播：三个开关的实际作用

前向传播是整篇文章的核心，逐行看清楚：

python复制# 文件路径：looped_transformer/model.py（续）
    def forward(
        self,
        idx: torch.Tensor,
        n_loops: int | None = None,
        detach_between_loops: bool = False,
        return_trajectory: bool = False,
    ):
        B, T = idx.shape
        assert T <= self.pos_emb.num_embeddings, "序列长度超过 max_seq_len"
        n_loops = self.max_loops if n_loops is None else n_loops
        assert 1 <= n_loops <= self.max_loops, "循环次数超出训练范围"

        pos = torch.arange(T, device=idx.device)
        x0 = self.emb_dropout(self.tok_emb(idx) + self.pos_emb(pos)[None, :, :])
        x = x0

        trajectory = []
        for loop in range(n_loops):
            h = x

            # 开关一：步嵌入，让共享参数知道当前是第几轮
            if self.use_step_emb:
                step_ids = torch.full((B,), loop, dtype=torch.long, device=idx.device)
                h = h + self.step_emb(step_ids).unsqueeze(1)

            # 开关二：输入注入，每轮把原始输入重新加回残差流
            if self.inject_input:
                h = h + x0

            # 同一组 Block 反复调用，这就是"循环"
            for block in self.blocks:
                h = block(h)

            x = h

            # 开关三：跨轮 detach，用截断 BPTT 换显存
            if detach_between_loops and loop < n_loops - 1:
                x = x.detach()

            if return_trajectory:
                trajectory.append(self.lm_head(self.norm_f(x)))

        logits = self.lm_head(self.norm_f(x))
        if return_trajectory:
            return logits, trajectory
        return logits

这段代码里有三处决策，每一处都影响训练能不能收敛。

第一处：步嵌入到底要不要加。 如果完全不加，网络每轮看到的输入分布都是从同一个函数映射出来的，它有可能学成一个"恒等变换"——反正改不改输出都一样，梯度自然就消失了。加上步嵌入相当于给每一轮一个不同的"身份标识"，破坏这种退化解。实践中我发现加步嵌入的收敛速度明显更稳，代价只是多 max_loops × d_model 个参数，非常划算。

第二处：输入注入的作用。 循环多轮之后，残差流里的信息会被反复覆盖，原始输入的信息可能被"冲淡"。每轮重新注入 x0，相当于给网络一个稳定的锚点。但要注意，注入是累加的，循环 8 轮就意味着残差流里叠了 8 份 x0，量级会变大。如果发现训练不稳，可以把它改成一个带门控的写法：

python复制gate = torch.sigmoid(self.inject_gate)   # 可学习标量
h = h + gate * x0

第三处：detach 换显存。 循环 (T) 次会构建 (T) 倍深度的计算图，反向传播时所有中间激活都得留着，显存占用随循环次数线性增长。detach_between_loops=True 会把每一轮之间的梯度连接切断，只保留最后一轮的梯度回传，这就是截断的 BPTT。代价是前面几轮学不到直接信号，通常需要配合"每轮都算一次 loss"的辅助监督来弥补。这个取舍在第 5 节会详细展开。

4.3 玩具任务设计：排序任务

为了验证实现是否正确，需要一个"解需要迭代"的任务。这里选排序：给模型一个随机打乱的不重复 token 序列，让它输出升序排列。

选它的理由有三个：目标完全确定，可以精确算准确率；答案依赖全局信息，必须看完整条序列；序列越长，需要的"计算步数"越多，正好能观察循环次数的影响。

python复制# 文件路径：train_sort.py
import argparse
import random

import torch
import torch.nn.functional as F

from looped_transformer.model import LoopedTransformer


def make_batch(batch_size: int, seq_len: int, vocab_size: int, device: str):
    """构造一个排序任务 batch。

    步骤：
    1. 从 vocab_size 个 token 中随机取 seq_len 个不重复的；
    2. 打乱顺序作为输入 x；
    3. 升序排列作为标签 y。
    """
    scores = torch.rand(batch_size, vocab_size, device=device)
    subset = torch.argsort(scores, dim=1)[:, :seq_len]          # 随机子集
    perm = torch.argsort(torch.rand(batch_size, seq_len, device=device), dim=1)
    x = torch.gather(subset, 1, perm)                            # 打乱
    y = torch.sort(x, dim=1).values                              # 升序标签
    return x, y

这里必须用"随机子集 + 打乱"两步，而不是直接把 0..seq_len-1 打乱。因为如果词表固定就是那 8 个 token，模型只要输出常量序列 [0,1,...,7] 就能拿到满分，根本没学排序。加上随机子集之后，标签每次都不一样，模型必须真的比较大小。

4.4 训练脚本：随机循环次数是关键

下面是完整的训练循环。核心设计是每次 batch 随机采样一个循环次数，而不是固定跑满 max_loops。

python复制# 文件路径：train_sort.py（续）
@torch.no_grad()
def evaluate(model, n_loops, vocab_size, device, seq_len=8,
             batch_size=512, batches=8):
    model.eval()
    correct, total, exact, seen = 0, 0, 0, 0
    for _ in range(batches):
        x, y = make_batch(batch_size, seq_len, vocab_size, device)
        logits = model(x, n_loops=n_loops)
        pred = logits.argmax(dim=-1)
        correct += (pred == y).sum().item()
        total += y.numel()
        exact += (pred == y).all(dim=1).sum().item()
        seen += batch_size
    model.train()
    return correct / total, exact / seen


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--steps", type=int, default=3000)
    parser.add_argument("--batch-size", type=int, default=128)
    parser.add_argument("--seq-len", type=int, default=8)
    parser.add_argument("--vocab-size", type=int, default=16)
    parser.add_argument("--d-model", type=int, default=128)
    parser.add_argument("--n-heads", type=int, default=4)
    parser.add_argument("--n-unique-layers", type=int, default=2)
    parser.add_argument("--max-loops", type=int, default=8)
    parser.add_argument("--lr", type=float, default=3e-4)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--device", type=str,
                        default="cuda" if torch.cuda.is_available() else "cpu")
    args = parser.parse_args()

    torch.manual_seed(args.seed)
    random.seed(args.seed)

    model = LoopedTransformer(
        vocab_size=args.vocab_size,
        d_model=args.d_model,
        n_heads=args.n_heads,
        n_unique_layers=args.n_unique_layers,
        max_loops=args.max_loops,
        max_seq_len=args.seq_len,
        dropout=0.0,
        causal=False,          # 排序任务需要看到全序列，用双向注意力
        use_step_emb=True,
        inject_input=True,
        tie_embeddings=False,  # 任务里输入输出语义不同，不绑定更稳
    ).to(args.device)

    print(f"参数量: {model.num_params() / 1e6:.2f}M")
    print(f"唯一层数: {args.n_unique_layers}, 最大循环次数: {args.max_loops}")

    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.01)

    for step in range(1, args.steps + 1):
        x, y = make_batch(args.batch_size, args.seq_len, args.vocab_size, args.device)

        # 关键：随机循环次数，让同一组参数适应不同的计算预算
        n_loops = random.randint(1, args.max_loops)

        logits = model(x, n_loops=n_loops)
        loss = F.cross_entropy(logits.reshape(-1, args.vocab_size), y.reshape(-1))

        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()

        if step % 200 == 0:
            acc, exact = evaluate(model, args.max_loops, args.vocab_size,
                                  args.device, args.seq_len)
            print(f"step {step:5d} | loss {loss.item():.4f} | "
                  f"loops {n_loops} | token_acc {acc:.3f} | seq_acc {exact:.3f}")

    # 训练结束后，对比不同循环次数下的表现
    print("\n不同循环次数下的准确率：")
    for loops in range(1, args.max_loops + 1):
        acc, exact = evaluate(model, loops, args.vocab_size, args.device, args.seq_len)
        print(f"  loops={loops:2d} | token_acc {acc:.3f} | seq_acc {exact:.3f}")


if __name__ == "__main__":
    main()

clip_grad_norm_(..., 1.0) 这一行不要删。循环结构等价于深层网络，梯度范数在早期很容易突然飙高，一次没裁剪就可能导致 loss 变成 nan，前功尽弃。

random.randint(1, max_loops) 这一步也不是装饰。固定跑满循环次数，模型会过度依赖"每一轮都在干有用的事"；随机采样相当于对每一条训练样本施加了不同的计算预算，逼着模型在更少的循环里也给出合理答案。这个技巧通常叫做 loop dropout 或随机深度循环，是训练能否稳定收敛的关键之一。

4.5 运行与结果观察

执行命令：

bash复制python train_sort.py --steps 3000 --batch-size 128 --seq-len 8 --max-loops 8

启动后先会看到模型规模：

text复制参数量: 0.35M
唯一层数: 2, 最大循环次数: 8

之后每 200 步打印一次进度。由于机器、随机种子和超参不同，具体数值会有差异，下面的输出只是示意，用来描述变化趋势，请以你实际运行的结果为准：

text复制step   200 | loss 2.0143 | loops 5 | token_acc 0.312 | seq_acc 0.002
step   400 | loss 1.2437 | loops 2 | token_acc 0.587 | seq_acc 0.031
step   800 | loss 0.6231 | loops 7 | token_acc 0.812 | seq_acc 0.184
step  1600 | loss 0.2104 | loops 3 | token_acc 0.946 | seq_acc 0.612
step  2400 | loss 0.0872 | loops 8 | token_acc 0.981 | seq_acc 0.846
step  3000 | loss 0.0413 | loops 6 | token_acc 0.993 | seq_acc 0.923

观察这段日志，有两点值得注意。

第一，token_acc 远高于 seq_acc。这是排序任务的正常现象：逐个 token 判断"这个位置放对了没"很容易凑对几个，但整句 8 个位置全对才算真正排序成功。评估这类任务时如果只看 token 级准确率，会严重高估模型能力，一定要看整句准确率。

第二，训练过程中不同 loops 的 loss 交错出现。这说明模型确实在同时适应不同的循环次数，而不是只记住某一个固定值。如果你发现前几百步 loss 几乎不动，可以优先检查是不是循环次数固定、或者步嵌入没生效导致表示坍缩。

4.6 推理：用循环次数当"计算预算"

训练完成之后，最有趣的实验是固定参数、改变循环次数，看效果怎么变。这正是循环结构相对普通 Transformer 的独特之处——同一个模型文件，可以在推理时按需分配算力。

python复制# 文件路径：eval_loops.py
import torch
from looped_transformer.model import LoopedTransformer
from train_sort import make_batch


@torch.no_grad()
def eval_at_loops(model, n_loops, vocab_size, device, seq_len=8,
                  batch_size=512, batches=16):
    model.eval()
    correct, total, exact, seen = 0, 0, 0, 0
    for _ in range(batches):
        x, y = make_batch(batch_size, seq_len, vocab_size, device)
        pred = model(x, n_loops=n_loops).argmax(dim=-1)
        correct += (pred == y).sum().item()
        total += y.numel()
        exact += (pred == y).all(dim=1).sum().item()
        seen += batch_size
    return correct / total, exact / seen

把训练好的模型加载进来，依次跑 loops=1..8，一般会看到一条先升后平的曲线：循环太少时准确率明显偏低，说明"想得不够"；循环增加到某个值后收益趋缓，再往上加甚至可能因为数值累积而轻微下降，这就是常说的"过度思考"。这条曲线的拐点，就是你这个模型最合适的推理循环数。

这个特性在工程上很有价值。线上服务可以按请求难度分档：简单请求走 loops=2，复杂请求走 loops=8，用一套权重覆盖不同成本档位。相比部署两个不同大小的模型，显存和运维成本都要低得多。

5. 训练稳定性：循环结构最容易踩的坑

5.1 为什么它像"很深很深的网络"

从计算图的角度看，循环 8 次的 2 层模型，梯度回传的路径长度和 16 层网络是一样的。所以所有深层网络的毛病它都有：梯度消失、梯度爆炸、表示坍塌、训练初期震荡。

更麻烦的是一个循环特有的问题：退化。因为所有轮次共享同一组参数，最优解可能出现"每轮都是恒等映射"的情况——循环 100 次和循环 1 次效果一样，等于白白浪费算力。步嵌入和输入注入就是用来对抗退化的两个手段：前者让每一轮的输入分布不同，后者保证原始信息不会被反复覆盖掉。

如果你发现模型"循环次数增加了但效果完全不变"，先别怀疑实现，先怀疑退化。检查方法是打印每一轮输出的余弦相似度：

python复制logits, traj = model(x, n_loops=4, return_trajectory=True)
h1 = traj[0][0]   # 第 1 轮输出
h2 = traj[-1][0]  # 最后一轮输出
sim = torch.nn.functional.cosine_similarity(h1.reshape(-1), h2.reshape(-1), dim=0)
print(f"首末轮输出相似度: {sim.item():.4f}")

如果相似度长期逼近 1.0，说明后面几轮几乎没在做事，需要回头检查步嵌入是否被正确使用、学习率是否过小、以及 inject_input 的缩放是否过大。

5.2 随机循环次数与课程学习

前面提到的随机采样循环次数，本质是一种正则化。除此之外还有个更温和的做法：课程学习。训练前 20% 的步数只跑 loops=1，中间逐步加到 max_loops/2，最后才放开到 max_loops。

这样做的好处是早期优化目标简单、梯度路径短，模型能先学会"单轮该干什么"，再逐步学习"多轮怎么配合"。如果一上来就跑满 8 轮，很容易出现早期 loss 长时间不降、然后突然崩掉的情况。

两种策略可以组合：前期用课程学习打基础，后期用随机采样增强鲁棒性。代价是超参数变多，需要多试几组。

5.3 中间层辅助监督与 early exit

如果用了 detach_between_loops=True，前面几轮拿不到来自最终输出的梯度，学习信号会变得很弱。这时候需要给每一轮都加一个监督：

python复制logits, traj = model(x, n_loops=4, return_trajectory=True)
loss = 0.0
for i, loop_logits in enumerate(traj):
    w = 0.4 ** (len(traj) - 1 - i)    # 越靠后的轮次权重越大
    loss = loss + w * F.cross_entropy(
        loop_logits.reshape(-1, vocab_size), y.reshape(-1)
    )

这样每一轮都要对最终答案负责，前面几轮就不会"摆烂"。权重按轮次衰减，是因为早期轮次的输出本来就粗糙，要求它一步到位反而会拖累整体。

辅助监督的另一个副产品是 early exit：推理时逐轮计算置信度，一旦某个轮次的输出足够自信就提前停止，不必跑满 max_loops。这在延迟敏感的线上场景里能省下可观的算力。

5.4 显存与计算量的账

这一块容易被低估。假设 n_unique_layers=2、max_loops=16：

• 参数量：约等于 2 层模型的量。
• 前向计算量：约等于 32 层模型的量。
• 显存占用（训练）：反向传播需要保存循环中每一轮的中间激活，近似与 n_unique_layers × max_loops 成正比。

也就是说，循环带来的参数节省，是用计算时间和训练显存换来的。如果你的瓶颈在显存而不是参数量，循环反而可能更难受。三个缓解手段：

1. 打开 detach_between_loops，显存立刻降下来，代价是前面几轮学得慢，需要辅助监督补偿。
2. 降低 batch size，同时把梯度累积步数提上去，等效 batch 不变但峰值显存下降。
3. 用梯度检查点（torch.utils.checkpoint）包住每一个循环体，用额外一次前向换显存。

这三招可以叠加，具体组合按你的硬件和任务调。

6. 常见问题与排查清单

下表汇总了实现和训练循环 Transformer 时最常遇到的问题。每一项后面都有对应的展开说明，建议按顺序排查。