论文由加州大学圣迭戈分校和Together AI的研究者共同完成聚焦于如何稳定地训练循环Looped架构的大语言模型并将其发展为一种独立于模型参数和数据量的、可预测的扩展计算规模的新维度。核心贡献是提出了新型循环架构Parcae通过动态系统理论和工程化改进解决了该类模型长期存在的训练不稳定问题并首次系统性地推导了其在训练和推理阶段的扩展定律。1. 研究背景与问题定义背景传统Transformer模型提升性能主要靠增加参数或训练数据但这会大幅增加内存和部署成本。循环架构如Universal Transformers通过让激活在同一个层块中反复循环来增加计算量FLOPs旨在用固定参数实现更强的性能尤其适用于边缘计算和推理场景。核心问题现有循环模型如RDM训练极不稳定主要表现为残差状态爆炸隐藏状态范数在循环中急剧增长导致数值溢出。损失尖峰训练过程中损失突然飙升导致训练失败。超参数敏感对学习率等超参数极其敏感难以大规模稳定训练。2. 核心方法论Parcae架构的设计与稳定化作者将循环过程视为一个非线性时变动态系统并通过对该系统线性化从控制理论角度揭示了不稳定的根源并据此设计了Parcae。2.1 不稳定性根源分析理论创新将循环更新表示为h_{t1} \overline{A} h_t \overline{B} e 非线性项关键发现系统的稳定性取决于状态转移矩阵 \overline{A} 的谱范数 ρ(\overline{A})。现有方法如残差连接中ρ(\overline{A}) 要么等于1临界稳定要么不受约束这导致了系统的内在不稳定性。2.2 Parcae的稳定化策略技术创新约束 \overline{A} 的谱范数将 \overline{A} 参数化为负对角矩阵并进行离散化从数学上强制其谱范数 1确保系统是严格稳定的。引入“前奏归一化”Prelude Norm对循环单元的输入嵌入进行归一化有效防止了大型模型如1.3B参数在训练后期出现的状态爆炸和损失尖峰。改进训练算法每序列深度采样在同一个批次中为每个序列独立采样循环深度替代原有的每微批次采样更准确地逼近训练目标显著减少了损失尖峰。修正梯度截断策略修正了先前工作中导致的分布偏移问题使模型能更好地泛化到训练时未见过的循环深度。3. 实验与主要发现论文在语言建模任务上进行了大量实验验证了Parcae的稳定性、有效性和扩展潜力。3.1 稳定性与端到端质量提升稳定性验证Parcae在不同学习率下均能稳定收敛而基线模型则频繁发散。消融实验证实其三个核心改进约束A、每序列采样、前奏归一化均对稳定性有显著贡献。性能提升优于RDM在相同参数和计算量下Parcae的验证困惑度降低了6.3%下游任务平均分提升了1.8分。优于标准Transformer在固定参数和数据的条件下7.7亿参数的Parcae模型在Core基准上性能可媲美13亿参数的Transformer。其参数效率最高可达后者的87.5%。3.2 训练时扩展定律计算最优新扩展维度首次证明了“循环次数μ_rec”是继“参数量”和“数据量”之后的第三个正交的计算扩展轴。最优分配定律在固定的FLOP预算下增加循环次数并相应减少训练数据能取得比固定深度模型更低的验证损失。最优的循环次数和训练数据量遵循可预测的幂律关系。3.3 测试时扩展定律推理加速饱和性测试时增加循环次数可以降低损失但收益是递减的并最终趋近于一个由训练决定的“不可约损失”。可预测的指数衰减测试时的性能提升遵循饱和指数衰减曲线这使得我们可以准确预测在给定推理预算下的最佳循环次数。统一扩展定律论文成功将训练时和测试时的扩展定律统一起来形成了一个能够预测模型在不同训练FLOP和推理FLOP下性能的完整公式。4. 结论与意义结论Parcae成功解决了循环架构的训练不稳定性难题使其成为一种可行、高效且可预测的模型扩展方案。意义理论贡献为理解和设计稳定的循环神经网络提供了基于控制理论的新视角。实践价值为在资源受限环境下如移动端、边缘设备部署高性能大模型以及在不增加参数的情况下提升模型推理能力提供了新的思路和技术基础。这里是自己的论文阅读记录感兴趣的话可以参考一下如果需要阅读原文的话可以看这里如下所示摘要传统的固定深度架构通过增加训练浮点运算量FLOPs来扩展质量通常是通过增加参数化程度但这会以更高的内存占用或更多数据为代价。一个潜在的替代方案是循环架构它通过让激活在层块中循环来增加 FLOPs。尽管前景广阔但现有的循环架构训练方法可能不稳定会出现残差爆炸和损失尖峰。我们通过将循环重新定义为残差流上的非线性时变动态系统来应对这些挑战。通过对该系统进行线性近似我们发现现有循环架构的不稳定性源于其注入参数中较大的谱范数。为了解决这些不稳定性问题我们提出了 Parcae一种新颖的稳定循环架构它通过负对角参数化的离散化来约束注入参数的谱范数。因此与先前的大规模循环模型相比Parcae 的验证困惑度降低了高达 6.3%。利用我们稳定的循环架构我们研究了循环作为一种通过增加训练和测试时的 FLOPs 来提高质量的手段的扩展特性。对于训练我们推导出可预测的幂律以便在保持参数数量固定的同时扩展 FLOPs。我们初步的扩展定律表明在给定的 FLOP 预算下循环和数据应协同增加。在测试时我们发现 Parcae 可以利用循环来扩展计算遵循可预测的饱和指数衰减。当扩展到 13 亿参数时我们发现在固定的参数和数据预算下与强大的 Transformer 基线相比Parcae 在 CORE 和 Core-Extended 基准上的质量分别提高了 2.99 分和 1.18 分达到了两倍规模 Transformer 相对质量的 87.5%。1 引言扩展定律已确定模型性能随着 FLOPs 的增加而可预测地提升 [33, 41]通常通过增加参数数量或训练数据来实现。这些扩展定律表明FLOP 最优的训练应遵循经验幂律同步增加参数和训练数据。因此为了与数据规模同步扩展最先进模型的深度和宽度都在增长这 subsequently 增加了部署这些模型的内存占用 [20, 46]。然而随着推理部署在计算中所占比例越来越大 [77]并且部署开始向边缘移动 [52, 54]在不增加参数的情况下扩展模型质量的兴趣日益增长。实现此目标的一种机制是层循环模型例如循环 Transformer [18, 27, 94]它们将激活迭代地循环通过一个层块。初步结果令人鼓舞循环模型在质量上可与更大的固定深度架构相媲美 [27, 94]。此外它们显示出潜在推理 [69, 87] 和每令牌自适应计算 [27, 49] 的潜力。不幸的是先前的研究 [27, 36, 49] 和我们的工作观察到这些模型的训练不稳定会出现残差状态爆炸和损失尖峰。由于这些模型循环的是复杂非线性架构如 Transformer 块 [78]的层因此很难从分析上理解循环模型中不稳定的来源。因此训练需要敏感的超参数选择和残差归一化例如后归一化来纠正这种不稳定性 [27]。此外即使在收敛的训练运行中当循环模型在随机深度上训练以增强测试时扩展时我们也会观察到损失尖峰 [4]。在本文中我们研究这种不稳定性并探究稳定这些模型是否能解锁循环作为一个可预测的、正交的计算扩展轴。为了分析不稳定性我们观察到先前的循环架构可以被重新表述为残差流上的非线性时变动态系统 [56]形式为我们在端到端质量、训练 FLOP 扩展和测试时扩展方面评估 Parcae端到端质量。我们将 Parcae 与参数和数据匹配的 RDM [27] 和 Transformer 进行比较。与 RDM 相比Parcae 将验证困惑度val. PPL降低了 6.3%。当扩展到 13 亿参数和 1000 亿令牌时Parcae 在 Core 和 Core-Extended [45] 基准上分别比参数匹配的 Transformer 高出 2.99 分和 1.18 分——在质量上匹配了高达两倍规模的 Transformer。训练 FLOP 扩展。为了评估 FLOP 训练扩展我们在参数匹配的 isoFLOP 设置中研究循环的扩展定律即是否用增加的数据或循环来扩展 FLOPs。我们发现循环引入了一个正交的扩展轴类似于参数和数据。具体来说FLOP 最优的训练会按照经验幂律同步增加循环和数据见图 1 [右]。测试时扩展。我们研究循环作为扩展测试时计算的一种机制观察到循环遵循可预测的指数衰减并存在一个不可约损失。我们进一步将测试时和训练时的幂律结合起来为 PARCAE 模型中的循环创建了一个统一的扩展定律。2 背景我们首先简要介绍循环模型第 2.1 节、LTI 系统第 2.2 节和建模扩展定律第 2.3 节的背景。先前的工作已经沿着几个设计轴研究了循环架构循环放置前、中或后循环[68]、停止机制显式路由器 [6, 94] vs. 隐式随机深度 [27, 49]、拓扑结构单块 [27] 或分层 [38, 79]以及微分显式或隐式反向传播 [7]。我们的工作侧重于使用显式微分的隐式停止中循环架构扩展综述见第 B 节。2.1 现有中循环架构2.2 线性时不变动态系统为了研究循环模型的不稳定性我们将使用 LTI 动态系统作为复杂非线性循环模型的可处理的线性代理。在控制理论中LTI 系统通过图 2循环架构的训练不稳定性。左前归一化循环模型发散而残差归一化和 PARcAE 收敛。右不稳定性源于爆炸的循环状态范数 ∣∣hT∣∣2即 TT 次循环后的隐藏嵌入范数。2.3 建模扩展定律我们遵循 Hoffmann 等人 [33] 的方法该方法通过抛物线和参数拟合针对不同的模型大小和训练令牌数量在固定的 FLOP 预算下对扩展定律行为进行建模。对于抛物线拟合对多个 FLOP 预算进行二次拟合以估计损失最优的模型大小或训练令牌数量。对于参数拟合使用 Huber 损失 [34] 在预测值和经验对数损失值之间针对变化的参数 N 和令牌 D拟合一个函数形式 L^(N,D)EX⋅N−xY⋅D−y并使用 L-BFGS [55] 进行最小化。3 理解循环架构中的不稳定性在本节中我们研究循环架构的不稳定性。使用残差上的 LTI 视角我们发现不稳定性源于不受约束的残差状态爆炸图 2表 2 [基线]第 F 节。虽然残差归一化有助于缓解这个问题但它需要敏感的超参数调整表 2 [Res. Norm]类似于固定深度 Transformer [83, 84]。使用这个 LTI 框架我们推导出 A 特征值的稳定性条件。我们发现先前的工作不满足这些 A 的条件我们凭经验验证这会造成严重的状态爆炸图 3。表 1基于 LTI 表示的先前更新规则稳定性比较。表 2超参数不稳定性。基线 RDM、残差归一化 RDM 和 PARCAE 在不同学习率下的收敛情况。PARCAE 对超参数选择更鲁棒。完整日志见第 F 节。4 Parcae一种稳定的循环架构使用我们的动态系统框架我们创建了 Parcae一种明确满足稳定性约束的循环架构第 4.1 节。此外我们提出了一种每序列深度采样方法以稳定由可变深度引入的方差第 4.2 节。4.1 Parcae 的块设计和稳定参数化4.2 Parcae 的稳定训练算法5 结果我们在端到端质量第 5.1 节、训练 FLOP 扩展第 5.2 节和测试时扩展第 5.3 节方面评估 Parcae。我们发现 Parcae 在性能上优于参数和数据匹配的 RDM 和 Transformer最优的循环和数据遵循可预测的幂律并且测试时循环遵循饱和指数衰减。表 3在 RDM 设置下训练的零样本和困惑度结果。比较 PARCAE 和 RDM [27] 在各种开源基准以及困惑度、留出验证集和 Wikitext [50] 上的表现。最佳结果已加粗。表 4在 Transformer 设置下训练稳定性结果。为了说明稳定性我们将一个 1.4 亿参数的基线 Transformer 改造为 RDM然后顺序添加我们的稳定性改进。5.1 Parcae 提升端到端质量我们将 PARCAE 与参数和数据匹配的 RDM 和 Transformer 进行比较发现 PARCAE 比先前的循环模型更稳定并且在质量上均优于两者。设置。对于 RDM我们遵循 Geiping 等人 [27] 的方法使用 Huginn 数据集和分词器进行训练。对于 Transformer我们遵循 Karpathy [42] 的方法在 FineWeb-Edu [60] 上进行训练。对于 RDM 和 Transformer 设置我们都对 RDM 和 Transformer 进行了超参数扫描然后将其用于 PARCAE即我们没有为 PARCAE 模型进行超参数扫描。扩展的模型定义、超参数选择和评估设置可以在第 P 节、第 Q 节和第 M 节分别找到。表 5Parcae 与固定深度 Transformer 的比较。我们在 nanochat 设置下以几种规模预训练 Transformer 和 PARCAE并在留出验证集、Lambada [58]、Core 和 Core-Extended [45] 上进行评估。最佳结果已加粗。5.2 循环作为训练中的正交扩展轴在本节中我们探讨在固定 FLOP 和参数预算下循环的 FLOP 效率。我们发现循环引入了一个正交的扩展计算轴其中计算最优的训练会按照经验幂律同步增加 μrec​ 和数据。设置。我们使用 nanochat 设置训练 140M 和 370M 的 PARCAE 模型在固定的 FLOP 和参数预算下改变训练令牌数量和平均循环次数 μrec​。额外的训练细节和 FLOP 估算可以在第 O 节和第 D 节找到。循环的扩展定律建模。在 140M 和 370M 规模下isoFLOP 曲线显示增加 μrec​ 同时按比例减少令牌产生的验证损失低于低循环图 6循环的帕累托前沿。我们观察到循环在 IsoFLOP 最优损失前沿上比固定深度、非循环模型更严格。点代表经验数据点。循环与固定深度的 IsoFLOP 比较图 6 显示了每个 FLOP 预算下没有循环的固定深度 PARCAE 模型。最优曲线实现了严格更低的损失这转化为 1.2 到 2.0 分更高的 Core 得分表 6。图 7Parcae 的测试时扩展。在评估表 5 中的 PARCAE 模型时我们观察到测试时循环遵循可预测的饱和趋势在各种模型大小上保持一致。5.3 Parcae 的测试时扩展定律我们研究循环作为扩展测试时计算的一种机制。我们发现测试时计算遵循可预测的饱和指数衰减可以与第 5.2 节结合连接训练和测试时的扩展定律。6 讨论与未来工作在本节中我们简要讨论局限性和未来的方向。循环架构。虽然围绕循环架构的若干设计选择已由小规模经验结果所指导但仍需在更大规模上对循环单元的放置[35]、组成例如循环单元中的参数数量以及不同架构的使用以及极端循环例如将平均循环次数增加到更深层次进行深入探究。在我们的动态系统框架内使用不同的离散化方法、满秩参数化和循环更新规则也值得研究以便在更深的循环深度上实现扩展。扩展。虽然我们发现Parcae为层循环引出了可预测的、最优的扩展定律但我们的观察仅限于小规模架构。当将这些观察结果扩展到更大的FLOP预算和参数化时Parcae是否仍具有优势还有待验证。我们也对参数、数据和循环作为正交轴的相互作用以及它们应如何被高效地协同扩展感兴趣。最后循环的一个局限性在于随着µrec的增加达到同等质量所需的测试时步数也会增加。研究如何在减少推理步数的情况下保持质量是一个有趣的未来方向。7 结论在这项工作中我们通过动态系统框架研究了循环模型的稳定性并提出了Parcae一种稳定的循环架构它通过约束注入参数的谱范数来防止残差爆炸。Parcae在性能上优于数据匹配和参数匹配的先前循环模型以及基线Transformer其下游任务质量可与高达两倍规模的模型相匹配。我们进一步建立了循环的扩展定律FLOP最优的训练会按照可预测的幂律同步增加循环和数据而测试时循环则遵循饱和指数衰减定律从而产生了一个连接训练和推理计算的统一扩展定律。