Simulating Training Spikes of Deep Neural Networks by Layer-wise Random Feature Models
Ran Tian, Courtney Paquette
cs.LG
2026-10-09
证明幂律重尾输入加近发散学习率就能在凸二次模型里复现 loss 尖峰,并给出从 checkpoint 预测尖峰的轻量模拟算法。
预训练大模型时,loss 偶尔猛地窜上去再落回来,这就是 training spike,PaLM 规模的训练里都有记录。它被当作病态信号:可能意味着训练不稳定、超参次优、最终模型变差;大模型普遍用低精度算术,剧烈波动还会直接触发数值问题。
难在复现。图 1 的对照很干净:普通波动跟着数据洗牌走,同一份 shuffle 下换初始化种子或微调结构,曲线形状基本不变;尖峰不遵循这个规律,只差初始化种子的两条轨迹,尖峰出现在不同步数。清洗数据 outlier 只能部分解决问题也是同一件事:清完之后 validation loss 里仍有随机尖峰,说明尖峰通常不是单个数据点造成的。此前理论大多围着 sharpness(Hessian 最大特征值 λmax)转,没有一个玩具模型的 loss 曲线能对上真实尖峰。
机理主张一句话:正常波动只由数据洗牌驱动,高维会抵消模型侧随机性;数据重尾时这个抵消失效,尖峰出现。
载体是 PLRF(Power-Law Random Feature)模型:固定随机矩阵映射输入,后面接二次损失的线性回归,输入各分量按 i^{-α} 幂律缩放,α 衡量重尾程度。没有深度结构,损失是凸的,但 α=0.9 时它就是能训出和真实 LLM 长相一致的尖峰。出尖峰要两个条件同时满足:重尾(大 α),加接近发散的大学习率;α=0.3 时调大学习率只会平滑发散,不出尖峰。
直觉上(SGD 情形)把更新拆成两部分:确定性的全批量下降,和逐步累积梯度噪声的随机游走项 r。高维下 r 的范数集中、方向随机,验证 loss 作为 r 的二次函数随之收敛到一条确定轨迹,这就是测度集中。前提是特征值不能太重尾:少数极大特征值压不住时,r 一旦随机指向这些方向,loss 就突然飙升,尖峰出现。学习率要够大让 r 攒出量,再大就发散,SGD 的发散阈值可以精确算出(式 2.2)。附录 A.6 证明测度集中只在 α≤0.5 成立,而自然语言正压在边界上:把输入 embedding 层看成 one-hot 到词向量的线性映射,特征值分布恰好是 unigram 概率的 Zipf 分布,对应 α=0.5。
接到真网络分两步。先把满足条件的层定义成独立线性单元,要求该层的输入和输出梯度都不依赖它自己的权重(多数 DNN 可这样拆,embedding 复用做输出头这类共享权重除外),再对验证 loss 在 step t 做二阶 Taylor 展开,得到每层的二次损失。定理 3.3 证明每层的 per-example 梯度可近似成 Random Feature 模型的梯度,靠 NTK 式的超高维隐空间线性化。Prediction 3.5 给出梯度二阶矩的显式公式,在 Transformer 上经验误差大多在 5% 以内,而公式两端的数值本身跨多个数量级。
模拟算法(Algorithm 3.1)每步采样匹配均值和二阶矩的高斯伪梯度,喂给优化器。初始只要四样东西,过一遍 validation set 即可:激活相关矩阵的特征分解、平均梯度、批梯度范数、Hessian 迹,前置量测合计约两小时 TPU 时间。之后全程不跑 DNN 反向,1500 步单层模拟 CPU 上 5 秒完成。
PLRF 玩具实验(D=1024,V=20000,batch 8):
| 设置 | 行为 |
| α=0.3,逐步调大学习率(SGD/Adam) | 平滑发散,无尖峰 |
| α=0.9,学习率接近发散阈值 | SGD 与 Adam 都先出尖峰再发散 |
Transformer 案例是重头。1.3B decoder(24 层)在 C4 上训到 Chinchilla 最优 token 数,共 25.5B tokens、TPUv4×32 上 40 小时,AdamW 超参继承自已调优配置,再把 β2 从 0.95 改成 0.999 制造尖峰。Spiky1 与 Spiky2 只差初始化种子。从 21k 步 checkpoint 模拟到 22.5k:两个 Spiky 配置在多个 FFN 层出现明显 bump(Spiky1 的 block 0/4/8 约在 22k 步,Spiky2 的 block 0/12 约在 21.5k 步),Normal 的模拟相对平滑。模拟出的尖峰位置也随种子变。实际训练在这一窗口只有 Spiky2 出现尖峰,模拟给出的是配置会不会尖峰的判断,不是尖峰时刻。
配套测量支撑机理:尖峰配置下激活相关矩阵的特征值分布显著更重尾;1.3B 模型的等效批量因子约等于 1,DNN 的逐点损失会放大梯度噪声,把大 batch 的降噪效果抵消掉。结论直接:训 Transformer 不能忽略梯度噪声,λmax 只是间接相关,起作用的是整个特征值谱。作者还从观测指标推导了一个 β2 日程,在他们的案例里提升了表现,定位是给理论做证明,不替代现有稳定化手段。
这是第一个把训练尖峰和重尾数据分布显式挂钩的研究,玩具模型是凸的,代码在 Colab 上两分钟跑完。对从业者有三点:从任一 checkpoint 出发,CPU 几秒即可预检某个配置会不会尖峰;激活相关矩阵特征值谱的幂律指数可以当训练健康度指标;大 batch 不等于低梯度噪声,等效噪声水平相当于 batch 1,而「β2 贴近 1 容易出尖峰」这条经验现在有了机理层面的解释路径。要说清的是,这是机理加工具的论文,没有新的稳定化方法,实用配方属于渐进改进。
作者自己列了几条:层分解不支持共享权重,留给未来;层级二次模型解释不了一阶 Taylor 项,一阶项尺度约是二阶的 20 倍,下降速度远快于二次模型的预测,尖峰在其中表现为 bump、jump 或 dive,二次模型给不出这个量级,需要跨层相关和更高阶展开;模拟中也没有施加 weight decay。
假设链不短:构造出的隐变量分布在短步数内被当作不漂移的 i.i.d. 高斯(假设 3.4),输出梯度协方差被折叠成各向同性(假设 3.6),模拟的只是权重矩阵的一个平均切片。案例规模只有 1.3B 一档、诱因只有 β2 一条,数据或架构变化触发的尖峰是否同一机理,论文没有验证。β2 到底怎么改变 Hessian 特征值谱,作者自己称之为未解之谜。