全循环Transformer：简单稳定循环

# 通过全局循环稳定循环变压器 来源：https://arxiv.org/html/2605.18797 傅饶 香港浸会大学 csraofu@comp\.hkbu\.edu\.hk &杨子轩 吉林大学 zxyang25@mails\.jlu\.edu\.cn &张建坤 吉林大学 jiankun24@mails\.jlu\.edu\.cn &马靖 香港浸会大学 majing@comp\.hkbu\.edu\.hk &陈贺昌 吉林大学 chenhc@jlu\.edu\.cn &李宇 吉林大学 liyu90@jlu\.edu\.cn &张毅 吉林大学 yichang@jlu\.edu\.cn ###### 摘要 缩放模型性能通常需要增加模型规模。循环变压器（Looped Transformer）提供了一种引人注目的替代方案，通过迭代复用相同的Transformer模块，以额外计算换取性能提升，而不增加参数数量或上下文长度。由于循环迭代次数可以在推理时调整，它还提供了一种自然的机制来平衡性能和测试时计算量。然而，当循环迭代次数增加时，循环变压器仍然面临训练不稳定的问题。我们的分析揭示，这种不稳定性源于两个原因：梯度振荡和残差爆炸。为了解决这两个问题，我们提出了全局循环变压器（Fully Looped Transformer），引入了两种无参数的修改：（1）全局循环架构（Fully Looped Architecture），将循环间信号分布到所有层以缓解残差爆炸；（2）注意力注入（Attention Injection），重用现有的注意力模块以抑制梯度振荡。这些修改稳定了训练动态，使得全局循环变压器能够稳定地训练多达12次循环迭代，而其他基线循环模型在此情况下崩溃。在循环变压器不会崩溃的较温和设置中，全局循环变压器仍然将下游任务平均性能提升高达13.2%。总体而言，我们的实验表明，全局循环变压器提高了训练稳定性，增强了下游性能，并通过在推理时改变循环迭代次数，初步展现出在不同测试时计算预算下的适应性。 ## 1 引言 大型语言模型（LLMs）通过在互联网上大规模文本语料库上的预训练，在众多下游任务中展现出卓越的泛化能力。然而，这种缩放范式日益受到高质量公开文本数据有限供应的制约。根据Chinchilla缩放定律（Hoffmann等人，2022 (https://arxiv.org/html/2605.18797#bib.bib20)），模型参数和训练数据应成比例缩放。仅仅增加模型规模而不相应增加数据会导致训练结果次优。同时，Villalobos等人（2024 (https://arxiv.org/html/2605.18797#bib.bib1)）表明，可用于LLM预训练的公开人类生成文本数据存量每年仅增长10%，而实践中使用的数据集规模每年大约增长2.4倍。数据可用性与计算能力之间日益加剧的错配呼唤新的范式，能够在不依赖日益庞大的数据集的情况下，将多余的计算转化为性能提升。为了应对这一挑战，已有多个研究方向得到探索。一方面的工作侧重于通过改写整个语料库来提高数据质量（Niklaus等人，2026 (https://arxiv.org/html/2605.18797#bib.bib21)），但这种方法从根本上仍然依赖于高质量数据的可用性。另一个有前景的方向是测试时缩放（test-time scaling），利用强化学习在推理时引出扩展的思维链推理（Wei等人，2023 (https://arxiv.org/html/2605.18797#bib.bib51)）（团队等人，2024 (https://arxiv.org/html/2605.18797#bib.bib22)；Guo等人，2025 (https://arxiv.org/html/2605.18797#bib.bib23)）。尽管其经验性能强劲，但这种范式通常会大幅增加上下文长度和推理成本。将额外计算转化为性能的一种补充方式是为模型提供循环机制。循环变压器（Giannou等人，2023 (https://arxiv.org/html/2605.18797#bib.bib17)）遵循这一方向，通过迭代复用相同的Transformer骨干模块。它不是增加模型参数数量或扩展上下文窗口，而是将相同的共享模块展开为多次循环迭代。凭借循环机制，循环变压器参数高效（Saunshi等人，2025 (https://arxiv.org/html/2605.18797#bib.bib7)），天然兼容测试时计算调整（Koishekenov等人，2025 (https://arxiv.org/html/2605.18797#bib.bib31)），并在推理任务中表现出色（Saunshi等人，2025 (https://arxiv.org/html/2605.18797#bib.bib7)）。然而，尽管有这些优势，当循环迭代次数增多时，循环变压器仍然难以训练（Geiping等人，2025 (https://arxiv.org/html/2605.18797#bib.bib19)；Zhu等人，2025 (https://arxiv.org/html/2605.18797#bib.bib18)）。我们在实验中也观察到了同样的问题：简单地增加循环变压器的循环迭代次数会导致训练崩溃。这促使我们检查循环变压器的训练动态。我们的诊断识别出两种不稳定模式。首先，在训练的早期阶段，梯度可能强烈振荡。其次，残差状态范数可能随着循环迭代次数的增加而迅速增长。我们将这两种现象分别称为梯度振荡和残差爆炸。 参见图注 图1：循环变压器（LT）与全局循环变压器（FLT）的比较。FLT采用全局循环架构，所有层都参与循环，并采用注意力注入，通过将自注意力模块复用为交叉注意力模块，将上一迭代产生的残差流重新引入模型。我们在所有注意力模块中使用因果掩码。 基于这一诊断，我们提出了全局循环变压器（Fully Looped Transformer），这是对循环变压器的一种简单的无参数改进。它包含两项修改：（1）全局循环架构，将循环信号分布到所有层，缓解残差爆炸；（2）注意力注入，重用现有注意力模块以受控方式注入循环信息，抑制梯度振荡。这两项修改共同稳定了训练，且无需引入额外的可学习参数。我们的主要贡献有三点：（1）我们诊断了循环变压器的训练不稳定性，并识别出两种相关现象：梯度振荡和残差爆炸。（2）我们提出了全局循环变压器，这是一种简单而有效的改进，可以在不增加任何参数的情况下稳定训练。（3）通过循环缩放和消融实验，我们表明全局循环变压器在基线循环模型崩溃的区间内能够稳定训练，将下游任务平均性能提升高达13.2%，并初步展现出对不同推理时计算预算的适应性。我们的代码可在GitHub (https://github.com/FuRuF-11/FullyLoopedTransformer) 获取。 ## 2 相关工作 ### 2.1 权共享模型 许多循环或循环语言模型可以概念化为权共享架构，利用一种基础技术来增强深度神经网络的参数效率。通用变压器（Universal Transformer）（Dehghani等人，2019 (https://arxiv.org/html/2605.18797#bib.bib13)）首先将这一思想扩展到Transformer架构，通过跨层重复应用相同的骨干网络，使模型能够执行迭代细化，并赋予Transformer模拟通用计算过程的能力。循环变压器（Looped Transformer）（Giannou等人，2023 (https://arxiv.org/html/2605.18797#bib.bib17)）通过使用更通用的仅解码器Transformer架构简化了这一思想。大量现有工作（Saunshi等人，2025 (https://arxiv.org/html/2605.18797#bib.bib7)；Geiping等人，2025 (https://arxiv.org/html/2605.18797#bib.bib19)；Koishekenov等人，2025 (https://arxiv.org/html/2605.18797#bib.bib31)）已经证明了循环变压器在测试时缩放和推理任务中的显著优势。与此同时，Parcae（Prairie等人，2026 (https://arxiv.org/html/2605.18797#bib.bib37)）尝试通过动力系统框架来约束残差状态的谱范数，以解决循环模型的不稳定性。相比之下，我们的工作引入了一种通用的架构修改，无需额外的可学习参数或计算开销，允许通过最少的改动集成到现有的Transformer架构（Vaswani等人，2023 (https://arxiv.org/html/2605.18797#bib.bib44)）中。 ### 2.2 循环模型中的训练困难 深度循环模型中的训练困难由来已久，最关键的当属梯度爆炸和梯度消失问题。Bengio等人（1994 (https://arxiv.org/html/2605.18797#bib.bib33)）系统地指出，在使用时间反向传播（BPTT）训练循环网络时，梯度信号会随着每个时间步呈指数级衰减或爆炸，使得模型难以学习长距离依赖。Pascanu等人（2013 (https://arxiv.org/html/2605.18797#bib.bib32)）进一步从理论上分析了该问题的根源，并提出了梯度裁剪等缓解策略。LSTM（Hochreiter和Schmidhuber，1997 (https://arxiv.org/html/2605.18797#bib.bib30)）和GRU（Cho等人，2014 (https://arxiv.org/html/2605.18797#bib.bib34)）引入了门控机制来解决此问题，通过选择性过滤和控制信息流，有效缓解了问题。然而，在循环变压器的背景下，这些挑战依然存在。由于模型使用共享权重的循环结构，它本质上相当于一个极深的循环网络，使得梯度稳定性再次成为训练的核心挑战。 ## 3 诊断循环变压器的不稳定性 先前的研究表明，增加循环迭代次数会显著加大LT的训练难度（Geiping等人，2025 (https://arxiv.org/html/2605.18797#bib.bib19)；Zhu等人，2025 (https://arxiv.org/html/2605.18797#bib.bib18)），但它们并未解释这种不稳定的原因。为了探究这种失败模式，我们检查了LT在完全收敛之前的早期训练动态。在我们的诊断运行中，相关的不稳定性在优化的早期阶段就已经出现，因此我们使用前2000个优化器步骤作为诊断窗口111此窗口仅用于识别训练不稳定性的出现和持续；最终性能比较将在第5节（https://arxiv.org/html/2605.18797#S5）中使用完全训练的模型进行。。在图2（https://arxiv.org/html/2605.18797#S3.F2）中，我们比较了使用相同6层Transformer骨干的6次、9次和12次循环迭代的LT模型，并将12次循环的全局循环变压器（FLT）作为稳定对照组。所有模型使用相同的预训练设置。我们监控三个量：训练损失、残差状态范数以及第一个FFN模块的梯度L2范数。额外的实验结果和训练细节见附录A.1.1（https://arxiv.org/html/2605.18797#A1.SS1.SSS1）和A.2（https://arxiv.org/html/2605.18797#A1.SS2）。我们将训练崩溃定义为模型无法再产生可用检查点的失败模式。在实践中，这包括进入持续的高损失平稳期、出现严重的损失波动，或产生无法使用的验证性能。在此定义下，12次循环的LT在早期诊断运行中崩溃：在初始下降后，其训练损失停止改善并停留在高损失平稳期，而其残差状态范数持续增加。除了这种崩溃情况外，损失曲线还揭示了LT存在一种较轻微但持续的优化困难。随着循环迭代次数从6增加到9，LT并未立即崩溃，但出现了明显的损失差距：9次循环的LT在诊断窗口的大部分时间内保持比6次循环的LT更高的训练损失。这表明，即使在出现绝对的高损失平稳期之前，增加循环次数就已经使得优化变得更加困难。相比之下，12次循环的FLT在相同的2000步窗口内保持稳定，损失下降更平滑，残差状态范数更小，梯度动态更稳定。基于这些观察，我们提出两个假设： 参见图注 图2：LT和FLT在前2000个优化器步骤中的训练动态。左图：残差状态范数。中图：训练损失（为便于阅读，以0.9因子平滑）。右图：第一个FFN模块的梯度L2范数。 梯度振荡可能导致早期优化困难。循环变压器在循环迭代中反复应用相同的Transformer骨干，使其训练动态类似于通过BPTT算法优化的循环模型（Werbos，2002 (https://arxiv.org/html/2605.18797#bib.bib52)）。尽管循环迭代次数远小于传统长序列RNN的时间步数，但梯度仍然可以通过共享的循环模块累积。因此，我们假设这种循环维度的梯度累积是早期梯度振荡的一个可能来源，这使得优化更加敏感，并增加了进入不稳定训练区间的可能性。 残差爆炸可能成为高循环次数的瓶颈。对固定深度的骨干进行循环增加了计算图的有效循环深度，这意味着在训练期间展开循环模型使其在计算图上等价于一个深得多的模型。在这种观点下，即使每个循环迭代对残差状态有轻微的放大，也会在多次应用相同共享层后累积。因此，我们假设高循环次数的LT可能在循环维度上产生残差放大，导致残差爆炸。由于这种现象在初始瞬态之后仍然持续，我们将其视为一种独立的不稳定模式，它本身可能不能完全解释早期崩溃，但在循环次数很大时可能成为主要的优化瓶颈。 总体而言，早期诊断结果揭示了LT训练中的两种不稳定模式：早期优化阶段的梯度振荡和高循环次数下的残差爆炸。这些观察结果启发了下一节中介绍的设计选择，我们的目标是稳定循环训练，而不增加可学习参数。 ## 4 全局循环变压器 在第3节（https://arxiv.org/html/2605.18797#S3）中，我们识别了与循环变压器（LT）训练不稳定性相关的两种模式：梯度振荡和残差爆炸。我们设计全局循环变压器（FLT）来缓解这两个问题。FLT引入了两项改进：（1）全局循环架构（FLA）和（2）注意力注入（AI）。图1（https://arxiv.org/html/2605.18797#S1.F1）展示了FLT与LT的整体架构比较，以及FLT中FLA和AI的具体实现。 ### 4.1 全局循环架构 尽管我们所有的模型已经在每一层中包含了RMSNorm（Zhang和Sennrich，2019 (https://arxiv.org/html/2605.18797#bib.bib45)）来调节输出的残差状态，但残差爆炸仍然发生。这表明仅靠归一化是不够的，必须从更深层次的结构角度来解决这个问题。