[论文解读] Large-Scale Distributed Second-Order Optimization Using Kronecker-Factored Approximate Curvature for Deep Convolutional Neural Networks

Kazuki Osawa, Yohei Tsuji+4 more|arXiv (Cornell University)|Nov 29, 2018

Advanced Neural Network Applications参考文献 30被引用 19

一句话总结

该论文提出了一种大规模分布式二阶优化方法，利用Kronecker-Factored Approximate Curvature（K-FAC）在ImageNet上训练ResNet-50，采用极大规模的mini-batch。通过采用半精度计算、对称Kronecker分解以及对Batch Normalization层的Fisher信息矩阵（FIM）进行对角近似，该方法仅用978次迭代（100个周期）就实现了75%的top-1准确率，mini-batch大小为131,072——证明了二阶方法在泛化性能上可与一阶方法相媲美，同时收敛速度更快。

ABSTRACT

Large-scale distributed training of deep neural networks suffer from the generalization gap caused by the increase in the effective mini-batch size. Previous approaches try to solve this problem by varying the learning rate and batch size over epochs and layers, or some ad hoc modification of the batch normalization. We propose an alternative approach using a second-order optimization method that shows similar generalization capability to first-order methods, but converges faster and can handle larger mini-batches. To test our method on a benchmark where highly optimized first-order methods are available as references, we train ResNet-50 on ImageNet. We converged to 75% Top-1 validation accuracy in 35 epochs for mini-batch sizes under 16,384, and achieved 75% even with a mini-batch size of 131,072, which took only 978 iterations.

研究动机与目标

• 解决由于有效mini-batch大小增加而导致的深度神经网络大规模分布式训练中的泛化差距问题。
• 证明像K-FAC这样的二阶优化方法可以在ImageNet上实现与高度优化的一阶方法（例如，带有自适应学习率的SGD）相当的泛化性能。
• 通过采用激进的近似技术（如对角FIM和陈旧Fisher矩阵更新），降低分布式设置下K-FAC的计算和内存开销。
• 在保持高验证准确率和快速收敛的前提下，实现高达131,072的mini-batch大小训练。

提出的方法

• 作者采用半精度浮点数算术，实现同步全worker分布式K-FAC优化器，以减少内存和通信开销。
• 利用曲率近似中Kronecker因子的对称性，最小化冗余计算和通信。
• 将Batch Normalization层的Fisher信息矩阵（FIM）近似为对角矩阵，将ResNet-50的内存消耗从1017 MiB降低至587 MiB。
• 在500次迭代后降低Fisher矩阵的更新频率（陈旧FIM更新），显著降低计算成本，且未影响准确率。
• 采用混合数据-模型并行策略，在1,024块Tesla V100 GPU上实现扩展，支持高效的超大批次训练。
• 采用动态FIM更新间隔：前13个周期每1次迭代更新一次，之后每20个周期更新一次，以平衡准确率与效率。

实验结果

研究问题

• RQ1 当使用极大规模的mini-batch训练深度网络时，像K-FAC这样的二阶优化方法是否能与一阶方法（如学习率缩放的SGD）一样实现良好的泛化性能？
• RQ2 如何在不降低模型性能的前提下，降低K-FAC在大规模分布式训练中的计算和内存开销？
• RQ3 将Batch Normalization层的Fisher信息矩阵（FIM）近似为对角矩阵，对训练稳定性和准确率有何影响？
• RQ4 能否有效利用陈旧的Fisher矩阵更新来减少计算量，同时保持收敛性和泛化能力？
• RQ5 在使用大mini-batch训练时，ResNet-50的FIM结构如何演变？这为优化动力学提供了哪些见解？

主要发现

• 该K-FAC优化器在使用131,072的mini-batch大小时，仅用978次迭代（100个周期）就在ImageNet上实现了75.0%的top-1验证准确率，证明了在如此大规模批次下的最先进性能。
• 当mini-batch大小达到16,384时，模型仅用35个周期就收敛至75.2%的准确率，显著快于需要更多周期才能达到类似准确率的一阶方法。
• 对Batch Normalization层的FIM采用对角近似，将内存使用量从1017 MiB降低至587 MiB，且对训练准确率无明显影响。
• 陈旧FIM更新（在第13个周期后每20个周期更新一次）显著降低了计算成本，同时在131,072的mini-batch大小下仍保持75%的准确率。
• 使用1,024块Tesla V100 GPU，该方法在10分钟内实现了74.9%的top-1准确率，训练速度和可扩展性均优于现有的基于SGD的方法。
• 研究表明，即使在极大规模的mini-batch下，二阶方法也未表现出比SGD更差的泛化性能，挑战了以往的假设。