[論文レビュー] AdaX: Adaptive Gradient Descent with Exponential Long Term Memory

Wenjie Li, Zhaoyang Zhang+2 more|arXiv (Cornell University)|Apr 21, 2020

Advanced Vision and Imaging参考文献 27被引用数 24

ひとこと要約

AdaXは、過去の勾配の指数移動平均を指数的長期記憶メカニズムに置き換えることで、訓練の安定性と収束性を向上させる、新しい自己適応最適化アルゴリズムを提案する。凸および非凸設定の両方で収束が理論的に証明されており、視覚および自然言語処理のベンチマークにおいて、Adamを上回り、SGD with momentumと同等の性能を発揮する。特に汎化性能およびハイパーパrameter選択に対するロバスト性が優れている。

ABSTRACT

Although adaptive optimization algorithms such as Adam show fast convergence in many machine learning tasks, this paper identifies a problem of Adam by analyzing its performance in a simple non-convex synthetic problem, showing that Adam's fast convergence would possibly lead the algorithm to local minimums. To address this problem, we improve Adam by proposing a novel adaptive gradient descent algorithm named AdaX. Unlike Adam that ignores the past gradients, AdaX exponentially accumulates the long-term gradient information in the past during training, to adaptively tune the learning rate. We thoroughly prove the convergence of AdaX in both the convex and non-convex settings. Extensive experiments show that AdaX outperforms Adam in various tasks of computer vision and natural language processing and can catch up with Stochastic Gradient Descent.

研究の動機と目的

• 非凸最適化におけるAdamの不安定性および非収束問題に対処すること、特に局所最適解に収束する傾向があること。
• 2次モーメント計算におけるAdamの指数移動平均の限界を克服すること、これにより学習率が不安定になり、一般化性能が低下する。
• 高速な収束を維持しながら一般化性能を向上させる理論的に整合性のある自己適応最適化手法の設計。
• 多様なディープラーニングタスクにおいて、AdaXがAdam、AdamW、AMSGrad、SGD with momentumを上回ることを実証的に検証すること。
• AdaXがハイパーパrameterチューニング、特に$eta_2$および学習率の選択に対してロバストであることを示すこと。

提案手法

• AdaXは、2次モーメントの計算において、Adamの過去の勾配の指数移動平均を、指数的長期記憶メカニズムに置き換える。
• AdaXにおける2次モーメントは、すべての過去の2乗勾配の指数重み付き和として計算され、歴史的な勾配行動の持続的記憶が保証される。
• 学習率は、長期記憶行列の逆平方根を用いて自己適応的にスケーリングされ、Adamと同様の仕組みだが、より高い安定性を有する。
• AdaXは、1次モーメンタムを維持しつつ、より安定した累積的勾配記憶に置き換えた修正された更新ルールを採用する。
• アルゴリズムは、凸および非凸設定の両方で収束が理論的に証明されており、AMSGradと同等の収束速度を示す。
• AdaXは$eta_2 = 10^{-4}$をデフォルトとして実装されており、これがロバストで計算的に効率的であることが示されている。

実験結果

リサーチクエスチョン

• RQ1 非凸問題におけるAdamの高速収束が、2次モーメント推定の不安定性により、最適でない解に収束する原因となっているか？
• RQ2 Adamの指数移動平均を指数的長期記憶メカニズムに置き換えることで、非収束問題が解消され、一般化性能が向上するか？
• RQ3 視覚および自然言語処理タスクにおいて、AdaXはAdam、AdamW、AMSGrad、SGD with momentumと比較して、収束速度および最終性能でどのように差をつけるか？
• RQ4 実用的な訓練シナリオにおいて、AdaXは$eta_2$および初期学習率といったハイパーパrameterの選択に対してロバストか？
• RQ5 AdaXは、自己適応手法の高速収束を維持しながら、SGD with momentumと同等の性能を達成できるか？

主な発見

• CIFAR-10では、AdaXがAdamWを上回り、テストtop-1正解率94.5%を達成したのに対し、AdamWは92.1%であった。
• ResNet-18を用いたImageNetでは、AdaX-Wがtop-1正解率75.58%を達成し、AdamW（68.27%）およびAMSGrad(W)（68.76%）を大きく上回った。
• ResNet-50を用いたImageNetでは、AdaX-Wが75.58%のtop-1正解率を達成し、SGDMの77.12%に近づきながらも、より高速に収束し、訓練正解率も高かった。
• AdaXは、学習率および$eta_2$ハイパーパrameterに対してロバストであり、さまざまな設定において性能低下が最小限に抑えられ、特に$eta_2 = 10^{-4}$のとき顕著であった。
• One Billion Wordデータセットにおける言語モデリングでは、AdaXはSGD with momentumと同等の性能を示し、優れた一般化性能と高速収束を示した。
• 実行時間ベンチマークでは、AdaX-WはAdamWとほぼ同等の効率を示し、ImageNetでは50.74時間（AdamWは50.10時間）と、わずか1.8%の増加に抑えられ、計算オーバーヘッドは最小限であった。