把更新幅度与梯度范数解耦,会如何影响优化器的扩展?

归一化 SGD 先加速幂律学习,随后触及稳定性上限,决定更大批量或更宽模型是否划算。

直接答案

把更新幅度与原始梯度范数解耦,会让有效学习率随梯度收缩而增大,这首先加速幂律学习,但最终把主导模式推向稳定性阈值。在这套理论中,该反馈产生三个后期阶段——宽度受限、随机稳定性边缘和确定性稳定性边缘——只有其中一部分阶段允许更大的批量或更宽的模型在算力相当的情况下缩短串行训练时间。

1篇文献引用

本文由 WisPaper 驱动的搜索和论文分析生成。

归一化 SGD 把梯度收缩变成学习率上升

RMSProp、Adam、SOAP 和 Muon 等现代自适应优化器即使基础学习率固定,也会通过依赖梯度的预条件来重新缩放更新。为了把这种幅度控制与方向分离,本研究分析归一化 SGD:它保留随机梯度方向,但固定更新范数。在本文采用的归一化下,它等价于带有内生标量学习率 ηeff(t) = η / ‖g(t)‖₂ 的 SGD。

这一等价关系就是全部机制。随着越来越弱的谱模式被学到,梯度范数衰减,因此有效学习率增大。论文把这描述为一种张力:谱尾部控制瞬态加速,而谱边缘控制这种加速何时崩溃。

加速把 SGD 指数映射为 2r/(1−r)

论文在具有幂律教师和数据协方差的随机特征模型中研究归一化 SGD,并推导出一个动态平均场理论(DMFT),刻画损失如何联合依赖训练时间、模型宽度和批量大小。在该模型中,归一化最初通过把幂律指数 rSGD < 1 映射为 2rSGD/(1 − rSGD) 来加速 SGD。

同一映射给出两个边界情形。当 rSGD = 1 时,理论预测指数收敛;当 rSGD > 1 时,它预测形式上的有限时间收敛。这些是连续时间流的瞬态预测,而不是离散优化器的最终行为。

有限步长下同一反馈会破坏加速

加速不会持续。在有限步长下,不断增大的有效学习率最终把主导谱模式带到离散时间稳定性阈值,动力学随之偏离流预测。确定性稳定性边缘(EoS)对应主导 Hessian 模式的临界稳定,而随机稳定性边缘(EoSS)是由小批量波动施加的稳定性极限。

论文的随机矩阵分析让 EoS 条件变得具体:有效学习率满足 ηeff = 2/λmax,其中 λmax 是 Λ^(1/2) Aᵀ A Λ^(1/2) 的最大特征值。因此瞬态加速及其崩溃来自谱的不同部分——尾部带来增益,边缘终结增益。

后期分为宽度受限、EoSS 与 EoS 三个阶段

对 DMFT 求解长时间行为时,作者假设有效学习率稳定到一个稳态值 η∞eff,并找到三个候选稳态,对应三个阶段。阶段 1 是宽度受限,阶段 2 是随机稳定性边缘,阶段 3 是确定性稳定性边缘。

在过参数化模型中,一个依赖宽度的临界批量把 EoSS 与 EoS 分开;欠参数化模型则受有限宽度误差限制。论文还把分析扩展到无维度极限,在该极限中维度从每个方程中消失,只剩 N、B、η 和谱。在该极限下,阶段 1 和阶段 3 仍然存在,但阶段 2 不存在:对任意 N、B 和 η,宽度受限状态都最先到达,只留下一个在临界批量处从宽度受限到 EoS 的转变。

只有部分阶段能把额外资源换成更少步数

这些阶段决定了更大的批量或更宽的模型何时能在算力相当的情况下缩短串行训练时间。论文量化了在哪个阶段增加批量大小或宽度,可以通过减少达到目标损失所需的优化步数,来补偿每步多出的算力消耗。

负面结果和正面结果一样具体。在阶段 1,固定宽度和谱时,宽度约束固定了 η∞eff ρ,因此稳态有效学习率与批量大小无关;小步长渐近给出 ∂log η∞eff / ∂log α → 1/2,所以批量翻倍只能把达到相同有效时间所需的更新次数减少 √2 倍,并把总计算量增加 √2 倍。在阶段 2,固定批量大小和谱时,EoSS 条件独立于宽度地固定 y,响应方程给出 ∂log η∞eff / ∂log ν = α/(ν+α) < 1/2,因此宽度增益慢于平方根。在阶段 3,确定性阈值 η∞eff = 2/λmax(ν) 与批量大小无关,因此额外增加批量不提供稳态加速。

论文的正面主张更窄。在 EoSS 阶段,一个近乎线性的批量缩放区间能在算力几乎不变的情况下减少达到目标损失所需的步数。在宽度受限阶段,作者识别出稳态有效学习率随宽度超线性增长的谱区间,从而带来潜在的宽度加速——但这种增益最终会被 EoS 分支上的稳定性逆转,而当该分支在 ν = 1 以下到达时,增长与下降两个区间意味着效率最大值出现在欠参数化宽度处。

CIFAR-5M 上的线性化 ResNet 支持预测趋势

该理论用 CIFAR-5M 上的线性化 ResNet 实验检验,谱指数从神经正切核和目标对齐中测得。测得的指数给出 rSGD = 0.095,预测的瞬态指数为 2rSGD/(1 − rSGD) = 0.21。

DMFT 对 k ≤ 1024 采用实测谱,超出部分采用拟合幂律,并延续到 d = 10^10;取更大的 d 差异可忽略。扫描中宽度取 4、16、64、128 和 256(N 从 1,120 到 4,135,936),批量大小取 2、10、100、1,000 和 10,000(T 从 2.5·10^6 降到 500),学习率取 0.01、0.1、0.2 和 0.5,每个格子两个随机种子,共 241 次运行。评估报告在固定的 4,096 张图像留出集上、150 个对数间隔步处的均方误差,并记录每一步的 ‖gt‖₂ 和 ηteff;受内存限制,DMFT 最多积分到 T = 3800 步。

实验在探索的区间内支持预测的加速、崩溃和资源缩放趋势。它们是线性化网络实验,因此检验的是理论的谱预测,而不是大型网络的完整端到端训练。

按阶段而非资源规模判断批量与宽度收益

这篇论文的实用内容是阶段划分,而不是对更大批量或更宽模型的普遍背书。同样的资源增加,可能几乎算力中性、慢于平方根,或在稳态下毫无用处,取决于系统处于 EoSS 阶段、宽度受限阶段还是确定性 EoS 阶段。

作者明确指出,使用 norm-SGD 是为了分离梯度归一化与学习动力学之间的反馈,而不是把它推荐为实用优化器。因此可迁移的主张是关于机制的:当更新幅度与梯度范数解耦时,有效学习率是一个随梯度收缩而增大的内生量,而批量大小或宽度的算力价值取决于这种增长先触及哪个稳定性极限。论文中最具决策相关性的阈值,是过参数化模型中把 EoSS 与 EoS 分开的依赖宽度的临界批量,以及在无维度极限下只剩宽度受限到 EoS 转变这一发现。

本文引用的文献

1

Scaling

在幂律随机特征模型中,归一化 SGD 首先通过把 SGD 指数 rSGD < 1 映射为 2rSGD/(1 − rSGD) 来加速幂律学习,但在有限步长下,同一反馈把主导模式推向临界稳定,产生宽度受限、随机稳定性边缘和确定性稳定性边缘三个阶段,这些阶段决定了更大的批量或更宽的模型何时能在算力相当的情况下缩短串行训练时间。