[论文解读] Sketchy Empirical Natural Gradient Methods for Deep Learning

Minghan Yang, Dong Xu+3 more|arXiv (Cornell University)|Jun 10, 2020

Sparse and Compressive Sensing Techniques参考文献 27被引用 5

一句话总结

该论文提出了一种名为Sketchy Empirical Natural Gradient（SENG）的可扩展二阶优化方法，用于深度学习。SENG利用随机化图示（randomized sketching）在低秩子空间中高效近似经验Fisher信息矩阵（EFIM）。通过将图示应用于正则化最小二乘子问题和低秩矩阵近似，SENG在神经正切核（NTK）框架下实现了全局收敛和快速线性收敛，其在ResNet50/ImageNet-1k上的训练表现优于SGD和KFAC，在41个周期内达到75.9%的Top-1准确率，每个周期耗时663秒，同时在分布式设置中保持了高可扩展性。

ABSTRACT

In this paper, we develop an efficient sketchy empirical natural gradient method (SENG) for large-scale deep learning problems. The empirical Fisher information matrix is usually low-rank since the sampling is only practical on a small amount of data at each iteration. Although the corresponding natural gradient direction lies in a small subspace, both the computational cost and memory requirement are still not tractable due to the high dimensionality. We design randomized techniques for different neural network structures to resolve these challenges. For layers with a reasonable dimension, sketching can be performed on a regularized least squares subproblem. Otherwise, since the gradient is a vectorization of the product between two matrices, we apply sketching on the low-rank approximations of these matrices to compute the most expensive parts. A distributed version of SENG is also developed for extremely large-scale applications. Global convergence to stationary points is established under some mild assumptions and a fast linear convergence is analyzed under the neural tangent kernel (NTK) case. Extensive experiments on convolutional neural networks show the competitiveness of SENG compared with the state-of-the-art methods. On the task ResNet50 with ImageNet-1k, SENG achieves 75.9\% Top-1 testing accuracy within 41 epochs. Experiments on the distributed large-batch training show that the scaling efficiency is quite reasonable.

研究动机与目标

• 解决尽管经验Fisher信息矩阵（EFIM）具有低秩特性，但自然梯度方法在大规模深度学习中计算和内存成本过高的问题。
• 开发高效的随机图示技术，以在不牺牲收敛性保证的前提下降低EFIM近似的计算负担。
• 在大批次训练设置中实现可扩展的分布式训练，具备高通信与计算效率。
• 在较弱假设和神经正切核（NTK）框架下，建立全局收敛性和快速线性收敛性。
• 在ImageNet-1k和CIFAR10等标准基准上，与当前最先进的的一阶和二阶方法相比，展示出具有竞争力的性能。

提出的方法

• 利用Sherman-Morrison-Woodbury（SMW）公式，将自然梯度方向表示为子采样梯度的线性组合，从而利用EFIM的低秩结构。
• 对于参数量适中的层，对由子采样梯度构成的正则化最小二乘子问题应用图示，以降低维度。
• 对于高维层，将梯度表示为矩阵乘积，对矩阵进行低秩近似，并利用图示加速求解所得子问题。
• 将基于SMW的方向计算与随机算法相结合，以高效处理大规模矩阵运算。
• 设计SENG的分布式版本，采用通信缩减策略，以在多GPU环境下保持高可扩展性。
• 通过理论证明确保收敛性：在弱假设下实现全局收敛，并在NTK框架下实现线性收敛。

实验结果

研究问题

• RQ1 随机图示技术是否能有效降低深度学习中经验自然梯度方法的计算和内存成本，同时保持收敛性？
• RQ2 在大规模基准测试中，所提出的SENG方法与SGD等一阶方法以及KFAC等二阶方法相比，在训练速度和准确率方面表现如何？
• RQ3 与现有方法相比，SENG在分布式大批次训练设置中的可扩展性如何？
• RQ4 SENG是否在神经正切核（NTK）框架下实现快速线性收敛？其成立的前提假设是什么？
• RQ5 在大批次设置下，SENG是否能保持高性能和低每周期耗时，避免大批次训练中常见的性能下降问题？

主要发现

• 在ResNet50/ImageNet-1k上，SENG仅用41个周期就达到了75.9%的Top-1准确率，总训练时间优于经过充分调优的SGD（76个周期）和KFAC（42个周期）。
• 在ResNet50/ImageNet-1k上，SENG的每周期耗时降低至663.17秒，显著低于KFAC的712.29秒，也远低于KFAC在更高矩阵更新频率下的1007.31秒。
• 在CIFAR10上使用VGG16_bn时，SENG单个周期耗时仅18秒，远低于KFAC的113.37秒，展现出卓越的计算效率。
• 在4至32块GPU的范围内，SENG保持了超过90%的GPU可扩展效率，表明其在分布式大批次训练中具有强大潜力。
• 在所有批量大小（包括4,096）下，SENG均仅需41个周期，表明其在大批次训练中具备出色的鲁棒性和有效性。
• 理论分析证实，在弱假设下实现全局收敛，并在NTK框架下实现线性收敛，支持其实际有效性。