Scaling Collapse — 计算最优训练的普适动力学

论文来源：arXiv:2507.02119 (ICML 2025)
作者：Shikai Qiu, Lechao Xiao, Andrew Gordon Wilson, Jeffrey Pennington, Atish Agarwala

摘要

这篇论文揭示了一个令人惊讶的现象：尽管神经网络训练涉及架构、训练算法和数据的复杂交互，但计算最优训练的模型展现出极其精确的普适性。具体表现为：当训练计算量和损失被归一化到训练结束时的单位值时，不同尺寸模型的损失曲线会精确地塌缩到同一条普适曲线上。配合学习率衰减，这种塌缩效果会进一步增强，使得跨模型尺度的归一化曲线差异低于由随机种子引起的个体损失曲线噪声floor，论文称之为 Supercollapse（超塌缩）。

---

1. Scaling Collapse 的定义与数学描述

1.1 核心观察

Scaling Collapse 指的是：当模型大小 $p$ 不同但均以计算最优（compute-optimal）方式训练时，其损失曲线经过简单的仿射变换后会完全重合。这一现象类似于统计物理中的标度塌缩（scaling collapse）——在适当的重标度后，不同尺寸系统的可观测量会塌缩到同一条普适曲线上。

1.2 归一化损失曲线的定义

设 $L(t,p,\omega )$ 为模型大小为 $p$、随机种子为 $\omega$ 时，训练 $t$ 个token（与步数成正比）后的测试损失。定义归一化损失曲线 $\ell$ 为：

$$\ell (x,p,\omega )=\frac{L(x{t}^{\star }(p),p,\omega )-\hat{L}}{L(t^{\star }(p),p,\omega )-\hat{L}},x\in [0,1]$$

其中：

• $t^{\star }(p)=(p/p_0{)}^{\gamma }$ 是大小为 $p$ 的模型的计算最优训练时长
• $\hat{L}$ 是偏移量，论文发现 $\hat{L}=L_0$（不可约损失）时塌缩效果最佳
• $x=t/t^{\star }(p)$ 称为归一化计算量（normalized compute）

当不同 $p$ 的 $\ell (x,p,\omega )$ 在取平均后完全相同时，称这些曲线发生了 collapse。

1.3 塌缩质量的量化

论文定义了塌缩偏差（Collapse Deviation） $\Delta$：

$$\Delta (x)=\frac{\sqrt{V_{p,\omega }[\ell (x,p,\omega )]}}{E_{p,\omega }[\ell (x,p,\omega )]}$$

其中 $E_{p,\omega }$ 和 $V_{p,\omega }$ 分别表示对随机种子 $\omega$ 和模型大小 $p$（在对数均匀分布下）的期望和方差。

作为对比，定义单个模型的噪声floor $\sigma$：

$$\sigma (x,p)=\frac{\sqrt{V_{\omega }[L(x{t}^{\star }(p),p,\omega )]}}{E_{\omega }[L(x{t}^{\star }(p),p,\omega )]}$$

这衡量了每个模型在不同随机种子间的相对波动。

---

2. Supercollapse 的具体条件

2.1 定义

Supercollapse 是比普通 collapse 更强的形式，指的是在训练的大部分阶段，塌缩偏差低于噪声floor：

$$\Delta (x)<\sigma (x,p),\text{当 }x>1-\delta \text{ 时}$$

这意味着不同模型间的归一化损失差异，比我们预测任何单一模型损失的能力还要精确。

2.2 产生条件

Supercollapse 并非普遍存在，论文发现其产生需要满足以下条件：

2.2.1 学习率衰减（必要条件）

使用恒定学习率时，$\Delta (x)$ 很快上升到与 $\sigma (x,p)$ 相当的水平，两者大致相等。这意味着普通塌缩的变异主要来源于种子间的波动，而非模型间的差异。

然而，当引入学习率衰减时，情况发生了戏剧性变化：

• 衰减阶段 $\Delta (x)$ 显著降低
• 对于某些调度（如 cosine 衰减），$\Delta (x)$ 可以低于 $\sigma (x,p)$ 长达训练时间的 50%（即 $\delta$ 可达 0.5）

论文测试了多种学习率调度，均观察到 supercollapse：

• 线性衰减：$\eta (x)=1-x$
• 余弦衰减：$\eta (x)=\frac{1}{2}(1+\cos (\pi x))$
• 快速余弦衰减：$\eta (x)=\frac{1}{2}(1+\cos (3\pi x))$
• 带 warmup 的组合调度

2.2.2 计算最优数据指数（充分条件）

Supercollapse 还要求模型以计算最优的方式训练，即：

1. 正确的参数化：使用 $\mu$P 初始化和相应学习率缩放
2. 正确的数据指数：数据量与模型大小的关系 $t^{\star }(p)\propto p^{\gamma }$ 中，$\gamma$ 需接近最优值

2.2.3 归一化方式

关键在于使用随机种子的经验最终损失进行归一化，而非期望最终损失。这种做法利用了沿优化轨迹的时间相关性，将噪声转化为”控制变量”，显著降低了归一化曲线的方差。

2.3 Supercollapse 的数学解释

在训练末期 $\tau ={\tau }^{\star }-\delta \tau$（其中 $\tau$ 是梯度流时间），论文推导了塌缩偏差的近似表达式：

$${\tilde{\Delta }}^2(\tau )\approx {\bar{L}}^{-2}(\tau )\cdot g(\tau {)}^{\top }\eta (\tau )\bar{\Sigma }(\tau )g(\tau )\cdot \delta \tau +O(\delta {\tau }^2)$$

这表明 ${\tilde{\Delta }}^2(\tau )\propto \eta (\tau )\cdot (1-\hat{\tau })$，其中 $\hat{\tau }=\tau /{\tau }^{\star }$ 是归一化梯度流时间。

这一公式精确解释了为什么学习率衰减能导致 supercollapse：

• 学习率 $\eta$ 的降低直接减少了每步的优化噪声
• $(1-\hat{\tau })$ 项捕捉了训练终点附近的噪声累积
• 衰减越快的调度，在固定的归一化计算量下，$\eta$ 和 $(1-\hat{\tau })$ 都越小

---

3. SGD 噪声动态模型

3.1 梯度流时间下的 SDE

论文将 SGD 建模为梯度流时间 $\tau (t)={\int }_0^t\eta (s)ds$ 下的随机微分方程：

$$\frac{dw}{d\tau }=-\nabla L(w)+{\Sigma }^{1/2}(w)\xi (\tau )$$

其中：

• $w(\tau )$ 是参数
• $\xi (\tau )$ 是高斯白噪声，满足 $E[\xi (\tau )\xi ({\tau }^{\prime }{)}^{\top }]=\delta (\tau -{\tau }^{\prime })I$
• $\Sigma (w)$ 是 minibatch 梯度噪声的协方差矩阵

3.2 二次损失下的精确求解

对于二次损失 $L(w)=\frac{1}{2}{w}^{\top }Hw$（$H$ 为正定 Hessian），参数轨迹为：

$$w(\tau )=e^{-H\tau }w(0)-{\int }_0^{\tau }ds{e}^{-H(\tau -s)}{\Sigma }^{1/2}(w(s))\xi (s)$$

期望损失分解为：

$$\bar{L}(\tau )=\underset{\text{确定性项 }F(\tau )}{\underbrace{\frac{1}{2}E[\Vert e^{-H\tau }w(0){\Vert }_H^2]}}+\underset{\text{噪声项 }E(\tau )}{\underbrace{\frac{1}{2}{\int }_0^{\tau }ds\eta (s)\text{Tr}(H{e}^{-2H(\tau -s)}\bar{\Sigma }(s))}}$$

• 确定性项 $F(\tau )$：在无噪声极限 $\eta \Sigma \to 0$ 下的损失曲线，与学习率调度无关
• 噪声项 $E(\tau )$：由 SGD 噪声引起的额外损失，是学习率调度效应的载体

3.3 慢变近似下的简化

当 $\eta \bar{\Sigma }$ 变化比指数移动平均的时间尺度慢时：

$$\bar{L}(\tau )\approx F(\tau )+\frac{1}{4}\eta (\tau )\text{Tr}(\bar{\Sigma }(\tau ))$$

3.4 跨调度预测

论文利用上述模型预测不同学习率调度下的损失曲线。设参考轨迹为 $(\bar{L},\eta ,\bar{\Sigma })$，目标轨迹为 $({\bar{L}}^{\prime },{\eta }^{\prime },{\bar{\Sigma }}^{\prime })$，则：

$$\delta \bar{L}(\tau )\approx \frac{1}{4}\text{Tr}(\delta [\eta (\tau )\bar{\Sigma }(\tau )])$$

其中 $\delta [\eta \bar{\Sigma }]={\eta }^{\prime }(\tau ){\bar{\Sigma }}^{\prime }(\tau )-\eta (\tau )\bar{\Sigma }(\tau )$。

分解 $\delta (\eta \Sigma )=\delta \eta \cdot {\Sigma }^{\prime }+\eta \cdot \delta \Sigma$，发现第一项通常比第二项大 3-10 倍。忽略第二项，得到简化的预测公式：

$$L^{\prime }(\tau )\approx L(\tau )+\alpha \cdot \delta \eta (\tau )\cdot \text{Tr}({\Sigma }^{\prime }(\tau ))$$

其中 $\alpha$ 是跨调度、模型大小和训练时长共享的唯一超参数，论文发现 $\alpha \approx 0.21$（CIFAR-5M）或 $\alpha \approx 0.26$（MLP）即可获得惊人的精确预测。

3.5 自适应优化器的处理

对于 Adam 等自适应优化器，论文引入预处理协方差 $\tilde{\Sigma }=P^{-1/2}\Sigma P^{-1/2}$，其中 $P$ 是 Adam 的预处理矩阵。这样 $\mu$P 协议中的逐层学习率被吸收到预处理矩阵中，$\eta (\tau )\in [0,1]$ 仅反映调度本身。

---

4. 与经典 Scaling Laws 的联系

4.1 幂律 Pareto 前沿的必要条件

论文证明了以下定理：

定理 E.1：对于可微损失曲线 $L(c,p)$（$c$ 为计算量），若归一化损失曲线 $\ell$ 与 $p$ 无关（即发生 collapse），则 Pareto 前沿 $L^{\star }(c)={\min }_{p}L(c,p)$ 必须是幂律形式 $L^{\star }(c)=a{c}^{-\delta }$。

证明概要：

• 由归一化条件，$\frac{\partial \log L(c,p)}{\partial \log c}{|}_{c=c^{\star }(p)}=-\delta$ 对所有 $p$ 相同
• 由包络定理，$L^{\star }(c)$ 的导数等于该偏导数
• 因此 $\frac{d\log L^{\star }(c)}{d\log c}=-\delta$ 为常数
• 积分得 $L^{\star }(c)=a{c}^{-\delta }$

4.2 幂律损失曲线的充分性

对于常见的幂律 scaling law 形式：

$$L(t,p)=L_0+t^{-\mu }+p^{-\nu }$$

论文证明：

1. 计算最优性：平衡两个幂律项的导数，得到 $t^{\star }(p)=r^{-1/\mu }{p}^{\nu /\mu }$，其中 $r=\nu /\mu$
2. 精确塌缩：归一化后

$$\ell (x,p)=\frac{r{x}^{-\mu }+1}{r+1}$$

完全与 $p$ 无关！

4.3 更一般的多幂律形式

对于 $L(t,p)=L_0+{\sum }_{i=1}^m{a}_i{t}^{-{\mu }_i}{p}^{-{\nu }_i}$，论文证明：

定理 F.1：

1. 计算最优性迫使”平局”：至少两个幂律项共享最小指数 ${\beta }_i={\mu }_i\gamma +{\nu }_i$
2. 渐近塌缩：归一化损失曲线趋向于仅由这些主导项决定的极限，与 $p$ 无关
3. 有限大小误差：非主导项贡献 $O(p^{-ϵ})$ 的误差，其中 $ϵ$ 是主导项与次主导项的指数差

4.4 与计算最优 Scaling Laws 的对应

论文的数据集上拟合的 scaling law 形式为 $L(c,p)=L_0+a{c}^{-b}$，这正是 Kaplan 等人和 Hoffmann 等人提出的经典形式。塌缩现象与这一经典形式完美契合：

• $L_0$ 对应不可约损失，需从归一化中减去
• $c^{\star }(p)\propto p^{1+\gamma }$ 确定了最优计算量与模型大小的关系
• 幂律结构是塌缩发生的根本原因

---

5. 作为缩放质量指标的实践应用

5.1 诊断超参数配置

Supercollapse 提供了一种诊断缩放配置是否正确的方法：

5.1.1 模型参数化

使用 $\mu$P 协议（正确缩放初始化和学习率）时观察到 supercollapse，而使用恒定学习率跨宽度缩放时塌缩被打破。这表明 collapse 对参数化方式极为敏感，比单纯比较最终损失更敏感。

5.1.2 数据指数调优

当数据指数 $\gamma$ 偏离最优值时：

• 即使 scaling law 本身只有微小恶化
• 塌缩也会被显著破坏

这提供了一个新的应用场景：与其在 Pareto 前沿上拟合稀疏点来调整 $\gamma$，不如最大化塌缩质量来调整 $\gamma$，从而利用整个损失曲线的统计功效。

5.2 跨架构泛化

论文在多种设置中验证了 supercollapse 的普适性：

任务	架构	学习率调度
CIFAR-5M	Transformer	线性衰减、余弦衰减
Chess	Transformer	线性衰减
幂律 Fourier Features	MLP	线性衰减、余弦衰减

5.3 深度缩放

论文还探索了深度缩放（通过调整分支乘数），发现在深度变化时也存在一定程度的塌缩，但会出现小幅偏移。这提示深度缩放可能需要更细致的调整。

---

6. 理论框架总结

6.1 塌缩的两阶段解释

论文提供了两层解释机制：

第一层：模型间差异

$$V_{p,\omega }[\ell ]=\underset{\text{模型间差异}}{\underbrace{V_p{E}_{\omega }[\ell ]}}+\underset{\text{种子间差异}}{\underbrace{E_p{V}_{\omega }[\ell ]}}$$

• 幂律 scaling law 确保第一项在计算最优训练下消失
• SGD 噪声模型解释第二项在学习率衰减下的抑制

第二层：噪声抑制机制

归一化损失 $\ell (t)\approx \bar{\ell }(t)(1+\psi (t)-\psi (t^{\star }))$ 中，$\psi$ 是相对波动。关键是：

• 不是 $\psi (t)$ 本身，而是 $\psi (t)-\psi (t^{\star })$ 决定 $\ell$ 的方差
• 这捕捉了 $t$ 和 $t^{\star }$ 之间累积的优化噪声
• 学习率衰减精确地减少了这段期间的噪声

6.2 跨调度的普适性

论文证明，若塌缩在一种调度（如恒定学习率）下成立，则在满足一定条件时对其他调度也成立。这基于：

1. 梯度噪声比例的普适性：$\text{Tr}(\Sigma (\tau ))/L(\tau )\approx h(x)$，仅依赖归一化计算量 $x$
2. 调度间损失的近似关系（式18）

---

7. 核心洞察

1. 计算最优训练不仅仅是最终损失的最优，更是整个训练动态的最优。当 $\gamma$ 匹配最优值时，模型大小和训练时间的缩放达到精确平衡，导致损失曲线在整个训练过程中保持一致。

2. Supercollapse 揭示了一种新的缩放极限——模型大小和训练时间共同增长的极限，这与传统的无限宽度/深度极限不同。那些理论预测初始动态一致，但随训练进展逐渐发散；而塌缩揭示的联合极限在整个训练过程中保持一致。

3. 简单的 SGD 噪声模型出人意料地有效。这暗示学习率调度的本质效应是通过调节注入的噪声规模来影响损失曲线，而这种调节方式在不同架构和数据集间高度一致。

4. 归一化的选择至关重要。使用随机种子的经验最终损失而非期望最终损失进行归一化，是实现 supercollapse 的关键，这利用了优化轨迹中噪声的时间相关性。

---

参考文献