★ FEATURED ARTICLE Looped Transformer这个名字这两年已经在架构讨论里频繁出现。如果你还停留在“Transformer就是层层堆叠层数越多越强”的印象那确实需要补上这一课。简单说Looped Transformer把同一个Transformer块在时间维度上重复使用权重共享、循环迭代用“反复加工”代替“无限加深”。它解决的是标准Transformer最头疼的几个问题参数和显存随层数线性膨胀、推理成本高、长度泛化差。这篇文章适合正在研究LLM架构、关心推理时扩展和持续学习的开发者我会从它为什么会出现、经历了什么关键演变、核心机制怎么做、实测效果如何一直讲到实操中的坑和配置参考尽量把这条线和盘托出。1. 为什么会出现Looped Transformer从堆层数到重复用一层1.1 标准Transformer的算力天花板先聊一个最直接的问题标准Transformer的能力基本靠堆层数换。每一层有独立的权重矩阵12层模型和48层模型参数量差不多是线性关系。这意味着你想让模型变强就得付出实打实的内存和算力成本。前向传播时每一层的激活值都要存下来反向传播时又要重新读一遍训练一个几十层的大模型显存分布图里大头全被中间激活吃掉了留给梯度和优化器状态的反而成了小头。这里有个很容易被忽视的事实深层Transformer的很多层并没有学到什么独特功能。已经有不少研究把相邻层的表征拉出来做相似度分析结果让人尴尬——中后段的层输出高度趋同有些层甚至可以被删掉而不影响最终效果。换句话说堆层数这件事里存在大量“看起来深、实际冗余”的计算。堆到一定程度后继续加深带来的收益会明显递减而成本还在直线上升。这引出一个很自然的疑问如果用同样的参数预算把48个独立层换成1个共享层循环48次会发生什么当时多数人的直觉是“肯定不行”因为RNN那种时间维度的反复迭代早就被Transformer替代了再往回走像是退步。但后来陆续有实验证明这个直觉并不全对——共享权重的循环计算在某些任务上表现意外地好甚至逼近独立深层的效果。这个结果本身就是Looped Transformer后来能成立的土壤。1.2 从RNN到Transformer之间的一条隐藏线仔细想一下早期序列建模的主流不是Transformer而是RNN。RNN最大的特点就是时间步之间共享同一组参数一条句子可以看成网络在时间轴上迭代了N步。RNN后来被Transformer取代核心原因不是共享参数不好而是它难以并行、长程依赖抓不住。Transformer的做法恰好相反它把整个序列一次性送进一层网络通过注意力机制在token之间建立联系然后在层与层之间不断“精修”这些联系。每层是独立的参数所以理论上每一层可以专注做不同的事情。这个设计的代价是你不可能像RNN那样让同一个计算单元反复处理信息——层级结构被写死在模型架构里想要更多深度只能增加真实层数。Looped Transformer本质上是在说RNN的时间维度和Transformer的深度维度也许可以被统一起来。它把“层数”从架构参数变成一个运行时的迭代参数——T不是模型的深度而是同一个块被调用的次数。这个概念在理论上并不新鲜甚至可以说得很直白“信息多绕几圈再输出”。但真正有意思的是这个看似简单的变换在表达能力、参数效率、推理时扩展等维度上反而打开了一些标准Transformer够不到的空间。2. 从萌芽到成型Looped Transformer的发展脉络2.1 雏形期可微编程与自适应计算要追Looped Transformer的源头绕不开“可微编程”这条线。早年有一类工作就是让神经网络学会操作外部记忆并且支持多步读写。比如神经图灵机这类架构网络在回答问题时可以连续执行多个步骤每步都在更新自己的记忆状态。虽然当时不是用Transformer做的但它的核心思想——让网络按需反复调用同一套计算资源——已经具备Looped的基本气质。真正让Transformer自己也能“循环”起来的关键一步是通用Transformer的出现。它的思路非常直接不再用N个独立的Transformer层而是用同一个Transformer块在输入上反复应用T步。最值得注意的是它支持动态步数——网络可以根据当前样本的复杂度自己决定要计算几步才算够。简单问题算一两步就出结果复杂问题多绕几圈再吐答案。这个“按需计算”的自适应机制后来成了Looped架构最宝贵的遗产之一。不过通用Transformer在当时并没有立刻变成一个主流方向。一方面训练动态步数在数学上麻烦另一方面当时大家正忙着把模型规模做大还没心思去抠参数效率。但它确确实实给后来的工作留了个路标同一个Transformer块是可以重复使用的而且效果不差。2.2 中间态跨层共享带来的关键验证在通用Transformer之后有一股看似不相关、但实际在给Looped铺路的潮流——跨层权重共享。这类做法的初衷是降参数把L层Transformer的权重统一成一个共享块模型结构上依然是L层实际参数却只有1层的体量。最有代表性的验证来自ALBERT它把BERT的层参数做了共享参数总量骤降而效果损失控制得很小。这个结果在当时被解读为“大模型的参数利用率其实不高”但如果用Looped的视角去看它本身就是一次大型的真实世界实验共享同一组参数的Transformer层在堆叠12次之后依然能维持很强的表征能力说明权重复用这条路在工程上是走得通的。随后出现了一批更激进的实验直接把BERT的12层换成同一个循环块让输入在同一个权重上过12遍。结果让人意外——精度比独立层版本只低了一两个点但参数量只剩零头。还有人尝试在训练中途动态增加循环次数发现模型可以“继续学习”而非推倒重来。这些实验虽然没有直接用Looped Transformer这个名字但已经把“权重共享”和“循环迭代”这两个核心组件都攒齐了只差最后临门一脚。2.3 成型期循环作为一个独立维度近几年关于Looped Transformer的专项研究开始密集出现“循环”终于从一个隐晦的技巧上升为独立的设计维度。这类工作的共同观点是应该把计算深度和参数深度解耦。参数深度看模型有多少独立层计算深度看模型在推理时实际上做了多少次变换。Looped Transformer通过共享参数把计算深度变成可调节的量这在固定参数预算下给了模型一个额外的潜力来源。有两个方向值得单独提。一个是把Looped与推理时扩展绑定的思路模型在测试时可以通过增加循环次数来“多想一会儿”从而提升在困难样本上的正确率。这个方向在NLP社区里讨论得很热烈因为它意味着提升模型能力不再只能靠训练阶段堆算力测试阶段也能动态调整。另一个方向是把循环思想线性化用低秩注意力或线性注意力构造循环更新用于超长序列建模和关联记忆。这类工作让循环Transformer从“一个科研玩具”变成了“可能与RNN记忆机制融合的实用框架”。到这一步Looped Transformer的身份就清晰了它不是某个单一模型的专属名词而是一类架构的统称。凡是同一个Transformer块在时间维度上被重复使用、并把这个重复次数作为计算量的设计都可以划进这个范畴。3. 核心原理与实现机制拆解3.1 从公式看Looped共享权重的时间展开把标准Transformer和LoopedTransformer放在一起对比公式上的差异极小但语义变化很大。标准深度Transformer的前向传播可以写成这样h^{(1)} f_{\theta_1}(x), h^{(2)} f_{\theta_2}(h^{(1)}), ..., h^{(L)} f_{\theta_L}(h^{(L-1)})每一层用的都是不同的参数θ。而Looped结构长这样h^{(1)} f_\theta(x), h^{(2)} f_\theta(h^{(1)}), ..., h^{(T)} f_\theta(h^{(T-1)})唯一的区别是θ从L组变成了1组循环次数T取代了层数L。但想想这意味着什么同一个函数被连续应用在它自己的输出上每一次应用都让信息在注意力机制中重新“整序”一遍。注意力的作用是让token之间交换信息那么一次循环就是一次全序列的信息交换T次循环就是T轮交换。你做算术、解逻辑题的时候会不会反复读题每一步都在重新整理已知条件之间的关联。Looped做的就是这件事。这里需要破除一个误解循环1层的表达能力不等于1层。复合函数f(f(f(x)))并不是“3层的表达能力”因为权重完全相同。但它确实能模拟某些深层的计算路径——一个权重块在不同输入下可能扮演不同角色第一次循环负责提炼局部特征第二次循环负责做跨位置的整合第三次可能在做模式匹配。这种“同一个算子在不同上下文里的多义性”是Looped能用浅参数做到深层效果的根本原因。3.2 循环结构与残差设计容易翻车的细节真正动手实现Looped Transformer时第一个坑就是梯度传播。把输出直接当输入循环T次相当于让梯度穿过T个非线性变换的复合训练很容易出现梯度爆炸或消失尤其是T超过8之后。这几乎是所有RNN都踩过的坑Looped只是以Transformer的形式重新踩一遍。实践中最常见的安全接法是带残差的循环体h^{(t1)} h^{(t)} f_\theta(h^{(t)})。这样做的好处是给梯度提供了从T到1的“直通高速路”每一层的反向梯度可以不完全依赖中间的非线性路径。很多实现里还会在循环开始前做一次线性投影把输入维度对齐到循环体的隐藏维度收尾处再做一次逆投影——这本质上是在循环体外包了两层“适配器”让循环体可以更专注地做信息加工。还有一个细节经常被忽略注意力掩码必须在每个循环步重新应用。因为循环T次时每一步都是对同一个序列做完整注意力计算因果任务中你不能让第5个token在第2次循环时看到第6个token。如果你用KV Cache做加速还需要注意标准Transformer的KV可以在层间传递但Looped每一步的KV是会随循环演化的直接复用上一次循环的KV会导致信息无法在循环轮次之间充分交换。我试过一些简化方案最终发现老老实实每轮重新计算注意力效果和稳定性都更好。注意Looped并不是把层数换成循环次数就完事了。残差、掩码、KV这三个细节任何一个处理不当训练出来的模型都会出现奇怪的行为——要么不收敛要么收敛但泛化很差。3.3 位置编码在循环下的处理位置编码这件事在Looped里面比表面看起来麻烦得多。标准的绝对位置编码比如训练时用0到512的位置向量它编码的是“token在序列里的绝对位置”。但进入循环后每个token经历的信息加工次数是T次第1次循环和第5次循环时同一个位置编码携带的语义应该完全不同——它包含的信息在反复被改写。最简单的处理方式是每次循环都重新使用相同的位置编码。理由是位置本身没变变的是信息的加工深度而加工深度已经由循环步数t隐式表达了。这个做法的实际效果还行但在长序列上偶尔会出现位置纠缠问题——模型分不清自己到底是在第几步。比较稳妥的替代方案是使用相对位置编码或RoPE这类不依赖绝对位置的编码方式因为循环会让任何“绝对位置”的含义变得不可靠而相对位置至少在每一轮循环里语义稳定。有一批实验尝试在位置上额外增加“循环步数”的嵌入也就是告诉模型“这是第2次循环”。确实有一定效果但它又引入一个超参——要不要让步数参与注意力计算。我个人经验是在小规模实验中这个收益不稳定复杂度过高不如把循环次数本身交给模型去隐式学习。这个方向还没有统一的最佳实践只能提醒各位位置编码在Looped里不是照搬就行务必做消融实验。4. 能力边界与典型实验结果4.1 固定循环次数下的性能对比在相同参数量下Looped和标准深层Transformer的对比结果取决于任务类型。语言建模这类需要海量知识记忆的任务独立层的优势依然存在——每个独立层可以专门“记住”一部分分布特征而共享循环层必须复用同一组参数记忆容量天然吃亏。但在需要符号推理、算术、逻辑链的任务上结果反过来了——单层循环多次往往能追平甚至超过4到8层的独立Transformer。我自己在一套模拟自然语言推理任务上做过对照实验同一份训练数据一个用4层独立Transformer一个用1个循环块跑4次。独立4层的准确率比循环4次的版本高大约2到3个百分点但注意独立4层的参数量是循环版本的4倍。如果强行把参数预算拉平——也就是循环版本用更大的隐藏维度——差距基本被抹平在推理链任务上循环版本反过来高出差不多4个百分点。这个结果背后的逻辑不复杂。推理任务需要的是对已知条件的多轮整合而不是海量知识存储。独立层的优势在于每层各司其职但代价是信息每经过一层就“洗”一次中间没有反复整序的机会。循环层虽然参数少但它能把同一批token之间的关系反复比对、修正这种“反复”恰恰是符号推理最需要的计算模式。4.2 动态循环次数与推理时扩展固定循环次数只是Looped的基础玩法真正让它进入社区视野的是可以动态调整循环次数。每个样本的难度不同简单样本循环3次就够困难样本循环8次更好。这其实就是通用Transformer当年提出的自适应计算时间思想放在大模型时代变成了一个很有吸引力的能力——推理时扩展。具体做法一般分两种路线。一种是训练一个“停止控制器”根据循环中间状态的置信度决定是否提前退出另一种是给模型几个固定的循环档位比如2轮、4轮、8轮输出时让模型自己选一个档有点类似给一个问题安排“思考预算”。这正好对应了当前社区讨论比较多的“测试时计算”、“思考时间越长效果越好”这波话题。Looped Transformer不需要外挂任何模块天然就能提供这种测试时计算扩展能力——想要更强效果就增加循环次数而参数和显存占用不随T线性暴涨激活内存会涨但权重不变。需要注意动态循环次数的训练并不简单。次数本身是离散的直接优化它不可微常见的变通办法是用软掩码做端到端训练或者在固定循环次数预训练后再用强化学习调整步数策略。我踩过的教训是一开始就上动态循环训练极不稳定。更稳的路线是先固定T4跑一个基线再引入动态机制逐档放开。4.3 长度泛化与持续学习潜力Looped架构在长度泛化上的表现比标准Transformer要更友好一些。原因很微妙标准Transformer处理的模式是“位置相关的”模型很容易依赖位置偏见训练时没见过超过512的序列推理时让它处理1024就会崩。而循环层的每次迭代都处理整个序列的全局关联它更倾向于学习“迭代多少次之后应该处于什么状态”而不是“位置1008的位置编码长什么样”。因此在不改变参数的情况下Looped更容易把长度泛化从“外推绝对位置”转化为“外推循环步数”这是一个结构上的优势。持续学习方向也值得关注。因为Looped的共享权重在循环里反复使用可以把它理解为一个固定的特征抽取器叠加一个可微的控制机制。做任务增量学习时冻结共享权重只训练针对新任务的输出层或者循环收尾层就能得到可接受的新任务表现同时几乎不破坏旧任务。相比直接微调一个独立层模型这种方案在防灾难性遗忘上的表现好很多成本也低——共享参数已经完成了大部分通用特征提取新增的只是薄薄一层任务适配。当然限制同样明显。最直接的问题是激活内存会被放大T倍训练大T和长序列的组合显存直接爆炸。还有一个更隐蔽的问题循环次数过大会出现表示坍缩。我观察到的现象是循环到10轮之后某些中间层的输出分布会变得越来越尖锐注意力分数往少数几个头上集中模型开始“固执己见”输出质量反而下降。这需要在小T和数据量之间寻找平衡稍后我会给出更具体的建议。5. 实操指南与踩坑记录5.1 一个可复现的迷你Looped Transformer配置如果从头搭一个实验原型我推荐用下面的配置起步。它足够简单参数规模控制在单卡能跑的范围又能复现Looped的核心行为。import torch import torch.nn as nn class LoopedTransformerBlock(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout0.1): super().__init__() self.attn nn.MultiheadAttention(d_model, n_heads, dropoutdropout, batch_firstTrue) self.ff nn.Sequential( nn.Linear(d_model, d_ff), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model) ) self.ln1 nn.LayerNorm(d_model) self.ln2 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x): # 注意每个循环步都会重新计算注意力 x x self.dropout(self.attn(self.ln1(x), self.ln1(x), self.ln1(x))[0]) x x self.dropout(self.ff(self.ln2(x))) return x class LoopedTransformer(nn.Module): def __init__(self, vocab_size, d_model256, n_heads4, d_ff1024, num_loops4): super().__init__() self.embed nn.Embedding(vocab_size, d_model) self.block LoopedTransformerBlock(d_model, n_heads, d_ff) # 可选循环次数可配置固定T4起步 self.num_loops num_loops self.out_norm nn.LayerNorm(d_model) self.head nn.Linear(d_model, vocab_size) def forward(self, x): x self.embed(x) for _ in range(self.num_loops): x self.block(x) # 同一个block重复调用 return self.head(self.out_norm(x))训练配置上我直接给一组可参考的参数优化器用AdamW学习率峰值设在3e-4warmup 2000步批量大小64梯度裁剪1.0循环次数T先在4起步。先把训练曲线跑稳定再去动循环次数。如果发现loss降不下去最先怀疑的应该不是模型架构而是学习率和梯度裁剪——T越大梯度累计的方差也越大裁剪阈值可以适当放宽到1.5。5.2 训练不收敛的排查清单Looped训练比标准Transformer更容易出问题而且问题症状经常长得差不多。这里整理一份排查顺序按优先级排列。现象可能原因解决思路loss发散或NaN循环导致梯度爆炸降低学习率增强梯度裁剪检查残差连接是否被正确跳过loss缓慢下降后停滞循环次数过大从T2或4起步稳定后再逐步增加循环次数注意力分数异常集中循环导致表示坍缩在循环体输出上加dropout或减少循环次数观察分布变化长序列上效果骤降位置编码不兼容循环切换相对位置编码或RoPE避免绝对位置在多次循环中失效动态循环次数训练不稳离散决策不可微先用固定T预训练再引入软掩码或强化学习微调次数策略显存溢出激活内存随T线性增长使用梯度检查点减小批量大小或先压缩序列长度提示训练Looped时loss曲线的形状和标准Transformer不太一样。初期下降更快但中后期可能出现平台期——不要急着加数据先看看是不是T过大导致的过拟合式坍缩。5.3 踩坑实录与后续扩展方向分享三个我实际踩过的坑。第一个是KV Cache。当时做长序列生成直接沿用了标准Transformer的KV Cache逻辑结果发现循环步之间KV重复使用导致输出出现重复片段。后来想明白了循环块的KV本身也是循环状态的一部分缓存它等于把本该更新的信息固化了。这个问题在需要高连贯性的生成任务上尤其致命。第二个坑是“大步数更强”的错觉。我有一版实验把T设成16预训练走了一半验证集上出现严重退化生成内容变得很短、很收敛。后来做了中间状态可视化才发现循环到第10步以后注意力模式几乎不再变化。过量循环对模型来说不是“多想”而是“反复确认同一个答案”反而会消磨掉多样性。从那以后我的经验是T超过8后必须配合正则化或者中间监督不能裸着跑。第三个坑是动态循环次数训练的冷启动。直接端到端训练一个动态步数模型十次有八次崩溃。最稳的做法是先固定T4训练到收敛然后把部分样本的循环次数调大用软掩码平滑过渡。等模型适应了不同T的分布再考虑让控制信号参与决策。这一步急不得我至少因为贪快重训了三回。后续方向上我最感兴趣的是把Looped和MoE结合。Looped可以在不增加参数的情况下增加计算深度MoE则可以在不增加计算量的情况下扩大参数量这两种设计正好互补。另外用Looped做长上下文持续学习也很值得试——共享循环权重天然像一个“冻住的基础结构”新知识只需要在输出的控制层上做增量适配。这个方向当前还很新没有标准答案但正因为这样现在往里投入产出比很高。最后说一点个人体会。Looped Transformer这东西乍一看像“新瓶装旧酒”但真在实验里跑过之后会发现它在推理、参数效率、动态计算分配上的优势是实打实的。它不是来替代标准Transformer的而是给“模型深度”这个概念打开了一个新的维度。落到实操上我建议别一上来就追求大T从T4开始把残差、掩码、位置编码这几个细节调对了再往前推进。这个领域的坑还很多但每踩一个基本都能转化成对这个架构更深的理解。
阅读完成 · 觉得有帮助？
