[论文解读] Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis

George Thomas, César Laurent+3 more|arXiv (Cornell University)|Jun 11, 2018

Advanced Neural Network Applications参考文献 26被引用 37

一句话总结

本文提出EKFAC，一种Fisher信息矩阵的新近似方法，通过在Kronecker分解特征基中追踪对角方差，实现了比KFAC更优的曲率估计，同时支持高效的局部更新。实验表明，EKFAC在多种深度神经网络架构中均能加速优化过程并提升泛化性能，在训练损失和实际运行时间方面优于KFAC和SGD。

ABSTRACT

Optimization algorithms that leverage gradient covariance information, such as variants of natural gradient descent (Amari, 1998), offer the prospect of yielding more effective descent directions. For models with many parameters, the covariance matrix they are based on becomes gigantic, making them inapplicable in their original form. This has motivated research into both simple diagonal approximations and more sophisticated factored approximations such as KFAC (Heskes, 2000; Martens & Grosse, 2015; Grosse & Martens, 2016). In the present work we draw inspiration from both to propose a novel approximation that is provably better than KFAC and amendable to cheap partial updates. It consists in tracking a diagonal variance, not in parameter coordinates, but in a Kronecker-factored eigenbasis, in which the diagonal approximation is likely to be more effective. Experiments show improvements over KFAC in optimization speed for several deep network architectures.

研究动机与目标

• 为解决大规模深度神经网络中精确自然梯度下降因Fisher信息矩阵规模过大而导致的计算不可行性问题。
• 通过特征基变换，提供比KFAC更精确的Fisher信息矩阵近似，实现可证明的改进。
• 通过保持固定特征基的同时仅更新对角方差，实现曲率估计的高效局部更新。
• 证明在Kronecker分解特征基中的曲率估计可实现更快的收敛速度和更好的泛化性能。

提出的方法

• 该方法通过Kronecker分解特征基（KFE）对参数空间进行变换，使曲率近似作用于对角方差而非原始参数。
• 通过Kronecker分解协方差矩阵的特征分解计算KFE，将昂贵的计算分摊到多个迭代中。
• 算法在KFE中维护对角方差估计，从而无需重新计算特征基即可实现低开销的局部曲率矩阵更新。
• 在变换空间中使用自适应对角缩放（如RMSProp类方法）更新预条件矩阵，提升收敛稳定性。
• 该方法在理论上证明了其对Fisher信息矩阵的Frobenius范数逼近优于KFAC。
• 该方法支持全量和局部更新，可在计算成本与估计精度之间实现权衡。

实验结果

研究问题

• RQ1 在Frobenius范数意义下，Kronecker分解特征基中的曲率近似是否能提供比KFAC更优的Fisher信息矩阵逼近？
• RQ2 在KFE中追踪对角方差是否能实现比标准KFAC更有效且高效的优化？
• RQ3 在KFE中对曲率估计进行局部更新是否能在降低计算成本的同时保持性能？
• RQ4 与KFAC和SGD相比，EKFAC是否能在深度网络训练中提升优化速度和泛化性能？

主要发现

• 由于KFE中引入了特征值校正，EKFAC在Frobenius范数下对Fisher信息矩阵的逼近被证明优于KFAC。
• 在全连接自编码器和CIFAR-10数据集上，EKFAC在使用动量的SGD和KFAC中均实现了比它们更快的优化收敛速度。
• 在VGG-11和ResNet-34上，当KFE计算被分摊时，EKFAC-ra（采用局部更新）在每轮训练损失和实际运行时间方面均优于KFAC和SGD。
• EKFAC保持了优异的泛化性能，其模型在不同小批量大小下均实现了低于KFAC和SGD基线的验证误差。
• EKFAC中的局部更新计算高效且保持高精度，支持在大规模训练中实际部署。
• 该方法在不同超参数设置下表现稳健，当按训练轮次或计算预算选择模型时，EKFAC配置通常优于最佳的KFAC配置。