SoftServe: A Scalable Quasi-Newton Method for Deep Learning
Joohwan Ko, Tetiana Parshakova, Diana Cai, Robert M. Gower
cs.LG, cs.AI
2026-10-02
软化 secant 约束、改用 Kronecker 结构,SoftServe 让拟牛顿法扩展到 136M 模型,病态任务常胜 Adam 与 Muon。
BFGS 这类拟牛顿(quasi-Newton)方法是大尺度凸优化的老牌强者,深度学习却几乎不用。卡点有两个。一是非凸:loss 曲面存在负曲率方向,逆 Hessian 的估计一旦带上负特征值,二阶更新方向可能变成上升方向。BFGS 靠曲率条件 sᵀy>0 保住正定性,非凸时该条件会破,硬守 secant 等式反而会把估计矩阵推成不定的;教科书解法是 line search,但它需要全批量 loss,在随机小批量训练里不可行。二是规模:稠密估计要存 D² 个数,L-BFGS 存 10 组差分向量相当于 20 倍模型状态的内存,除玩具模型外都不可用。
这两个卡点在病态任务上代价最大。PINN、RNN 这类问题的 loss Hessian 特征值量级悬殊,一阶方法走得慢,曲率信息有真实收益。
SoftServe 的底子是 Berglund 等 2025 年的 Soft QN:把 secant 等式 Hy=s 从硬约束改成最小二乘惩罚项,再加一个相对上一步估计的 log-det 散度正则。该目标永远有唯一正定解,负曲率出现也不怕;曲率为正且 λ 趋于无穷时退回 BFGS。代价是稠密求解要 O(D²) 的计算和内存,深度学习用不起。SoftServe 做的事是把这个目标限制在结构化矩阵族上求解:
解这些方程不走 SVD,因为 SVD 的常见实现不适合 GPU;改用 coupled Newton–Schulz 迭代,只做矩阵乘法和加法,实现里平方根跑 18 步迭代、求逆 10 步。这套数值基建与 Muon 的扩展路径同源。工程上还有三件标配:动量加 Nesterov 外推或偏差修正;参数更新做归一化,即 H 度量下的范数约束下降;曲率估计每 K=10 步刷新一次,窗口起点快照参数、梯度和全部随机性(小批量索引、dropout mask 等),窗口末尾用同一随机性重算梯度,拼出一对干净的 secant 数据,保证梯度差分反映的是同一个被采样 loss 的曲率。回放每 10 步多花一次梯度计算,约 10% 预算开销,论文的对比都把这笔账计入。
基线是 Adam、Muon、SOAP、K-FAC、K-BFGS(L),确定性任务加 L-BFGS;另有 SoftServe-I(把 H 固定为单位阵的消融,即纯归一化动量)。所有方法在等梯度预算下对比,各自扫学习率,三个新种子出结果。
| 任务 | 指标 | 结果 |
| RNN Adding(321 参数) | 测试 MSE | SoftServe-Kron 6.6×10⁻⁶;Muon 1.5×10⁻⁴;SoftServe-I 1.0×10⁻² |
| MNIST 自编码器(2.8M) | 训练目标 | 全批量下最低;批大小 1000 时与 Muon 接近 |
| 小型 PINN(三个方程 × 固定或重采样点) | 最终训练目标 | 6 组设置中 5 组最低;例外是固定点 Wave,Adam+L-BFGS 最低 |
| PirateNet(0.73M) | 平均损失 | KdV 上最低;Allen–Cahn 与 Muon 相当 |
| 物理扩散模型 PIDM(136M) | 10 万次梯度评估的训练损失 | 低于 Adam、Muon、SOAP |
| 16M GPT(200M FineWeb token) | 验证 NLL | SoftServe-Kron 4.661;AdamW 4.786;Muon 4.440;SOAP 4.416 |
RNN Adding 的 loss Hessian 在初始化处数值上近乎奇异,四个种子全部如此。这个任务里预条件是必需品:不带曲率的 SoftServe-I 收在 10⁻²,带上 Kronecker 曲率降到 10⁻⁶ 量级。
语言模型是明确的负结果:赢 AdamW,输 Muon 和 SOAP。
对做 AI4Science 的人,这是一个现在就能试的工具。PINN、PDE 求解、物理耦合生成模型这类病态任务上,曲率带来的收益比通用任务大,代码已开源。对做优化器研究的人,它示范了一种干净的分工:正定性由变分目标在构造上保证,不需要 damping、line search 这类补救;可扩展性交给 Kronecker 结构加 Newton–Schulz,与 Muon 共用同一套 GPU 数值积木。它还不是通用优化器:语言模型预训练上没有优势,136M 是测试上限。整体更接近一次扎实的系统工作,把「拟牛顿进深度学习」这条线推进到能跑 136M 模型的程度,适用边界也画得清楚。
作者自己列了两条:λ 在整个 run 内固定,自适应 λ 留作未来工作;现实设定下的收敛性保证还是开放问题。实验读下来还有几点。优势集中在病态任务,语言模型上落后 Muon 和 SOAP。Kron 版的精确刷新是层维度的立方复杂度,靠每 10 步刷新一次摊销,wall-clock 对比放在附录,正文只给梯度预算口径。λ 要逐任务调,选中的值从 1 跨到 9999,四个数量级,换任务的重调成本不低。回放机制要求保存并重放全部随机性,对数据管线的工程约束比一般优化器多。