用激活均值近似曲率，让优化更快更省内存
MAC: An Efficient Gradient Preconditioning using Mean Activation Approximated Curvature
- 用激活均值近似曲率矩阵，降低计算开销
- 在多种模型上比KFAC更快、更省显存
- 首次将此方法用于Transformer注意力层
二阶优化方法如KFAC通过利用损失曲面的曲率信息实现更优收敛，但计算成本高。本文分析了KFAC中用于层间Fisher信息矩阵（FIM）的两个组成部分：与激活相关的克罗内克因子和预激活梯度。基于对它们特征谱的实证观察，提出高效近似方法，构建出名为MAC的新型优化算法。据我们所知，MAC是首个将克罗内克分解应用于Transformer注意力层的FIM，并显式引入注意力分数进行预条件化的方法。我们还研究了MAC在非线性神经网络中的收敛性，给出了其收敛到全局最优的两个条件。在多种网络结构和数据集上的大量实验表明，该方法在准确率、端到端训练时间和内存占用方面均优于KFAC及其他先进方法。
原文摘要 · Abstract (English)
Second-order optimization methods for training neural networks, such as KFAC, exhibit superior convergence by utilizing curvature information of loss landscape. However, it comes at the expense of high computational burden. In this work, we analyze the two components that constitute the layer-wise Fisher information matrix (FIM) used in KFAC: the Kronecker factors related to activations and pre-activation gradients. Based on empirical observations on their eigenspectra, we propose efficient approximations for them, resulting in a computationally efficient optimization method called MAC. To the best of our knowledge, MAC is the first algorithm to apply the Kronecker factorization to the FIM of attention layers used in transformers and explicitly integrate attention scores into the preconditioning. We also study the convergence property of MAC on nonlinear neural networks and provide two conditions under which it converges to global minima. Our extensive evaluations on various network architectures and datasets show that the proposed method outperforms KFAC and other state-of-the-art methods in terms of accuracy, end-to-end training time, and memory usage.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。