Looped Transformers 深度循环实战：权重共享、参数效率与测试时计算扩展

Looped Transformers 这两年出现的频率越来越高，但它不太像一个能直接下载安装的项目，更像一种被反复验证的架构思路。说人话就是：把同一组 Transformer 层反复用上好几遍，用“循环次数”代替“堆更多层”，在参数量几乎不变的前提下把有效深度拉上去。Neural Breakdown with AVB 那期内容讨论的就是这件事——循环到底循环了什么，为什么共享权重还能涨点，以及在什么情况下会失效。

如果你之前只接触过标准 Transformer 那种“层数越深越好”的写法，第一反应大概率是：同一层跑两遍，参数没变，效果怎么会变？这恰恰是 Looped Transformer 最反直觉、也最值得动手验证的地方。它把“深度”从训练时写死的结构，变成了推理时可以拧的旋钮：同一个模型，跑 2 轮和跑 8 轮，表现可以完全不一样。这和当前大模型领域里“测试时计算扩展”的思路是同一条线上的东西，只不过后者靠多吐 token，前者靠多跑几遍内部层。

门槛方面要说清楚：这东西的显存胃口和同规模的标准 Transformer 差不多，甚至更小，因为唯一层数通常压得很低。真正吃资源的是串行时间，循环没法并行，跑 8 轮就是老老实实串 8 遍。所以在消费级显卡上跑一个几千万参数的小模型做验证完全可行，想拿它做高吞吐线上服务，就要先算清楚账。

下面这篇按“能跑起来、能测出差异、能判断值不值得继续投入”的顺序展开：先讲机制和它解决什么问题，再给一份可以直接跑的最小实现，然后用四组对比实验验证循环到底带来了什么，最后补上接口封装、批量实验矩阵、显存与耗时观察，以及实际踩过的坑。

1. Looped Transformers 核心能力速览

先把关键信息摆在前面，方便判断要不要继续往下看。下面这张表里的参数不是某个具体仓库的 README，而是这个架构思路本身的属性，实际数值必须按你自己的模型配置测。

理解这张表的关键在第三行。标准 Transformer 里，层数和参数量是强绑定的，加到 24 层就是 24 份参数。Looped Transformer 把这两件事拆开了：你可以只有 2 层独立的参数块，但让它们跑 6 轮，得到的有效深度接近 12 层，而参数量只有 2 层的量。省下来的参数预算可以拿去做更宽、更大的词表，或者干脆就省着。

但别急着兴奋，这个交换是有代价的。层与层之间原本是不同参数的串联，现在变成同一组参数反复作用，表达能力必然受到约束。为了让模型不至于在反复变换中把信息磨平，工业界和研究里普遍会加一个“输入再注入”的结构，也就是每一轮循环结束后，把最初的 embedding 或者某一层中间状态重新加回去。这个细节在不少实现里是成败关键，后面代码里会专门标出来。

另一个容易被忽略的点是位置编码。标准模型里每层的注意力都在同一套位置信息上做运算，循环结构里如果处理不当，多轮之后位置信号会被反复叠加放大，导致长序列上的行为变得很奇怪。稳妥的做法是在每轮开始前重置或重新注入位置信息，而不是让它自然累加。

2. 这个思路解决什么问题，适合谁

Looped Transformer 最直接的动机是参数效率。在参数量被卡死的情况下，想让模型有更深的表征能力，传统做法是加层，但这会线性增加参数和显存。循环结构提供了一个折中：用计算换参数。公开资料里有一类结论被反复提及，就是循环能显著改善长度外推和算法类任务的表现，比如排序、查表、简单的状态追踪，这些任务对“迭代次数”天然敏感，循环结构相当于把这个迭代能力显式地建进了架构里。

第二类动机是测试时计算扩展。现在主流做法是让模型输出更长的思维链，用 token 换效果，代价是解码时间线性增长、显存里要一直留着 KV Cache。循环结构走的是另一条路：不吐 token，在内部潜空间里多跑几轮。理论上同样的计算预算下，它不需要维护成倍增长的上下文，开销曲线更平缓。这也是为什么最近一批关于“递归深度”的工作会把它和推理能力联系在一起。

适合动手的人大致三类。第一类是做架构研究或者想复现论文的，需要一个小规模、可完全掌控的实验平台；第二类是参数量受限场景下的工程实践者，比如端侧、小设备、教学演示，需要在有限参数里榨出更多效果；第三类是对“深度到底意味着什么”这件事本身好奇的，想亲眼看到同一组参数跑不同轮数的差异。

不适合的场景也要说明白。如果你要的是一个能直接调接口的生产级模型，这个方向目前给不了你现成的东西，需要自己写代码、自己训练、自己评估。如果你追求极致吞吐，循环结构的串行特性会直接卡住你，批次再大也救不回来。如果你只是想把现成模型跑起来做业务，那这条线对你几乎没有价值，不如去看成熟的推理框架。

还有一条边界必须写清楚：任何涉及模型训练和内容生成的工作，都要确认训练数据的来源合法、授权清晰，涉及个人信息的数据必须做脱敏处理，生成的内容在对外发布前要做人工复核。技术验证阶段用公开数据集和自建小样本就够了，不要为了刷指标去碰来路不明的语料。

3. 环境准备与前置条件

这套东西对环境的要求其实不高，难点不在装依赖，而在于你想验证到什么程度。如果只是跑通前向、看看循环结构长什么样，一台普通笔记本就够；如果要做训练对比实验，最好有一张显存 8G 以上的独立显卡，参数量控制在千万级别以内会比较舒服。

基础环境建议用独立的虚拟环境隔离，避免和机器上其他项目的 PyTorch 版本打架。Python 版本选 3.10 或 3.11 都行，PyTorch 装带 CUDA 的版本，版本号不用死磕某一个具体数字，和你的显卡驱动匹配即可。这里不给死版本号，因为不同显卡、不同驱动对应的组合差异很大，硬写一个反而容易踩坑。

目录结构建议一开始就分开，后面做实验矩阵的时候会省很多事：

bash复制looped-lab/
├── configs/          # 实验配置
├── data/             # 训练与评测数据
├── src/
│   ├── model.py      # 循环结构实现
│   ├── train.py      # 训练入口
│   └── serve.py      # 推理服务
├── runs/             # 训练日志与检查点
└── outputs/          # 生成结果

环境自检建议按这个顺序来，别跳步：

bash复制# 1. 确认 Python 与虚拟环境
python -V

# 2. 确认 PyTorch 能识别到显卡
python -c "import torch; print(torch.__version__); print(torch.cuda.is_available()); print(torch.cuda.device_count())"

# 3. 确认显存总量，据此决定模型宽度
nvidia-smi --query-gpu=name,memory.total --format=csv

如果第二步返回的 is_available() 是 False，先别改代码，问题九成在驱动和 CUDA 版本不匹配上。这时候去看 nvidia-smi 输出的驱动版本，再对照 PyTorch 官方给出的对应关系重装，比在代码里瞎试有效得多。

显存规划上给一个粗略的经验判断：千万参数级别的小模型，用 batch size 8、序列长度 256 做训练，8G 显存通常够用；如果你把唯一层数堆到 12 层以上，同时循环次数开到 8，激活值占用的显存会明显上升，因为反传的时候要穿过所有循环轮次。这时候第一件事是降 batch size，第二件事是开梯度检查点，而不是急着换卡。

4. 最小可运行实现与启动方式

概念讲再多，不如自己写一遍。下面这份实现刻意做得很小，唯一层数、循环次数、模型宽度全部可配置，目的是让你能在几分钟内看到同一个参数块被复用多次的效果。

python复制import torch
import torch.nn as nn


class LoopedBlock(nn.Module):
    """单个 Transformer 参数块，会被循环复用多次"""

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

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


class LoopedTransformer(nn.Module):
    def __init__(
        self,
        vocab_size,
        d_model=256,
        nhead=4,
        dim_ff=1024,
        n_unique_layers=2,
        n_loops=4,
        max_len=512,
        dropout=0.1,
    ):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(max_len, d_model)
        self.blocks = nn.ModuleList(
            [LoopedBlock(d_model, nhead, dim_ff, dropout)
             for _ in range(n_unique_layers)]
        )
        self.n_loops = n_loops
        # 关键结构：每轮回注输入，防止表征在反复变换中坍缩
        self.inject = nn.Linear(d_model, d_model, bias=False)
        self.norm_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, idx, n_loops=None):
        b, t = idx.shape
        loops = self.n_loops if n_loops is None else n_loops
        pos = torch.arange(t, device=idx.device).unsqueeze(0)
        h = self.tok_emb(idx) + self.pos_emb(pos)
        x0 = h                                # 保存初始状态用于回注

        for _ in range(loops):
            for blk in self.blocks:
                h = blk(h)
            h = h + self.inject(x0)           # 每轮结束再注入一次

        return self.head(self.norm_f(h))


if __name__ == "__main__":
    model = LoopedTransformer(vocab_size=1000, n_unique_layers=2, n_loops=4)
    total = sum(p.numel() for p in model.parameters())
    print(f"参数量: {total / 1e6:.2f}M")

    x = torch.randint(0, 1000, (2, 16))
    for k in (1, 2, 4, 8):
        y = model(x, n_loops=k)
        print(f"loops={k}, 输出形状={tuple(y.shape)}")

这段代码跑起来你会发现一件很有意思的事：loops 从 1 改到 8，参数量一个字节都没变，但前向的计算量翻了八倍。这就是循环结构的全部魔力所在，也是它全部麻烦所在。

训练脚本可以长这样，重点是两处：梯度裁剪和随机循环次数。共享权重在反传时会被穿过多次，梯度范数很容易失控，裁剪几乎是必须的。

python复制import random
import torch
import torch.nn as nn
from model import LoopedTransformer  # 按实际路径调整


def sample_loops(min_loops=2, max_loops=8):
    """训练时随机采样循环次数，提升推理阶段的可调性"""
    return random.randint(min_loops, max_loops)


def train_step(model, batch, optimizer, device, max_grad_norm=1.0):
    model.train()
    idx, target = batch["input_ids"].to(device), batch["labels"].to(device)

    loops = sample_loops()
    logits = model(idx, n_loops=loops)

    loss = nn.functional.cross_entropy(
        logits.view(-1, logits.size(-1)), target.view(-1)
    )
    loss.backward()

    # 共享参数被多轮反传，梯度范数波动大，裁剪不要省
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)
    return loss.item(), loops


if __name__ == "__main__":
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = LoopedTransformer(vocab_size=1000).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
    print(f"运行设备: {device}")

启动顺序建议是：先跑 model.py 里的自检，确认前向能出结果；再拿一个极小的随机数据集跑几十步训练，确认 loss 在下降而不是变 NaN；最后才接真实数据。跳过前两步直接上全量数据，是最常见的浪费时间的做法。

5. 功能测试与效果验证

这一章是重点。循环结构到底有没有用，不能靠感觉判断，得设计能对照的实验。下面四组是我建议必做的，顺序按投入从小到大排。

5.1 参数效率对照实验

第一组要回答的问题是：在参数量对齐的情况下，循环结构和标准堆层结构谁更强。

做法很简单，准备两个模型：A 是 2 层唯一参数块加 6 轮循环，B 是 12 层不同参数的普通堆叠。把两者的参数量打印出来，你会发现 A 大概是 B 的六分之一。这不公平，所以还要再补一个模型 C，把 A 的宽度调大，让它的参数量和 B 对齐，然后再比。

评价指标用你自己的任务指标就行，语言建模用困惑度、分类任务用准确率都可以。判断标准不是“谁一定赢”，而是看两者在参数量对齐后差距有多大。如果 A 在六分之一参数的情况下能达到 B 的九成效果，这个交换就是划算的；如果差出一大截，说明你的任务不吃纯深度，可能更适合直接堆层。

常见失败原因是数据量太小。参数量少的模型在小数据上反而容易欠拟合，跑出来的结论不可信。建议先确保数据集至少能支撑几百步有效训练。

5.2 循环次数与效果的曲线

第二组是循环结构独有的实验，也是最能说明问题的。训练完之后固定模型权重，只改推理时的 n_loops，从 1 扫到 12，每个点重复几次取平均。

预期会看到一条先升后平甚至先升后降的曲线。上升段说明多轮迭代确实带来了额外的表征能力，转折点说明模型已经在该任务上饱和，再往下只是徒增计算。这个拐点的位置对你的实际使用非常关键，因为它决定了推理时该开多少轮。

判断成功与否的标准：如果曲线完全平，说明训练时模型没有学会利用多轮信息，大概率是回注结构没起作用或者训练时的循环次数太单一；如果曲线一开始就掉，通常是每轮回注把噪声也注进去了，需要调小注入层的初始化尺度。

5.3 随机循环训练

第三组验证的是训练策略。同样是上面的模型，一次训练时固定跑 4 轮，另一次训练时每步随机采 2 到 8 轮，其他条件完全一致，最后都用同一个 n_loops 序列去测。

这一步的意义在于，固定轮数训练出来的模型往往对那个特定数字过拟合，你改成其他轮数效果就崩。随机轮数相当于做了一种结构上的数据增强，模型必须学会“不管跑几轮都能给出合理结果”，推理阶段的可调性才真正可用。

失败原因通常是采样区间设得太宽，比如从 1 采到 16，短轮数那几步梯度信号太弱，长轮数那几步又容易炸。建议区间从 2 到 6 起步，稳定之后再放宽。

5.4 算法任务与长度外推

第四组针对的是这类结构最被看好的能力。设计一个需要迭代才能解的任务，比如序列反转、括号匹配、简单的排序，训练序列长度控制在 32 以内，测试时逐步拉长到 128。

如果循环真的在承担迭代计算，你会看到它在外推长度上的衰减比同参数量的标准模型平缓。这个实验不需要大模型，几百万参数的规模就能看出趋势，非常适合做快速验证。

要注意的是，长度外推还受位置编码影响很大，做这个实验的时候务必控制位置编码方案一致，否则你测出来的是位置编码的差异，不是循环结构的差异。

6. 接口封装与批量实验

单个模型跑通之后，下一步是把它变成能反复调参的东西。研究类项目的效率瓶颈往往不在训练，而在于手动改参数、手动记录结果、手动对比，来回折腾。

最省事的方式是把推理封成一个小服务，把 n_loops 暴露成请求参数，这样你可以用脚本批量打请求，把不同轮数的结果一次性收集齐。

python复制from fastapi import FastAPI
from pydantic import BaseModel
import torch

from model import LoopedTransformer

app = FastAPI()
device = "cuda" if torch.cuda.is_available() else "cpu"
model = LoopedTransformer(vocab_size=1000).to(device).eval()


class GenerateRequest(BaseModel):
    input_ids: list[int]
    n_loops: int = 4
    max_new_tokens: int = 32


@app.post("/generate")
def generate(req: GenerateRequest):
    idx = torch.tensor([req.input_ids], device=device)
    with torch.no_grad():
        logits = model(idx, n_loops=req.n_loops)
    next_id = int(torch.argmax(logits[0, -1]))
    return {"n_loops": req.n_loops, "next_token": next_id}

启动和调用：

bash复制# 启动服务，端口按需修改
uvicorn serve:app --host 127.0.0.1 --port 8000

bash复制# 换不同循环次数打同一个请求，对比返回
curl -X POST http://127.0.0.1:8000/generate \
  -H "Content-Type: application/json" \
  -d '{"input_ids": [1, 2, 3, 4], "n_loops": 6, "max_new_tokens": 16}'

批量实验不要手写循环，用配置文件驱动。建议把每次实验的关键参数记全，包括唯一层数、循环次数、学习率、随机种子、数据版本，后面复现的时候能省大量时间。

json复制{
  "exp_name": "loop_scan_v1",
  "data": "./data/train_small.jsonl",
  "seed": 42,
  "grid": {
    "n_unique_layers": [2, 4],
    "n_loops_train": [4, 8],
    "d_model": [256, 512]
  },
  "eval_loops": [1, 2, 4, 6, 8],
  "output_dir": "./runs",
  "save_last_k": 1
}

配套的跑批脚本逻辑很简单：读配置、展开网格、逐个启动训练、训练完自动在 eval_loops 上做扫描、把结果追加写进 results.csv。失败重试建议放在最外层，单次实验失败就记一笔继续下一个，不要让一个坏配置卡住整晚的机器时间。日志建议按实验名分文件，不要把几十个实验的输出混在一个终端里，出问题根本查不出来。

7. 资源占用与性能观察

这是很多人会忽略、但对循环结构特别重要的一环。因为它把参数和计算彻底解耦了，你必须同时盯住两个指标，否则很容易得出错误结论。

显存方面，训练时的主要占用来自激活值，而激活值跟循环次数是线性关系，因为反传要保留每一轮的中间结果。观察方法用这一句就够：

bash复制# 训练过程中的显存峰值观察
python -c "
import torch
print('allocated:', torch.cuda.memory_allocated() / 1024**3, 'GB')
print('reserved :', torch.cuda.memory_reserved() / 1024**3, 'GB')
print('peak     :', torch.cuda.max_memory_allocated() / 1024**3, 'GB')
"

如果在训练脚本里加一行 torch.cuda.reset_peak_memory_stats() 放在每个 epoch 开头，你能清楚看到不同 n_loops 下峰值的变化曲线。我的经验判断是，把 n_loops 从 4 提到 8，峰值显存大致会涨一半以上，涨多少取决于序列长度，序列越长涨得越明显。

耗时方面要分开看。训练时每步时间基本随循环次数线性增长，这没什么取巧空间；推理时如果用固定轮数，也是线性的。真正的坑在于你可能同时开了大的 batch 和大的轮数，两边相乘，单步时间直接失控，然后误以为是代码有问题。

降低占用的手段按优先级排：第一是开梯度检查点，用时间换显存，训练能省下相当一部分激活；第二是降序列长度，它和循环次数是乘性关系，效果好于降轮数；第三是降唯一层数，但这一项会真正损失表达能力，放最后考虑。

CPU 推理和 GPU 推理的差异在循环结构上会被放大。小模型在 CPU 上跑 8 轮可能就是几秒钟，看起来还行，但一旦你想扫 12 个轮数点、每个点跑几百条样本，时间就不可接受了。做轮数扫描实验务必放 GPU 上，功能性验证放 CPU 上没问题。另外记得每轮实验结束把进程清干净，CUDA 上下文残留会让下一轮一开始就报显存不足，看起来像是模型的问题，其实是上一个进程没退。

8. 常见问题与排查方法

下面这些是这类实验里高频出现的现象，按“现象—原因—排查—解决”整理成表，方便对号入座。

表格之外再补两个容易被误判的情况。一个是“模型不收敛”，很多时候不是循环结构的问题，而是学习率沿用了标准 Transformer 的配置。共享权重意味着同一个参数被更新多次，等效学习率偏大，起步阶段把学习率降一半往往就好了。另一个是“效果比基线差”，先确认基线是不是在同等参数量下对比的，拿 2 层循环模型去和 12 层普通模型比效果，输了是正常的，这个对比本身没有意义。

9. 最佳实践与使用建议

第一件要做的事是把最小可运行配置固定下来。一个 2 层、宽度 256、循环 4 轮的模型，加上一份几千条的小数据集，能在几分钟内跑完一轮完整流程。这套配置看着简陋，但它是你后面所有改动的参照物，没有它你无法判断某次“效果变好”是真的改进还是随机波动。建议把它单独存成一个配置文件，永远不要删。

第二件事是管住实验变量。循环结构涉及的可调项比普通模型多：唯一层数、循环次数、回注方式、位置编码处理、随机轮数区间。每一项都会影响结果，同时改两项以上，最后你根本不知道是谁起的作用。建议一次只动一个，而且固定随机种子，至少跑三个种子看方差。

第三件事是记录循环次数在训练和推理时的对应关系。这两者不一致是很常见的隐性 bug，训练跑 4 轮、推理跑 4 轮看起来对齐了，但如果你训练时用了随机采样，推理时其实有更大的调参空间，不做扫描就白白浪费了这个能力。建议在每次训练结束后自动跑一遍轮数扫描，把曲线存成图，长期积累下来能看出规律。

第四件事是关于落地判断。这个方向目前更适合当作研究和验证平台，而不是直接上生产的方案。如果你确实想在业务里用，比较务实的路线是把它当作一个可选的推理加速或参数压缩手段，先在离线评估里跑通效果对照，确认在同等质量下资源开销确实下降，再考虑接进流程。反过来，如果评估下来只是“效果差不多、速度更慢”，那就果断放弃，不要因为概念新就硬上。

最后是合规层面。训练数据要确认来源和授权，涉及用户生成内容的场景要做脱敏，模型输出的内容在对外使用前要有复核环节。这些不是套话，尤其是当你把模型接进任何面向真实用户的产品时，数据来源不清带来的风险远大于模型效果本身的那几个点差异。

10. 收尾

回到最初那个问题：同一组层跑多遍，凭什么效果会变好。跑完上面这几组实验，答案会具体很多——变好的部分来自迭代本身带来的计算深度，变差的部分来自权重共享带来的表达约束，两者的平衡点在哪里，取决于你的任务、你的数据量、你调的回注方式和位置编码处理。这些都没有统一的数字答案，只能自己测。

如果你准备动手，建议的顺序是：先跑通那份最小实现，确认轮数变化确实改变了输出；再做参数效率对照，判断这个方向对你的任务有没有意义；有意义再去做轮数扫描曲线，找到拐点；最后才考虑接口封装和批量实验。跳过前面几步直接上服务，最后大概率是在调一个根本没配好的东西。

最容易踩的三个坑分别是：忘了梯度裁剪导致训练发散、把固定轮数训练出来的模型拿去随便改轮数、以及在序列长度和循环次数上同时加码把显存撑爆。这三个都不难避免，但每一个都能让人白花一晚上。

值得继续挖的方向也不少，比如把循环用在检索增强的中间层、把轮数和任务难度做动态绑定、用早停机制让简单样本少跑几轮。这些都还是开放的，拿一个小模型就能试。