[論文レビュー] Adafactor: Adaptive Learning Rates with Sublinear Memory Cost

Noam Shazeer, Mitchell Stern|arXiv (Cornell University)|Apr 11, 2018

Stochastic Gradient Optimization Techniques参考文献 10被引用数 163

ひとこと要約

Adafactor は、行列パラメータの二次モーメント推定を行ごとおよび列ごとの和に分解することで、メモリ効率の高い適応最適化アルゴリズムを導入し、Transformer の学習において Adam と競争力のある性能を実現しつつ、サブリニアなメモリ使用量を可能にする。また、トレーニングを安定化させるためのアップデートクリッピングとデカイ・レートスケジュールを提案し、スケール認識的な更新のための相対的ステップサイズ（Adafactor）を導入する。

ABSTRACT

In several recently proposed stochastic optimization methods (e.g. RMSProp, Adam, Adadelta), parameter updates are scaled by the inverse square roots of exponential moving averages of squared past gradients. Maintaining these per-parameter second-moment estimators requires memory equal to the number of parameters. For the case of neural network weight matrices, we propose maintaining only the per-row and per-column sums of these moving averages, and estimating the per-parameter second moments based on these sums. We demonstrate empirically that this method produces similar results to the baseline. Secondly, we show that adaptive methods can produce larger-than-desired updates when the decay rate of the second moment accumulator is too slow. We propose update clipping and a gradually increasing decay rate scheme as remedies. Combining these methods and dropping momentum, we achieve comparable results to the published Adam regime in training the Transformer model on the WMT 2014 English-German machine translation task, while using very little auxiliary storage in the optimizer. Finally, we propose scaling the parameter updates based on the scale of the parameters themselves.

研究の動機と目的

• モデルサイズが大きくなるにつれて適応的勾配法のメモリ制約を動機づける。
• 行列パラメータについて per-parameter メモリを O(nm) から O(n+m) に削減する factorized 二次モーメント推定を提案する。
• 二次モーメントの減衰が遅すぎると不安定性が生じる問題と、その対処法としてアップデートクリッピングと適応減衰を提供する。
• Adafactor が Transformer の学習で Adam と同等の結果を、はるかに少ないメモリで達成することを示す。
• パラメータの大きさに合わせて更新をスケールさせる相対ステップサイズで最適化を拡張する。

提案手法

• 行列値パラメータに対して R ∈ R^{n×k} および S ∈ R^{k×m}（k ≪ n,m）による V ≈ RS の factored 二次モーメント表現を導入する。
• k=1 の核となる場合の解析解を導出し、行/列和の指数平滑化（V1_m1_n^T V / 1_n^T V 1_m）との互換性を確保する。
• per-row および per-column の累積を用いて事実上の second moments を実装し、R_t および C_t を用いた分解的な V_t を normalize して ¼actored¼ullV_t = (R_t C_t)/(1_n^T R_t) を形成する。
• 大きな未スケーリング更新を RMS(U_t) が閾値 d を超えるときにキャップするアップデートクリッピングを提案する。
• 二次モーメントの減衰スケジュールを増やす（最初の ast, Reddi らの継承）こと、安定化のための別のスケジュールを提案する。
• Adafactor を相対ステップオプティマイザとして定義し、実際のステップ alpha_t をパラメータスケールの RMS と相対ステップ rho_t から計算し、個々のパラメータ更新 U_t = G_t / sqrt(V_hat_t) を行い、必要に応じてクリッピングを適用する。

実験結果

リサーチクエスチョン

• RQ1 factored（行ごと/列ごと）二次モーメント推定が全二次モーメントアキュムレータと同等の性能を再現できるか。
• RQ2 factorization によるメモリ削減が Transformer 学習のような大規模タスクで収束とモデル品質に影響を与えるか。
• RQ3 運動量を分離し適応学習率を使用する際に生じる安定性問題は何か、アップデートクリッピングと減衰率スケジュールはそれをどう緩和できるか。
• RQ4 パラメータの大きさに応じてスケールする相対的ステップサイズは、異なる初期化を持つパラメータでのロバスト性を改善するか。
• RQ5 提案された救済策（アップデートクリッピング、減衰の増加、相対的ステップサイズ）は、現代の neural MT タスクで実務上どのように相互作用するか。

主な発見

• factored 二次モーメント推定は、行列パラメータに対するメモリを O(nm) から O(n+m) に削減し、Transformer における BLEU スコアが full-accumulator Adam に匹敵する。
• モメンタムを取り除くとトレーニングが不安定になるが、アップデートクリッピングと適切な減衰スケジュールが安定性を回復させる。
• 閾値 d を用いたアップデートクリッピングはウォームアップなしの設定で安定性を向上させた； d=1 のとき不安定性を大きく緩和し、d=2 は改善を示さなかった。
• 相対ステップサイズは Adafactor と組み合わせると競争力のある性能を維持し、埋め込みパラメータのスケールが異なる場合でもロバストだった。
• 二次モーメントの増加減衰スケジュール（例: 1 - t^{-c}）は、特に c=0.5（および関連するバリエーション）でクリッピングと組み合わせた場合に安定かつ収束する結果を得た。
• rank-1 または factorized 表現と提案されたスケーリングを使用することで、Transformer モデルを学習しつつ副次的なストレージをサブリニアに抑えつつ、Adam ベースの基準とほぼ同等の BLEU スコアを達成できた。