[論文レビュー] 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

ひとこと要約

本稿では、スケッチ化された低ランク部分空間における経験的フィッシャー情報行列（EFIM）を効率的に近似するために、確率的スケッチを用いるスケーラブルな2次最適化手法であるSketchy Empirical Natural Gradient（SENG）を提案する。正則化最小二乗問題と低ランク行列近似にスケッチを適用することにより、ニューラルトランジェントカーネル（NTK）の枠組み下で、グローバル収束と高速な線形収束を達成し、ResNet50/ImageNet-1kで75.9%のTop-1精度を41エポックで達成、1エポックあたり663秒で、SGD や KFAC より優れた性能を示した。分散環境でも高いスケーリング効率を維持した。

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.

研究の動機と目的

• 経験的フィッシャー情報行列（EFIM）の低ランク性にもかかわらず、自然勾配法の高い計算コストとメモリ消費を解決すること。
• EFIM近似の計算負荷を犠牲にすることなく収束保証を維持する、効率的な確率的スケッチ技術の開発。
• 大バッチ学習向けに、高い通信および計算効率を実現するスケーラブルな分散学習の実現。
• やや弱い仮定とNTKの枠組み下で、グローバル収束と高速な線形収束を確立すること。
• ImageNet-1k や CIFAR10 といった標準ベンチマークで、最先端の1次および2次最適化手法と比較して競争力のある性能を示すこと。

提案手法

• 自然勾配方向を、EFIMの低ランク構造を活用して、サブサンプルされた勾配の線形結合として表現するため、Sherman-Morrison-Woodbury（SMW）の公式を用いる。
• パラメータ数が中程度の層に対しては、サブサンプルされた勾配から形成される正則化最小二乗問題にスケッチを適用し、次元削減を実現する。
• 高次元の層に対しては、勾配を行列積として表現し、その行列に低ランク近似を適用した後、スケッチを用いて得られるサブプロブレムの解法を高速化する。
• 大規模行列演算を効率的に行うために、SMWに基づく方向計算と確率的アルゴリズムを統合する。
• 複数のGPU間で高いスケーリング効率を維持するための通信削減戦略を用いた、SENGの分散版を設計する。
• やや弱い仮定の下でグローバル収束を保証する理論的分析を行い、NTKの枠組み下で線形収束を示す。

実験結果

リサーチクエスチョン

• RQ1 確率的スケッチ技術は、深層学習における経験的自然勾配法の計算コストとメモリ消費を効果的に低減できるか、収束性は保持されるか？
• RQ2 大規模ベンチマークにおいて、SENGはSGD や KFAC といった1次および2次最適化手法と比較して、学習速度と精度で優れているか？
• RQ3 既存手法と比較して、SENG は分散的大バッチ学習環境でのスケーリング効率はいかがなっているか？
• RQ4 SENG はニューラルトランジェントカーネル（NTK）の枠組み下で高速な線形収束を達成するか？また、その仮定は何か？
• RQ5 大バッチサイズであっても、SENG は高い性能と低エポック時間（1エポックあたりの時間）を維持できるか？一般的に大バッチ学習で見られる性能劣化を回避できるか？

主な発見

• SENG は、ResNet50/ImageNet-1k で41エポックで75.9%のTop-1精度を達成し、最適化されたSGD（76エポック）や KFAC（42エポック）よりも総合的な学習時間で優れた性能を示した。
• ResNet50/ImageNet-1k において、SENG の1エポックあたりの所要時間は663.17秒にまで低減され、KFACの712.29秒を下回り、さらに高い行列更新頻度のKFAC（1007.31秒）よりも顕著に短い。
• VGG16_bn を用いたCIFAR10では、SENG は1エポックあたり18秒で完了したのに対し、KFAC は113.37秒を要し、計算効率に顕著な優位性を示した。
• 4〜32のGPUにわたり、SENG は90%を超えるGPUスケーリング効率を維持しており、分散的大バッチ学習への強い可能性を示している。
• 4,096を含む全バッチサイズにおいて、SENG は41エポックという同じエポック数を達成しており、大バッチ学習におけるロバスト性と有効性が裏付けられた。
• 理論的分析により、やや弱い仮定の下でグローバル収束が保証され、NTKの枠組み下で線形収束が確認された。これにより、実用的有効性が裏付けられた。