Parcae: Scaling Laws For Stable Looped Language Models——用于稳定循环语言模型的扩展定律

论文由加州大学圣迭戈分校和Together AI的研究者共同完成，聚焦于如何稳定地训练循环（Looped）架构的大语言模型 ，并将其发展为一种独立于模型参数和数据量的、可预测的扩展计算规模的新维度 。核心贡献是提出了新型循环架构 Parcae，通过动态系统理论和工程化改进，解决了该类模型长期存在的训练不稳定问题，并首次系统性地推导了其在训练和推理阶段的扩展定律。

1. 研究背景与问题定义

背景：传统Transformer模型提升性能主要靠增加参数或训练数据，但这会大幅增加内存和部署成本。循环架构（如Universal Transformers）通过让激活在同一个层块中反复循环来增加计算量（FLOPs），旨在用固定参数实现更强的性能，尤其适用于边缘计算和推理场景。
核心问题 ：现有循环模型（如RDM）训练极不稳定，主要表现为：
1. 残差状态爆炸：隐藏状态范数在循环中急剧增长，导致数值溢出。
2. 损失尖峰：训练过程中损失突然飙升，导致训练失败。
3. 超参数敏感：对学习率等超参数极其敏感，难以大规模稳定训练。

2. 核心方法论：Parcae架构的设计与稳定化

作者将循环过程视为一个非线性时变动态系统，并通过对该系统线性化，从控制理论角度揭示了不稳定的根源，并据此设计了Parcae。

2.1 不稳定性根源分析（理论创新）

将循环更新表示为：h_{t+1} = \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)=E+X⋅N−x+Y⋅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 节分别找到。

表 5：Parcae 与固定深度 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）。

图 7：Parcae 的测试时扩展。在评估表 5 中的 PARCAE 模型时，我们观察到测试时循环遵循可预测的饱和趋势，在各种模型大小上保持一致。

5.3 Parcae 的测试时扩展定律

我们研究循环作为扩展测试时计算的一种机制。我们发现测试时计算遵循可预测的饱和指数衰减，可以与第 5.2 节结合，连接训练和测试时的扩展定律。

6 讨论与未来工作

在本节中，我们简要讨论局限性和未来的方向。

循环架构。虽然围绕循环架构的若干设计选择已由小规模经验结果所指导，但仍需在更大规模上对循环单元的放置 $35$ 、组成（例如，循环单元中的参数数量以及不同架构的使用）以及极端循环（例如，将平均循环次数增加到更深层次）进行深入探究。在我们的动态系统框架内，使用不同的离散化方法、满秩参数化和循环更新规则也值得研究，以便在更深的循环深度上实现扩展。

扩展。虽然我们发现Parcae为层循环引出了可预测的、最优的扩展定律，但我们的观察仅限于小规模架构。当将这些观察结果扩展到更大的FLOP预算和参数化时，Parcae是否仍具有优势，还有待验证。我们也对参数、数据和循环作为正交轴的相互作用，以及它们应如何被高效地协同扩展感兴趣。最后，循环的一个局限性在于，随着µrec的增加，达到同等质量所需的测试时步数也会增加。研究如何在减少推理步数的情况下保持质量，是一个有趣的未来方向。

7 结论

在这项工作中，我们通过动态系统框架研究了循环模型的稳定性，并提出了Parcae，一种稳定的循环架构，它通过约束注入参数的谱范数来防止残差爆炸。Parcae在性能上优于数据匹配和参数匹配的先前循环模型以及基线Transformer，其下游任务质量可与高达两倍规模的模型相匹配。我们进一步建立了循环的扩展定律：FLOP最优的训练会按照可预测的幂律同步增加循环和数据，而测试时循环则遵循饱和指数衰减定律，从而产生了一个连接训练和推理计算的统一扩展定律。