[论文解读] M-FAC: Efficient Matrix-Free Approximations of Second-Order Information

Elias Frantar, Eldar Kurtic+1 more|arXiv (Cornell University)|Jul 7, 2021

Stochastic Gradient Optimization Techniques参考文献 35被引用 4

一句话总结

本文提出 M-FAC，一种无需存储完整 Hessian 矩阵的矩阵自由方法，用于在经验 Fisher 近似下高效计算逆 Hessian-向量积（IHVP），采用线性时间算法。该方法显著降低了现有二阶方法的计算成本，实现了神经网络剪枝与优化的最先进性能。

ABSTRACT

Efficiently approximating local curvature information of the loss function is a key tool for optimization and compression of deep neural networks. Yet, most existing methods to approximate second-order information have high computational or storage costs, which can limit their practicality. In this work, we investigate matrix-free, linear-time approaches for estimating Inverse-Hessian Vector Products (IHVPs) for the case when the Hessian can be approximated as a sum of rank-one matrices, as in the classic approximation of the Hessian by the empirical Fisher matrix. We propose two new algorithms as part of a framework called M-FAC: the first algorithm is tailored towards network compression and can compute the IHVP for dimension $d$, if the Hessian is given as a sum of $m$ rank-one matrices, using $O(dm^2)$ precomputation, $O(dm)$ cost for computing the IHVP, and query cost $O(m)$ for any single element of the inverse Hessian. The second algorithm targets an optimization setting, where we wish to compute the product between the inverse Hessian, estimated over a sliding window of optimization steps, and a given gradient direction, as required for preconditioned SGD. We give an algorithm with cost $O(dm + m^2)$ for computing the IHVP and $O(dm + m^3)$ for adding or removing any gradient from the sliding window. These two algorithms yield state-of-the-art results for network pruning and optimization with lower computational overhead relative to existing second-order methods. Implementations are available at [9] and [17].

研究动机与目标

• 解决深度神经网络中二阶优化与剪枝方法的高计算与存储成本问题。
• 在不存储大矩阵的前提下，实现经验 Fisher 近似下 IHVP 的精确、高效计算。
• 开发与模型维度 d 呈线性增长的算法，适用于大规模深度神经网络。
• 通过利用曲率信息的更新，提升网络剪枝的准确率与优化收敛速度。
• 在 CIFAR-10、GLUE 和 SQuADv2 等标准基准上验证所提方法的有效性。

提出的方法

• 该方法利用 Woodbury-Sherman-Morrison 公式，通过梯度外积的秩一分量递归更新逆 Fisher 矩阵。
• 对于静态算法，预计算成本为 O(dm²)，IHVP 计算成本为 O(dm)，对角线元素查询的单位成本为 O(m)。
• 对于动态算法，维护 m 个最近梯度的滑动窗口，支持增量更新，IHVP 计算成本为 O(dm + m²)，添加/移除梯度的成本为 O(dm + m³)。
• 算法仅操作向量与标量，避免显式存储矩阵，确保在模型维度 d 上的时间与空间复杂度均为线性。
• 该方法利用经验 Fisher 矩阵为秩一矩阵之和的特性，通过递归更新实现精确求逆。
• 该方法被应用于网络剪枝（通过 Optimal Brain Surgeon）与优化中的预条件 SGD。

实验结果

研究问题

• RQ1 我们能否在不存储完整 Hessian 矩阵的前提下，高效计算经验 Fisher 近似下的 IHVP？
• RQ2 无矩阵、线性时间的 IHVP 算法能否在神经网络剪枝与优化中实现最先进性能？
• RQ3 在动态优化设置中，逆 Hessian 的滑动窗口近似精度如何？
• RQ4 与现有二阶方法相比，该方法在计算效率与模型性能方面是否更具优势？
• RQ5 该方法能否在无需显著超参数调优的情况下有效应用于 BERT 等大规模模型？

主要发现

• M-FAC 在 ResNet20 上的 CIFAR-10 剪枝任务中达到最先进性能，计算成本显著低于先前方法。
• 在 SQuADv2 基准上，M-FAC 超过 HuggingFace 的 Adam 基线，在 BERT-tiny 上达到 52.18 的 F1 分数（对比 49.99），在 BERT-mini 上达到 61.35（对比 58.13）。
• 在 GLUE 基准上，M-FAC 在 BERT-tiny 上匹配或超越调优后的 AdamW 基线，SST-2 准确率达 83.4%（对比 83.2%），MRPC F1 达 81.9%（对比 81.1%）。
• 动态与静态 IHVP 方向之间的余弦相似度在整个训练过程中保持较高水平（>0.95），验证了滑动窗口近似的有效性。
• 该方法实现 O(dm) 的 IHVP 计算成本与 O(m) 的单位元素查询成本，支持高效剪枝与优化。
• 代码实现已发布于 [9] 和 [17]，支持在真实世界深度学习流程中的实际部署。