A Defense of the Quadratic Model
Alexandru Meterez, Pranav Ajit Nair, Depen Morwani, Cengiz Pehlevan, Sham Kakade, Alex Damian
cs.LG, cs.AI, math.OC, stat.ML
2026-07-24
150M 参数预训练中途做泰勒展开,二次方代理在后段能预测真实 loss 达 10% 训练量,且 Hessian 尾部谱呈不受 batch/Adam 影响的普适幂律。
训练大模型时,学习率怎么排、批大小怎么选、用哪个优化器,背后都有一套从凸优化搬来的理论。这些理论几乎都建立在一个最简化的假设上:把损失函数在某点附近当成一个二次碗(quadratic model)来分析。LLM 的真实损失地形极其复杂、高度非凸,这个二次碗到底能不能描述真实训练,一直没人系统测过。
现有出路两难。模型太理想,假设在真实网络上站不住;纳入足够多结构,又变得数学上不可解。二次方模型是折中:足够简单可解,又能按 Hessian 谱逐模式分析。作者要确认的是,这个折中在 LLM 预训练里到底准不准、能准多久。
核心实验是给训练过程做局部体检。训一个 150M 参数(含 embedding 共 168M)的 OLMO 风格 transformer,在 FineWeb 的 3B token 上预训练,批大小取 1、64、1024。每跑完 10% 训练量存一个检查点,从该点出发把模型和损失做泰勒展开,用展开后的代理损失继续训 10% 训练量,再拿代理轨迹和真实轨迹比。
两档代理:
确认近似成立后,换两个放大镜看这个局部二次问题:Hessian 谱(用极深的 Lanczos quadrature,深度 m=1200,单个检查点的基要 750GB,分摊在 16 块 H100 上),以及局部线性稳定性。
近似在后段训练非常准。从 70% 之后的检查点出发,prox 和 quad 都能在整整 10% 训练量内把真实 loss 走向预测到 10% 误差以内(论文 Table 1)。早期训练(前 30%)两档代理都很快偏离,quad 尤其要等到 60% 之后才贴合。
| 指标 | 测得值 | 对照 |
| Gauss-Newton 尾部容量指数 α | ≈0.96 | — |
| Hessian 尾部容量指数 α | ≈0.65 | — |
| Gauss-Newton 源指数 β | 0.10.2 | 文献假设 β>1 |
| Hessian 源指数 β | 0.30.6 | 文献假设 β>1 |
| loss 幂律 γ(拟合 L=L+AT^-γ) | ≈1,共享 L≈2.502 | R²≈0.99 |
谱里有两个意外发现。其一,谱在大约 i=8192(正好是词表大小)处一刀切成头和尾,头部前 V 个特征向量几乎全压在 unembedding 层,最强的几个离群特征值恰好对应 the、逗号、句号等高频 token,与神经塌缩(neural collapse,网络末期各类别表征坍缩成对称几何结构)的预测吻合。其二也是最意外的:无论批大小怎么变、加不加 Adam 预条件,谱的尾部幂律形状完全一致,只整体平移一个常数。Adam 在尾部基本不起作用。
稳定性这块,从各检查点出发冻结 Adam 的二阶矩,用预条件 SGD 在学习率和批大小的 5×5 网格上扫、跑 10 个种子看哪格发散。结论是训练整体运行在稳定性边缘的 2 倍以内:小到中等批(B=1、64)处在随机稳定性边缘(边界沿对角线走,受温度 η/B 限制),大批(B=1024)处在确定性边缘(学习率撞上 2/λ₁(H) 的硬上限)。
这篇的价值是给一套沿用十年的做法补上经验地基。学习率调度、批大小缩放、Adam 为什么 work,背后那套二次方理论一直是在信仰之跃下挪用到真实网络的。现在有了 10% 训练量级的局部贴合证据,加上尾部谱的普适性,说明这些理论站得住。
对从业者最直接的提示:尾部谱不随 batch 和 Adam 变,能当成一个稳定的诊断量;β<1 这个发现,暗示大网络的偏差项会在某个临界时间后指数衰减,而不是无限缓慢下降,这跟 scaling law 的形式是吻合的。这是一篇打地基的论文,不是给你换优化器的。
作者自己把界限划得很清。所有结论只在 150M、单一 transformer 架构、Adam 上验证过,换更大的模型、换二阶优化器,谱指数和贴合窗口会不会变还不知道。早期训练近似明显更差,说明高阶项在前期扮演重要角色,二次模型在那里不适用。最扎眼的是:常数学习率加 EMA 的设置下,二次代理跟真实轨迹贴合得很差(论文 Figure 12),作者只推测这跟该设置下的边缘稳定性行为有关,没给答案。
还有一处存疑:稳定性实验里他们滤掉了预条件 Gauss-Newton 范数前 1% 的离群序列,理由是稀有 token 和死神经元会让 Adam 二阶矩趋零、再被它除导致尖峰。这个过滤是必要的工程处理,但它也意味着「训练运行在稳定性边缘 2 倍以内」是剔掉最极端尖峰之后的结论,而真实训练里这些尖峰恰恰是引发 loss spike 的地方。