A Differential Equation for Modeling Nesterov's Accelerated Gradient Method: Theory and Insights
Weijie Su, Stephen Boyd, Emmanuel J. Candes
stat.ML, math.CA, math.OC
2015-03-04
把 Nesterov 加速梯度取步长极限,得到带 3/t 摩擦的二阶 ODE;常数 3 是加速的临界阻尼,速度重启还能在强凸下证明线性收敛。
训练一个大模型、拟合一个回归,底层常常归结为同一件事:最小化一个凸函数。在数据量爆炸之后,只用到梯度、不碰二阶信息(海森矩阵)的一阶方法重新吃香,因为它们每步便宜、扛得住大规模。一阶方法里,梯度下降的收敛率是 O(1/k),意思是要想把误差压低一个数量级,迭代次数得多花十倍。1983 年 Nesterov 提出加速梯度法(NAG),同样每步只算一次梯度,却把率提到 O(1/k²),这在只看梯度的一阶信息下界意义下是最优的,等于免费快了一档。
问题在于,这套加速「为什么有效」长期只有代数证明,没有直觉。Nesterov 递推里那个动量系数 (k−1)/(k+2) ≈ 1−3/k 是精心凑出来的常数,跑起来轨迹还会明显震荡(overshoot),这些都缺一个好懂的解释。这篇要做的,就是给 Nesterov 加速一个连续时间的物理模型,让它从一行魔法变成一个能复述、能推广的机制。
作者把步长 s 缩到 0,并令时间 t ≈ k√s(迭代数乘以步长的平方根),对 Nesterov 的离散递推取极限。结果是这样一个二阶常微分方程:
Ẍ + (3/t)Ẋ + ∇f(X) = 0,初始条件 X(0)=x₀、Ẋ(0)=0。
一个只看梯度的一阶格式,极限居然是二阶的,这是第一处反直觉。更关键的是它的物理读法:这正是质点在势能 f 里运动、受一个随时间衰减的摩擦 3/t 的方程,即一个阻尼振子。离散格式里那个动量系数 (k−1)/(k+2),对应到连续时间就是速度前面的摩擦系数 3/t。于是 Nesterov 加速可以当成一个被精心调过阻尼的物理系统来理解,后面所有结论都从这个 ODE 出发。
收敛率先对上了:这个 ODE 满足 f(X(t))−f ≤ O(‖x₀−x‖²/t²),和离散 Nesterov 的 O(1/k²) 同阶,而普通梯度下降只有 O(1/t)。
第一处洞见是震荡。把 3/t 当摩擦看,系统早期 3/t 很大,是过阻尼,平滑地滑向极小点;后期 3/t 变小,转成欠阻尼,开始来回震荡,幅度慢慢衰减。这正好解释了 Nesterov 轨迹后期那种「蝴蝶形」震荡和函数值曲线上一个个鼓包。在二次函数特例下,解能用 Bessel 函数写出来,震荡周期是可以定量算的。
第二处洞见是一个相变,也是全文最漂亮的结果。把方程推广成 Ẍ + (r/t)Ẋ + ∇f(X) = 0,加速率 O(1/t²) 当且仅当 r ≥ 3 才成立,离散版本同理,O(1/k²) 同样当且仅当 r ≥ 3。r < 3 拿不到加速率,r ≥ 3 时最坏情形常数在 r=3 处最小。也就是说,Nesterov 当年那个看似任意的「3」恰恰是加速的临界阻尼:多加摩擦(r>3)不破坏 O(1/t²),但常数会变差。
| 阻尼系数 r | 收敛率(一般凸) | 备注 |
| r < 3 | 拿不到 O(1/t²) | 摩擦不够,加速失效 |
| r = 3 | O(1/t²),最优常数 | 加速的临界点 |
| r > 3 | 仍 O(1/t²),常数更差 | 强凸下反而更快(可达 O(1/t^(2r/3))),常数随 r 涨 |
第三个结果是强凸下的实用收获。强凸问题里 Nesterov 反而吃亏:后期摩擦太小,过度震荡,卡在 O(1/poly(k)),不如普通梯度下降的线性收敛 O((1−µ/L)^k)。作者提出「速度重启」:一旦质点速度 ‖Ẋ‖ 开始下降,就把时间清零、把 3/t 复位,重新攒摩擦。定理 10 证明这样得到的解满足 f(X^sr(t))−f ≤ c₁L‖x₀−x‖²e^(−c₂t√L),即线性收敛,而且不需要事先知道强凸参数 µ(标准强凸版 Nesterov 必须知道条件数 µ/L,而 µ 在实际里几乎估不准)。
实验在四类问题上跑:500×500 强凸二次型(特征值 0.0011)、非强凸的 log-sum-exp、秩 5 矩阵补全(300×300、观测 10%)、稀疏 Lasso(5000×50000)。对比速度重启 Nesterov(srN)、梯度重启(grN)、原始 Nesterov(oN)、近端梯度(PG)。结论是两种重启都经验上呈线性收敛,把曲线上的鼓包抹平了,即便函数不满足定理假设(非光滑、非强凸)也成立;速度重启比梯度重启更稳。梯度重启为什么也线性收敛,至今没有证明,作者留作 open problem。
这是「把加速优化写成微分方程、变分原理」这条研究线的开山之作之一,后来 Wibisono、Wilson、Jordan 2016 的变分视角等都顺着这条路走。它的价值是把一坨代数凑出来的常数变成几何与物理直觉:加速就是合适的阻尼,震荡就是阻尼衰减后的欠阻尼,「3」就是临界点。
对从业者更直接:今天训练里的 momentum、Adam 这类加速器,底层逻辑都是这个阻尼振子。看懂了它,就知道为什么加速法会在窄谷里来回抖,为什么自适应重启(梯度重启或速度重启)能榨出强凸下的线性收敛又不依赖 µ。FISTA、AGD 的成熟实现现在普遍内置自适应重启,源头就在这。
作者自己交代得很清楚。第一,从连续 ODE 翻译回离散格式要靠参数调试和繁琐计算,目前没有一个把 ODE 性质系统映射到离散更新的通用理论,他们明说希望有人补上。第二,定理 10 给出的线性收敛常数 c₁、c₂ 不是最优(c₂=5Cµe^(−C̃µ/L)/(4L));作者猜测重启间隔 T=O(√(L/µ)),若成立 c₂ 能大幅改善,他们还自认 Lemma 12 紧而 Lemma 13 不紧。第三,实验在违反假设的函数上也呈线性收敛,说明现有理论覆盖不全,梯度重启的可证明线性收敛仍是 open problem。
适用范围也要说清:全文只处理凸和强凸问题,非凸(也就是深度学习)不在保证范围内。momentum 直觉经常被人借到非凸训练里,但本文不给任何那里的收敛保证。