Parcae：用更少参数实现更大性能的稳定循环语言模型

摘要

我们提出 Parcae，这是循环语言模型最早稳定的架构之一，其质量达到两倍规模的 Transformer，且训练干净、可预测。Parcae 通过增加循环次数而非纯粹扩展数据，创造了一种扩展质量的新途径，为训练内存受限的端侧模型开辟了高效前沿。

充分利用你的参数

传统的缩放定律告诉我们，要达到最佳性能，需要扩展 FLOPs，通常通过增加参数或数据。但随着模型向边缘迁移，推理成本飙升，我们不禁思考：能否在不增加内存占用的情况下扩展质量？

为此，我们一直在探索循环架构，即通过多次将激活值传递到相同层来增加计算量的模型。尽管前景广阔，但这些模型训练不稳定。我们直接解决这一问题，并引入 Parcae，一种稳定的循环架构，它：

优于先前的循环模型：Parcae 相比之前的大规模循环方案，验证困惑度降低高达 6.3%。以小博大：我们的 770M Parcae 匹配 1.3B 参数 Transformer 的质量，在相同数据上训练，以大约一半的参数达到相同性能。可预测缩放：我们建立了首个循环缩放定律，发现计算最优训练需要同时增加循环次数和数据。

循环模型很酷，但实际训练困难

随着模型向边缘迁移，推理部署占据更大计算份额，人们越来越关注在不增加参数的情况下扩展模型质量。我们兴奋的一种机制是层循环，早期工作已训练出匹配更大固定深度架构质量的循环模型。

要将普通 Transformer 转变为循环模型，我们遵循先前工作，将其层划分为三个功能块：前奏（$\mathcal{P}$）、循环（$\mathcal{R}$） 和 尾声（$\mathcal{C}$）。前向传播分三个阶段进行：

嵌入：前奏将输入转换为潜在状态 $e$。循环：循环块迭代更新隐藏状态 $h_t$，共 $T$ 次循环。为保持输入的影响，$e$ 被注入到每次循环中，通常通过加法 [1]（$h_{t+1} = \mathcal{R}(h_t + e)$）或带投影的拼接 [2]（$h_{t+1} = \mathcal{R}(W[h_t; e])$）。输出：尾声处理最终的 $h_T$ 以生成模型输出。

不幸的是，循环模型训练起来很麻烦[2][3][4]。我们个人发现它们存在残差状态爆炸和损失尖峰的问题。更棘手的是，循环模型中的循环块由多个普通Transformer块组成，这使得难以推断不稳定的根源。

理解循环的不稳定性

虽然不稳定性是一个难以捉摸的敌人，但我们观察到，一个简单的线性框架捕捉到了不稳定性的一个重要来源。具体来说，我们将循环重新表述为一个关于残差的非线性时变动力系统，其更新规则为：

其中 $\overline{A}, \overline{B}$ 执行注入，$\overline{\mathcal{R}}$ 是Transformer块对残差流的贡献。对于次二次序列混合的狂热者，请注意，如果我们忽略非线性项 $\overline{\mathcal{R}}$，得到的系统是一个关于残差状态（跨模型深度）的离散线性时不变（LTI）动力系统。

很酷的是，对于离散LTI系统，其稳定性和收敛性由 $\overline{A}$ 的特征值决定。具体来说，稳定性使用谱范数 $\rho(\overline{A})$（即 $\overline{A}$ 的绝对值最大的特征值）来分类，稳定系统（收敛）满足 $\rho(\overline{A})<1$，不稳定系统（发散）满足 $\rho(\overline{A})=1$。

虽然这种分析绕过了循环的非线性（例如，注意力机制和MLP单元），但上表和上图证实，我们的分析在经验上很重要：发散的运行学习到的谱半径 $\rho(\overline{A}) \geq 1$，而收敛的运行保持 $\rho(\overline{A}) < 1$。当我们通过Parcae维持LTI条件时，循环模型对超参数选择变得更加鲁棒。

Parcae：一种稳定、无麻烦的循环模型

那么，我们如何实现稳定呢？我们设计了一种新的循环模型Parcae，它通过构造显式地维持了上一节中观察到的稳定性条件。具体来说，我们使用连续形式 $A, B$ 参数化输入注入参数，并通过ZOH和欧拉方案（即 $\overline{A} = \exp(\Delta A)$ 和 $\overline{B} = \Delta B$）进行离散化，使用学习到的 $\Delta \in \mathbb{R}^{d_h}$。然后，我们将 $A := \texttt{Diag}(-\exp(\texttt{log}_A))$ 约束为负对角矩阵，其中向量的 $\texttt{Diag}(-\exp(\cdot))$ 强制为负，$\texttt{log}_A\in \mathbb{R}^{d_h}$ 是我们的可学习向量。这确保了 $\rho(\overline{A}) < 1$！

那么，我们是否解决了所有问题并稳定了循环模型？不幸的是，要获得Parcae的干净训练，还需要一些其他小技巧。有兴趣的读者可以查看我们的论文。

回到语言建模：扩展Parcae

Parcae 不仅训练更可靠，而且我们发现，与之前的 RDM 相比，它生成的模型质量更高。使用与 RDM [2]（一种先前的循环模型）完全相同的设置，我们与参数和数据匹配的 RDM 进行了对比测试，观察到 Parcae 将验证困惑度降低了高达 6.3%。

当将一个非常强大的 Transformer 基线改造为 RDM 时，无需任何超参数调整，我们发现 Parcae 比 RDM（直接发散）更稳健。

我们还采用了 Parcae，并将其作为标准固定深度 Transformer 的直接替代品。使用受 nanochat 启发的设置，我们在 FineWeb-Edu 上训练了一系列语言模型，参数规模高达 1.3B。我们发现 Parcae 优于所有参数和数据匹配的 Transformer，我们的 770M Parcae 模型几乎达到了两倍大小 Transformer 的下游质量！

循环还是不循环

但循环是否真的 FLOP 高效？为了研究这一点，我们探索了一种设置，在固定参数数量和 FLOP 预算下，我们在训练中权衡平均循环次数与数据量（例如，如果我们增加平均循环次数，则减少训练数据以维持固定的 FLOP 预算）。

在两个规模上，我们发现增加训练中使用的平均循环次数 $\mu_{\text{rec}}$ 同时按比例减少 token 数量，比使用低循环次数和更多数据训练能获得更低的验证损失。更酷的是，如果我们使用抛物线拟合来提取每个 FLOP 级别下的最优 $\mu_{\text{rec}}$ 和 token 预算，我们发现它们都遵循具有一致指数的幂律。

好了，好了。但我们能击败固定深度模型吗？ 使用我们的最优循环缩放定律，我们比较了固定深度 Parcae 模型（即 $\mu_{\text{rec}}=1$ 的模型）和遵循我们缩放定律的最优平均循环预测的循环 Parcae 模型。我们发现循环为验证损失创建了更严格的帕累托前沿（下图），这转化为更好的下游质量（下表）。

下一步与亲自尝试 Parcae

我们对参数效率能推到多高感到非常兴奋。随着推理过程中内存开销成本的不断增长，我们认为在参数重用方法（如层循环）方面有很多值得探索的地方。为了加速这一过程，我们发布了训练代码和模型。我们还没有结束；我们有很多新想法来进一步推动循环模型，敬请期待接下来的进展！

如果您有任何问题或想与我们合作推进 Parcae 的下一步发展，请联系 Hayden Prairie：[email protected]。

名称 PaRCae 致敬了三位罗马命运女神：诺娜（前奏块 $\mathcal{P}$），她初始化了计算的生命之线；德西玛（循环块 $\mathcal{R}$），她丈量丝线并随模型深度演化；以及莫塔（尾声块 $\mathcal{C}$），她通过剪断丝线来终结序列，产生最终输出。

致谢

我们感谢 Together AI 与我们合作并为这些实验提供计算资源。我们也感谢 Austin Silveria 和 Jonah Yi 对这篇博文提出的宝贵反馈。