[論文レビュー] GradInit: Learning to Initialize Neural Networks for Stable and Efficient Training

Chen Zhu, Renkun Ni+4 more|arXiv (Cornell University)|Feb 16, 2021

Advanced Neural Network Applications被引用数 15

ひとこと要約

GradInit は、SGD や Adam の1ステップ更新後の損失を最小化することで、ニューラルネットワークの各層に対して最適な初期化スケーリング係数を学習するアーキテクチャに依存しない手法である。収束が早くなり、テスト精度が向上し、Transformer の学習率ウォームアップなしで学習可能となり、ImageNet や IWSLT 習得ベンチマークで最先端の結果を達成した。

ABSTRACT

Innovations in neural architectures have fostered significant breakthroughs in language modeling and computer vision. Unfortunately, novel architectures often result in challenging hyper-parameter choices and training instability if the network parameters are not properly initialized. A number of architecture-specific initialization schemes have been proposed, but these schemes are not always portable to new architectures. This paper presents GradInit, an automated and architecture agnostic method for initializing neural networks. GradInit is based on a simple heuristic; the norm of each network layer is adjusted so that a single step of SGD or Adam with prescribed hyperparameters results in the smallest possible loss value. This adjustment is done by introducing a scalar multiplier variable in front of each parameter block, and then optimizing these variables using a simple numerical scheme. GradInit accelerates the convergence and test performance of many convolutional architectures, both with or without skip connections, and even without normalization layers. It also improves the stability of the original Transformer architecture for machine translation, enabling training it without learning rate warmup using either Adam or SGD under a wide range of learning rates and momentum coefficients. Code is available at https://github.com/zhuchen03/gradinit.

研究の動機と目的

• 特に複雑または新しいアーキテクチャにおいて、初期化が不適切なために生じる学習不安定性の課題に対処すること。
• 新しいまたは異種のアーキテクチャに移植できないアーキテクチャ固有の初期化ルールの限界を克服すること。
• アーキテクチャの変更を要せず、訓練の安定性と収束速度を向上させる汎用的かつ最適化器に依存する初期化手法を開発すること。
• SGD や Adam を用いて、1202層の ResNet やポスト-LN Transformer を学習率ウォームアップなしで学習可能にすること。
• 合計学習時間の1%未満で最適なスケーリング係数を学習できるスケーラブルで効率的な初期化手法を提供すること。

提案手法

• 畳み込みカーネルや線形層などのパラメータブロックごとにスカラーマルティプライヤーを導入し、ランダムに初期化された重みを再スケーリングする。
• 指定された最適化器（SGD や Adam）の1ステップ更新後の損失を最小化するために、これらのスカラーマルティプライヤーを最適化する数値スキームを採用する。
• 最適化ステップを微分可能プロセスとして扱い、最初の更新を逆伝播することで最適なスケーリング係数を学習する。
• ミニバッチサンプリングの確率性、学習率、および最適化器のダイナミクス（例：モーメンタムや適応的学習率）を考慮する。
• 学習済みのスケーリング係数を用いて、標準的な学習を開始する前に元のランダム重みを再スケーリングする。
• Adam と SGD の両方の最適化器をサポートし、それぞれに特化した異なるスケーリングパターンを学習する。これは、勾配分散低減戦略の差を反映している。

実験結果

リサーチクエスチョン

• RQ1 1つの自動化された手法が、多様なアーキテクチャにわたって学習安定性と収束速度を向上させる初期化スケーリング係数を学習できるか？
• RQ2 GradInit は、Adam や SGD で学習する Transformer において、学習率ウォームアップの必要性をどの程度排除できるか？
• RQ3 GradInit は、Xavier や He、Admin などの従来の初期化方式と比較して、最終テスト精度と学習安定性の面でどのように優れているか？
• RQ4 GradInit は、標準的な初期化では失敗するが、1202層の極めて深いネットワーク（例：ResNet）を安定して学習可能にするか？
• RQ5 学習済みのスケーリング係数に現れるパターンは、勾配分散と初期化に起因する不安定性に関する何を明らかにするか？

主な発見

• GradInit を用いることで、IWSLT-14 DE-EN で Post-LN Transformer を学習率ウォームアップなしに学習可能となり、極端なハイパーパramータ設定下で平均 BLEU スコア 36.0 を達成した。
• GradInit を用いることで、1202層の ResNet は ResNet-110 よりも顕著に高いテスト精度を達成し、標準的な初期化では失敗する極めて深いネットワークの学習に成功した。
• ImageNet において、Batch Normalization がなくても、GradInit は ResNet-50 の最終テスト精度を向上させ、より良い初期化を学習した。
• GradInit は、特にエンコーダー部と残差ブランチ部において、初期の勾配分散を低減し、スキップ接続のバランスを取るために出力投影層やFFN層の重みを低減する。
• GradInit は、Adam と SGD それぞれに特化した異なるスケーリングパターンを学習しており、最適化器に依存する初期化が勾配分散制御において重要であることが示された。
• 可視化結果から、GradInit は初期化段階でネットワークを滑らかな損失領域に配置することで、サンプル間の勾配分散を低減し、収束を加速していることがわかった。