← 返回 2026-06-02

统一神经缩放定律 Unified Neural Scaling Laws

Ethan Caballero, Priyank Jaini, David Krueger, Irina Rish 📅 2026-05-25 👍 7 2026-07-13 08:36
函数形式 多维缩放 大模型训练 模型外推 神经缩放定律

提出同时建模多维度变量并精确外推神经网络缩放行为的统一函数形式

前置知识

神经缩放定律 (Neural Scaling Laws)

研究深度学习模型性能指标(如损失、错误率)如何随规模因素——模型参数量、训练数据量、训练步数、计算量等——变化的经验规律。经典工作如 Kaplan et al. 2020 与 Hoffmann et al. 2022(Chinchilla)发现损失常随这些因素呈幂律 $L \propto x^{-c}$ 下降,由此可预测大模型在更大规模下的表现并指导资源分配。

本文正是要改进这些缩放定律的函数形式,使其能同时处理多个维度并精确外推;理解缩放定律的基本概念是读懂本文动机的前提。

条件熵不等式 (Conditional Entropy Inequality)

信息论基本不等式 $H(Y|X) \le H(Y)$,含义是在已知预测变量集合 $X$ 的条件下,目标变量 $Y$ 的不确定性不会增加。增加更多与 $Y$ 因果相关的预测变量 $X_i$ 只会让条件熵 $H(Y|X_1,\dots,X_m)$ 下降或保持不变。

本文用这个不等式论证'统一函数形式'的合理性:纳入的预测维度越多(参数、数据、步数、推理步数、超参数等),预测性能的不确定性就越低,这正是 UNSL 同时建模多维度变量的理论依据。

断点神经网络缩放定律 BNSL (Broken Neural Scaling Law)

Caballero et al. 2023 提出的单变量缩放形式,公式形如 $y = a + b\cdot x^{-c_0}\prod_{j=1}^{n}\left(1 + \frac{x^{c_j}}{d_j}\right)^{-f_j}$。它在双对数空间里表现为多段平滑相连的直线(折线),段间通过'断点 (break)'平滑过渡,能拟合单一缩放曲线上出现的拐点(如数据耗尽导致的瓶颈)。

UNSL 的最内层 Equation 4(MBNSL)就是把 BNSL 推广到多变量场景,理解 BNSL 是理解 UNSL 多层嵌套结构的基础。

多重对数空间 (Multi-log space)

把每个输入维度 $x_1,\dots,x_m$ 和输出 $y$ 都取对数后构成的 $(m+1)$ 维空间。幂律关系 $y = b\prod x_i^{-c_i}$ 在此空间中变成超平面(线性),便于用几何语言(超平面、hyperbreak、梯度、锐度)来描述缩放行为。

本文几乎所有几何解释都建立在多重对数空间之上,不熟悉这个视角会难以理解 Equation 4 中超平面、断点位置、梯度变化、sharpness 的含义。

计算最优缩放 / Chinchilla (Compute-optimal scaling)

Hoffmann et al. 2022 提出的方法:在固定计算预算 $C = C_0\cdot N\cdot D$($N$ 参数量、$D$ 数据量、$C_0 \approx 6$)下,求解使损失最小的参数与数据分配,发现以往模型普遍'参数过多、数据不足',应大致按 $N:D \approx 1:20$ 训练。

论文的 CF 基线函数形式即来自 Chinchilla,且 Appendix 12 专门讲如何从拟合好的 UNSL 用拉格朗日乘子法求计算最优分配,理解 Chinchilla 才能看懂这些对比与延伸。

普遍逼近定理 (Universal Approximation Theorem)

Leshno et al. 1993 等证明:带非多项式连续激活函数(如 softplus)的单隐层前馈网络可以逼近任意紧集上的连续函数。论文把 MBNSL 在多重对数空间写成 $\log y = \log b - \sum c_{i0}\log x_i - \sum f_j\cdot \text{softplus}(\dots)$,正是一个带 softplus 激活、线性跳跃连接、$n$ 个隐单元的单隐层网络。

论文 Appendix 15 用这一定理证明 A1/A2/A3/UNSL 的'极大表达力 (supremal expressivity)'完全相同——它们外推精度差异纯粹来自施加的对称性约束而非表达能力,这是理解消融实验结论的关键。

KFAC (Kronecker-factored Approximate Curvature)

一种二阶近似优化器(Martens & Grosse),把 Fisher 信息矩阵做 Kronecker 因子分解以降低存储和计算开销,常用于拟合复杂参数化模型。论文用 KFAC-JAX 实现来最小化均方对数误差 MSLE 拟合 UNSL 的众多常数参数。

了解拟合流程(20000 步、20 个随机种子、L2 正则)有助于评估复现难度,也能理解为何超参 $n, S, \lambda$ 需要单独的验证集调参。

研究动机

训练当今最先进的神经网络需要巨额算力和数据,而可选择的架构与方法众多。准确的性能预测对于在多种方案中选出真正能扩展到大规模的方案至关重要,因为'在小规模上表现最好的方法往往在大规模上无法保持优势'(Sutton 2019 的'苦涩教训';Tolstikhin et al. 2021 的 MLP-Mixer 等)。此外,预测大模型在规模放大时是否会涌现新能力,对 AI 安全和负责任部署同样关键。然而现有神经缩放定律的函数形式(如 Kaplan 2020、Hoffmann 2022 的 $y = a + b_1 x_1^{-c_1} + b_2 x_2^{-c_2}$,或 Muennighoff 2023 的 DC 形式)通常只建模一两个维度,且无法表达非单调过渡——例如过拟合、以及学习率、初始化权重标准差等超参数与性能之间的非单调关系。这些限制使得它们在多维同时变化时的外推精度很差。

本文的目标是本文目标是构造一个'统一函数形式' (Unified Neural Scaling Law, UNSL),能够同时、联合地建模多个维度(模型参数量、训练数据量、训练步数、推理步数、各种超参数)变化时的缩放行为,并在任务与架构集合上给出比已有形式显著更精确的外推结果。理想情况下,这个形式要能刻画那些其他形式根本无法表达的现象,特别是过拟合与超参数(学习率、初始化权重标准差)引起的非单调过渡。作者还希望借此探究'缩放行为可预测性的极限'。

与已有工作不同的是,本文的独特切入角度有两点。第一,理论层面:作者从条件熵不等式 $H(Y|X) \le H(Y)$ 出发,论证纳入更多因果相关的预测维度只会降低 $Y$ 的不确定性,因此'最统一'的函数形式应同时包含尽可能多的预测变量。第二,技术层面:作者没有发明全新函数,而是把单变量 BNSL(Caballero 2023)嵌套推广为多变量 MBNSL,再用四层递归结构(Equations 1–4)把'非瓶颈学习'、'瓶颈性能极限'、'超参数的对抗力 (oppositional force)'、'过拟合的对抗力'、'随机猜测等误性能极限'等不同物理含义的成分显式组合起来,并通过'加法对称性'(Section 2.1)专门刻画非单调过渡。关键在于,消融形式 A1/A2/A3 与 UNSL 的极大表达力完全相同,所以外推精度的提升来自更严格的对称性/期望 (desiderata) 约束,而非更强的拟合能力——这是一个非常巧妙的归纳偏置设计。

核心方法

整体思路是:先把性能 $y$(上游验证损失、下游零样本错误率等)和一组正预测变量 $(x_i)_{i=1}^m$(参数量、数据量、步数、推理步数、超参数值)放到'多重对数空间'里观察,发现真实缩放行为在该空间里是若干平滑相连的超平面(每段是个幂律),段间通过 'hyperbreak' 平滑过渡。技术路线由此展开:最内层 Equation 4 定义 MBNSL——把单变量 BNSL 推广到多变量,用 $n$ 个 hyperbreak 连接 $n+1$ 个超平面,每个超平面的梯度由 $(c_{i0k})$ 决定、第 $j$ 次梯度变化由 $(c_{ijk})$ 决定、断点位置由乘积 $\prod_{i\in M} x_i^{c_{ijk}} = d_{jk}$ 决定、过渡锐度由 $f_{jk}$ 控制。然后 Equation 3 把 MBNSL 分成'非瓶颈成分' $K(U_r, n_{r0}, \cdot)$(正常学习曲线)和若干'瓶颈成分' $K(\{t\}, n_{rt}, \cdot)$(每个对应被某个维度 $x_t$ 卡住的性能极限)。Equation 2 引入超参数的'对抗力'和误性能极限 $a_q$(如随机猜测的错误率)。最外层 Equation 1 引入过拟合的对抗力、全局性能下限 $a_0$(如不可约熵/Bayes 误差)以及指标相关的 $a_2$(交叉熵无界则 $a_2=\infty$,即 $a_2^{-1}=0$)。

核心创新不在于提高函数的表达能力,而在于通过'加法对称性'(Section 2.1)把归纳偏置写死进函数形式。对称性指的是 $y = b\prod x_i^{-c_{i0}} + g\prod x_i^{h_i}$(Equation 5),它等价于 Equation 4 的 $n=1$ 特例但带两个约束:相邻超平面间的梯度变化恒非负、且断点锐度被梯度变化量唯一决定。作者经验地观察到:非单调过渡(过拟合、学习率过大)总是被 Equation 5 刻画;而向零梯度区域的过渡则被 $h_i=0$ 的变体刻画。于是 UNSL 在所有发生加法的地方强制使用对应的对称性变体。这个设计与已有方法的本质区别是:CF/DC 等只是把几个幂律项相加,既不能表达非单调、也不能保证瓶颈极限只依赖单一维度;而 UNSL 通过四层嵌套 + 加法对称性,精确地编码了 8 条期望(Section 2.2),却保持与最简单消融形式相同的极大表达力——精度提升完全来自更强的结构先验。

方法步骤详情

完整步骤如下:(1) 收集缩放数据点 $(x_1,\dots,x_m, y)$,三角点用于拟合、圆点留作外推评估。(2) 选择 hyperbreak 数 $n$、对抗力极限数 $S$ 和 L2 正则系数 $\lambda$:把三角点再分成训练集与验证集,要求验证集的每个点都不同时大于训练集所有维度,取验证误差最低的 $(n, S, \lambda)$。(3) 用 KFAC-JAX 最小化均方对数误差 $\text{MSLE} = \frac{1}{N}\sum_{i=1}^{N}(\log(y_i+\epsilon) - \log(\hat y_i + \epsilon))^2$,其中 $\epsilon = 10^{-16}$,对指数施加相对权重为 $\lambda$ 的 L2 正则,跑 20000 步、20 个 LeCun Normal 初始化的随机种子,选训练误差最低的种子。(4) 拟合好后把验证集并回训练集,对圆点(留出点)计算外推 RMSLE 及根标准对数误差。(5) 如需求计算最优分配,对拟合好的 UNSL 解拉格朗日方程组 $\frac{\partial y}{\partial x_\ell}/x_\ell + \lambda C = 0$(计算维度 $\ell\in\mathcal D$)、$\frac{\partial y}{\partial x_v} = 0$(非计算超参数 $v\in\mathcal H$)、$C - C_0\prod_{\ell\in\mathcal D} x_\ell = 0$。输入是数据点,输出是拟合常数 $(a,b,c,d,f)$ 及其在任意新规模下的预测 $y$。

技术新颖性

技术新颖性体现在几方面。其一,把单变量 BNSL 严格推广为多变量 MBNSL(Equation 4),用乘积 $\prod x_i^{c_{ijk}}$ 自然实现'断点位置随其他维度平移'(Desideratum 2),这是 CF/DC 做不到的。其二,'加法对称性'(Equation 5)把非单调过渡的锐度与梯度变化量绑定,这与 BNSL 中 $f$ 解耦锐度的设计正好互补,二者在不同物理现象上各司其职。其三,用'对抗力 (oppositional force)'这一统一框架同时建模过拟合和超参数:每种对抗力都是非负、并 opposition 地作用于'好的学习'与'坏的过拟合'(Desideratum 8),对应不同误性能极限 $a_q$(随机猜测 vs. 远差于随机猜测)。其四,理论保证:Appendix 15 证明 A1/A2/A3/UNSL 极大表达力完全相同(都用 softplus 单隐层网络逼近任意正连续函数),所以 Table 1 中 UNSL 大幅领先纯粹归因于更完整的对称性约束,这是一个非常干净、可证伪的设计哲学。

An illustration of a Unified Neural Scaling Law (UNSL) with two input dimensions x1 and x2
Figure 1: An illustration of a Unified Neural Scaling Law (UNSL) with two input dimensions x1 and x2
An illustration of an example configuration of Equation 5 with two input dimensions
Figure 2: An illustration of an example configuration of Equation 5 with two input dimensions
Extrapolation Results on (n,k)-sparse parity; Right figure fits the form when f is constrained to be 1
Figure 7: Extrapolation Results on (n,k)-sparse parity; Right figure fits the form when f is constrained to be 1
Extrapolation Results of UNSL functional form (MLP with large weight init std)
Figure 9: Extrapolation Results of UNSL functional form (MLP with large weight init std)
Extrapolation Results of y=(a+sum b_t x_t^{-c_t})^{-1} on ImageNet test error of ViT/JFT
Figure 10: Extrapolation Results of y=(a+sum b_t x_t^{-c_t})^{-1} on ImageNet test error of ViT/JFT

实验结果

核心发现可以分成几组。首先是汇总(Table 1):在下游图像分类域,UNSL 在 60.87% 的任务上外推最优,次优形式(A3)仅 21.74%,而经典 CF/DC 都是 0%;在语言(上游+下游)域,UNSL 占 88.89%,次优仅 11.11%,CF/DC 同样为 0%。其次是三变量视觉(Table 2):在 Birds 上 UNSL 的外推 RMSLE 为 $4.03\times10^{-2}\pm5.51\times10^{-3}$,而 DC 高达 $2.65\times10^{-1}$,A3 为 $4.77\times10^{-2}$;ImageNet 上 UNSL 为 $1.70\times10^{-2}$,DC 为 $2.54\times10^{-1}$,差距约一个数量级。二变量视觉(Table 3)逐模型看,例如 ImageNet+ViT/B/16 上 UNSL 为 $8.57\times10^{-3}$,DC 为 $3.69\times10^{-1}$,A1 为 $1.41\times10^{-2}$;多数情况下 UNSL 与 A3 接近且都远好于 DC。语言三变量(Table 4):UNSL 为 $7.82\times10^{-3}\pm1.33\times10^{-3}$,DC 为 $6.24\times10^{-2}$,A3 为 $1.49\times10^{-2}$。语言二变量(Table 5):例如 Upstream+Transformer+Chinchilla 上 UNSL 为 $3.81\times10^{-3}$,CF 为 $1.72\times10^{-2}$;Upstream+Recurrent+Constant 上 UNSL 仅 $4.66\times10^{-3}$,而 A1 高达 $2.65\times10^{-1}$。Desideratum 7 的三变量(学习率×初始化标准差×步数)实验(Table 6):UNSL 外推 RMSLE 为 $5.11\times10^{-2}$,去掉对称性的基线为 $8.09\times10^{-2}$。此外 Figure 5 显示仅用 9 个点 UNSL 就能合理外推,Figure 6 显示它能外推到比拟合区域大一个数量级的规模。Appendix 还展示了 UNSL 在强化学习(Figure 13, StarPilot-hard)、推理时缩放(Figure 14, MATH 上的 CoT 长度)、宽度×深度(Figure 15)、批量大小(Figure 16)等多种多维场景下都能精确外推。

Percentage of tasks by domain where each functional form is the best for extrapolation
Table 1: Percentage of tasks by domain where each functional form is the best for extrapolation
Extrapolation Results for trivariate scaling behavior of downstream vision performance
Table 2: Extrapolation Results for trivariate scaling behavior of downstream vision performance
Extrapolation Results for bivariate scaling behavior of downstream vision performance
Table 3: Extrapolation Results for bivariate scaling behavior of downstream vision performance
Extrapolation Results for trivariate scaling behavior of language performance
Table 4: Extrapolation Results for trivariate scaling behavior of language performance
Extrapolation Results for bivariate scaling behavior of downstream (and upstream) language performance
Table 5: Extrapolation Results for bivariate scaling behavior of downstream (and upstream) language performance
Results on trivariate scaling behavior in which Desideratum 7 is true empirically
Table 6: Results on trivariate scaling behavior in which Desideratum 7 is true empirically
UNSL accurately Extrapolating Downstream Performance
Figure 3: UNSL accurately Extrapolating Downstream Performance
Varying the number of observed points used for fitting UNSL functional form from 9e0 to 9e2
Figure 5: Varying the number of observed points used for fitting UNSL functional form from 9e0 to 9e2
Extrapolation of UNSL ... scales an order of magnitude larger in multiple dimensions
Figure 6: Extrapolation of UNSL ... scales an order of magnitude larger in multiple dimensions
Extrapolation Results of UNSL on scaling behavior of reinforcement learning (StarPilot hard)
Figure 13: Extrapolation Results of UNSL on scaling behavior of reinforcement learning (StarPilot hard)
Extrapolation Results of UNSL on scaling behavior of inference scaling (CoT length, MATH)
Figure 14: Extrapolation Results of UNSL on scaling behavior of inference scaling (CoT length, MATH)
Extrapolation Results of UNSL as width and depth simultaneously vary
Figure 15: Extrapolation Results of UNSL as width and depth simultaneously vary
Extrapolation Results of UNSL as batch size and number of training steps simultaneously vary
Figure 16: Extrapolation Results of UNSL as batch size and number of training steps simultaneously vary
查看结构化数据
任务指标本文基线提升
下游图像分类外推(三变量:数据量×步数×参数量) RMSLE(越低越好) Birds 4.03e-2;ImageNet 1.70e-2 DC: Birds 2.65e-1, ImageNet 2.54e-1;A3: Birds 4.77e-2, ImageNet 2.20e-2 相比 DC 降低约一个数量级,相比次优 A3 进一步降低约 15–23%
语言外推(三变量:参数×处理 token×数据集 token) RMSLE 7.82e-3 DC 6.24e-2;A3 1.49e-2 相比 DC 降低约 8 倍,相比 A3 降低约一半
语言外推(二变量 Upstream/Recurrent/Constant) RMSLE 4.66e-3 A1 2.65e-1;A2 1.92e-2;CF 3.13e-2 相比 A1 降低约 57 倍
三变量超参数缩放(学习率×初始化标准差×步数) 外推 RMSLE 5.11e-2 去加法对称性基线 8.09e-2 降低约 37%
整体外推最优占比 各域任务百分比 视觉 60.87%,语言 88.89% 次优:视觉 21.74%(A3),语言 11.11%(A2);CF/DC 全为 0% 视觉约 2.8 倍,语言约 8 倍

局限与改进

作者明确承认的主要局限是'缩放行为可预测性的极限'(Section 5):要精确外推越过某个 hyperbreak,拟合点(凸包)到该 hyperbreak 的最短距离必须足够小。换言之,如果一个拐点(如 grokking、涌现能力、数据耗尽转折)尚未在拟合区域内被观察到,UNSL 也无法可靠预测它之后的行为——Figure 4 的稀疏奇偶校验实验就是例证:该任务前几百步几乎看不到损失下降,UNSL 只有在拟合点足够靠近真正的'断点'后才能外推。我自己的观察补充几点:(1) 实证验证的领域仍有限,主要是视觉(JFT-300M 子集上的 ViT/MiX/BiT)、语言(Muennighoff 2023 与 Hoffmann 2022 数据)和几个合成 MLP 任务,对扩散模型、多模态、RLHF 后训练等场景尚未检验;(2) 函数形式含大量需拟合的常数,虽然 Appendix 13 表明 9 个点可用,但选择 $n, S, \lambda$ 仍需单独验证集和人工/网格调参;(3) 拟合本身需要 20000 步 × 20 种子,对超大规模数据点集的成本未讨论;(4) $S>1$ 的极端场景(如 Figure 9 同时大学习率、极小步数)需要特殊处理,作者也承认这类场景'相对人为 (contrived)'。

独立分析的弱点

弱点一:外推强依赖'已观测到拐点'。改进方向是把 hyperbreak 的位置建模为先验分布(贝叶斯版本),或在拟合时主动探测潜在断点区域(主动学习)。弱点二:$(n, S, \lambda)$ 的选择是离散+连续混合超参,需要划分验证集并禁止验证点同时大于训练点,流程繁琐且对数据点布局敏感;可用可微的超参学习或基于信息准则的自动选择替代。弱点三:函数形式非常复杂、常数众多,物理可解释性弱(虽有'对抗力'等命名,但 $c_{ijk}, f_{jk}$ 等仍难直接对应机理);可结合理论分析(如损失景观谱、神经切核)给参数赋予机理含义。弱点四:仅在有限域验证,对扩散、多模态、后训练、MoE 等现代场景缺数据,应扩展评测。弱点五:拟合成本(20000 步×20 种子)对反复实验不友好,可探索更好的初始化、跨任务复用拟合常数的迁移策略,或可微架构搜索式的自动结构选择。

未来方向

作者直接提出的方向是继续探究'可预测性的极限'——即量化'到 hyperbreak 的距离'与外推误差的关系,并理解哪些拐点(涌现、grokking)原则上是不可预测的。基于本成果可延伸的方向包括:(1) 把 UNSL 用于 AI 安全中的能力涌现预测,提前估计危险能力出现的规模;(2) 把 Appendix 12 的计算最优求解推广到含推理步数(test-time compute)的全生命周期成本最优,兼顾训练与推理算力;(3) 用 UNSL 替代 Chinchilla 公式做大规模训练的资源规划,量化参数/数据/步数/批量/超参数的联合最优;(4) 理论上解释为什么加法对称性会经验成立——是否源于某种损失景观或泛化的对称性;(5) 把 UNSL 嵌入神经架构搜索或 AutoML,作为可微的性能预测器;(6) 扩展到非幂律维度(如数据质量、课程顺序、量化位宽)。

复现评估

复现评估总体较好。实现使用开源的 KFAC-JAX(Botev & Martens 2022),代码在 Appendix 20 给出。损失为标准 MSLE($\epsilon = 10^{-16}$)加 L2 正则,流程清晰:拟合 20000 步、20 个 LeCun Normal 随机种子、选训练误差最低的种子,$(n, S, \lambda)$ 用内部验证集选。数据方面,视觉部分来自公开的 ViT/16(Zhai 2022)与作者通信获得的 Alabdulmohsin 2022 子集;语言部分来自 Muennighoff 2023 与 Hoffmann 2022 的公开/通信数据;合成实验(稀疏奇偶、Greydanus & Kobak 2024 数据集)可自行训练(batch size 80000、单 epoch、Adam,$\beta_1=\beta_2=0$)。算力门槛不高——拟合是曲线拟合而非训练大模型,单机 GPU 应可完成;但部分真实缩放数据(Hoffmann 通信数据)需向原作者申请。主要复现难点在于:四层嵌套公式的正确实现、离散 $(n, S)$ 的搜索、以及 20 种子的随机性管理。整体属于'代码+数据基本齐备、中等工程量即可复现'的级别。