Hessian-Free

Hessian-Free 在深度学习和大规模机器学习中Hessian-Free无 Hessian 优化简称 HF是一种非常著名的二阶Second-order优化算法由 James Martens 于 2010 年将其引入深度学习领域最初用于攻克 RNN 的训练难题。传统的二阶优化如牛顿法因为需要显式计算和存储庞大的 Hessian 矩阵海森矩阵即二阶导数矩阵而无法应用在大模型中。而 Hessian-Free 的核心思想正如其名它能够完美利用二阶导数的全局曲率信息却不需要显式地计算出 Hessian 矩阵的任何一个元素。1. 为什么需要二阶优化对抗病态曲率在深度神经网络的训练中一阶梯度下降算法如 SGD、Adam经常会遇到病态曲率Ill-conditioned Curvature问题。想象一个极度狭长的“峡谷”地形在一个方向上坡度极陡梯度大在另一个方向上坡度极缓梯度小但这是通往最优解的方向。一阶算法会在这两个陡峭的峭壁之间疯狂震荡由于缺乏全局视野每次迈步都小心翼翼前进极慢。二阶算法牛顿法利用二阶导数Hessian 矩阵感知到了空间的弯曲程度通过H−1∇fH^{-1}\nabla fH−1∇f直接对梯度进行旋转和缩放一脚就能踩到峡谷的底部。然而如果网络有NNN个参数假设N107N 10^7N107Hessian 矩阵的大小就是N×N1014N \times N 10^{14}N×N1014。光是存储这个矩阵就需要几百 TB 的显存更别提去计算它的逆矩阵了。2. Hessian-Free 的核心魔法Hessian-Vector ProductHessian-Free 算法能够成立全靠一个数学上的绝技虽然我们无法承受计算整个 Hessian 矩阵HHH的代价但我们可以在极短的时间内计算出 Hessian 矩阵与任意向量vvv的乘积——HvHvHvHessian-Vector Product。在线性代数中HHH是一个N×NN \times NN×N的矩阵而vvv是一个N×1N \times 1N×1的向量两者的乘积HvHvHv依然是一个N×1N \times 1N×1的向量。怎么计算HvHvHv才能不碰HHH这里有两种主流方法本质上都是将计算复杂度从O(N2)\mathcal{O}(N^2)O(N2)降到了与前向传播相同的O(N)\mathcal{O}(N)O(N)① 有限差分法Finite Difference利用泰勒展开的定义可以通过两次梯度计算来逼近Hv≈∇f(wϵv)−∇f(w)ϵHv \approx \frac{\nabla f(w \epsilon v) - \nabla f(w)}{\epsilon}Hv≈ϵ∇f(wϵv)−∇f(w)​(其中ϵ\epsilonϵ是一个极小的正数。这意味着你只需要在参数www上加上一点点扰动ϵv\epsilon vϵv再算一次梯度两次梯度相减就能得到二阶信息。)② 自动微分法Pearlmutter 技巧现代深度学习框架如 PyTorch、JAX在底层实现了 Pearlmutter 技巧通过反向传播传播前向导数可以精确、无误差地计算出HvHvHv其计算开销大约仅相当于两到三次普通的前向反向传播。3. 算法全景HF 是如何跑起来的既然我们手里只有HvHvHv这个工具我们要怎么优化参数呢牛顿法的更新公式是HΔw−∇fH \Delta w -\nabla fHΔw−∇f既然不能直接解出H−1H^{-1}H−1Hessian-Free 选择把这个牛顿方程转化成一个子优化问题。在每一次大迭代Outer Loop中计算当前位置的一阶梯度g∇f(w)g \nabla f(w)g∇f(w)。构建一个二次目标函数q(d)dTg12dTHdq(d) d^T g \frac{1}{2} d^T H dq(d)dTg21​dTHd我们的目标是找到让q(d)q(d)q(d)最小的更新方向ddd。进入内循环Inner Loop使用共轭梯度法Conjugate Gradient, CG来迭代求解ddd。关键点CG 算法在求解线性方程组时只需要知道矩阵与向量的乘积。在 CG 的每一步中我们输入当前的搜索方向vvv通过 Pearlmutter 技巧计算出HvHvHv然后更新ddd。当 CG 迭代了十几步或几十步收敛后我们得到了近似的牛顿方向ddd。执行线搜索Line Search更新大模型的参数w←wαdw \leftarrow w \alpha dw←wαd。进入下一次大迭代。4. Hessian-Free 的优缺点优点每步信息量极大由于利用了精准的二阶曲率HF 可以跨越一阶算法无法逾越的极度平坦或高度不稳定的非凸区域如高度饱和的激活函数区。对超参数不敏感相比于 SGD 需要精细调优学习率、衰减率HF 具有自适应步长的特性收敛非常稳健。内存友好显存占用与普通一阶算法在同一个数量级O(N)\mathcal{O}(N)O(N)。缺点单步计算开销过大虽然外循环迭代次数少但由于每一个大步内部都要嵌套一个跑几十次 CG 的内循环每次 CG 都要算一次HvHvHv导致更新一次参数的时间显著长于 Adam。无法享受高度并行的流水线优化在现代分布式算力集群GPU/TPU上一阶算法如 AdamW的吞吐量流水线已经被优化到了极致。而 HF 密集的内循环和线搜索极难进行大规模的并行和异步分布式扩展。无法完美兼容随机小批量Mini-batch二阶信息对噪声极其敏感。如果每次 CG 迭代更换 Batch会导致 Hessian 矩阵不一致如果不更换 Batch在大规模数据集上又会陷入严重的局部过拟合。5. 现状与延伸在当前千亿参数大模型LLM制霸的时代标准的 Hessian-Free 已经很少直接用于全量参数的预训练。但是它的核心思想催生了许多极具生产力的变体K-FAC (Kronecker-factored Approximate Curvature)通过将 Hessian或 Fisher 信息矩阵近似分解为两个小矩阵的克罗内克积既保留了二阶优化对抗病态曲率的能力又获得了逼近一阶算法的计算速度被广泛应用于大规模分布式强化学习和视觉网络的加速中。大模型后量化与剪枝如 GPTQ/OBC在对已经训练好的大模型进行极致压缩如 4-bit 量化时需要计算参数之间的敏感度耦合。此时利用 Hessian-Free 中的HvHvHv乘积技术可以极其高效地计算出误差补偿确保模型剪枝后精度不崩盘。