中文

结构化免逆自然梯度:内存高效且数值稳定的 KFAC

机器学习 2024-07-24 v4 机器学习

摘要

诸如 KFAC 之类的二阶方法可用于神经网络训练。然而,它们通常内存效率低下,因为其预条件 Kronecker 因子是稠密的;并且在低精度下数值不稳定,因为它们需要矩阵求逆或分解。这些局限性使得此类方法在现代混合精度训练中不受欢迎。我们通过 (i) 提出一种免逆的 KFAC 更新以及 (ii) 在 Kronecker 因子中施加结构来解决这些问题,从而得到结构化免逆自然梯度下降 (SINGD)。在现代神经网络上,我们表明与 KFAC 相比,SINGD 具有内存高效性和数值鲁棒性,并且即使在半精度下也通常优于 AdamW。我们的工作填补了现代低精度训练中一阶方法与二阶方法之间的差距。

关键词

引用

@article{arxiv.2312.05705,
  title  = {Structured Inverse-Free Natural Gradient: Memory-Efficient & Numerically-Stable KFAC},
  author = {Wu Lin and Felix Dangel and Runa Eschenhagen and Kirill Neklyudov and Agustinus Kristiadi and Richard E. Turner and Alireza Makhzani},
  journal= {arXiv preprint arXiv:2312.05705},
  year   = {2024}
}

备注

A long version of the ICML 2024 paper, updated the text about a related work