Adam最適化の根本的な課題
大規模言語モデル（LLM）の訓練は、膨大な計算資源と時間を要します。その効率を左右する中心的な要素が、勾配降下法に基づく最適化手法です。現代のLLM訓練では、Adamが事実上の標準として広く採用されています。しかし、Adamはその優れた性能にもかかわらず、その理論的な基盤には長らく脆弱性が指摘されてきました。特に、特定の条件下での発散挙動は、研究者や開発者の間で大きな懸念事項でした。本研究は、このAdamが抱える根本的な課題に対し、原理的な解明と具体的な解決策を提供することを目指しています。
Adamの収束と発散の境界
本研究の最初の重要な成果は、Adamの収束と発散に関する長年の議論に終止符を打つものです。私たちは、Adamの挙動が問題の種類に依存する「相転移」を示すことを明らかにしました。具体的には、処理単位のデータ量に応じて適切に調整された設定値（ハイパーパラメータ）を用いることで、Adamは安定して収束します。これは、Adamの理論的な安定性を保証する上で極めて重要な発見です。一方で、特定の条件下、特に一部の設定値が小さい場合にAdamが発散する可能性も厳密に示されました。この知見は、Adamを実践的に利用する上で、設定値の選択がいかに重要であるかを明確に示唆しています。
TransformerにおけるAdamの優位性
次に、本研究はAdamがなぜSGD（確率的勾配降下法）と比較して、Transformerモデルの訓練において顕著に優れた性能を発揮するのかを深く掘り下げました。その鍵となるのは、訓練中にモデルの損失関数の形状を示す「ヘッセ行列」が特殊な構造へと進化することです。私たちは、訓練が進むにつれてヘッセ行列がほぼブロック対角形へと変化し、さらに各ブロック間で勾配の性質が大きく異なる「強いブロック異質性」を持つことを発見しました。この特異な構造こそが、Adamが持つ「対角前処理器」の機能を非常に効果的にする要因であることを数学的に証明しました。
特殊なヘッセ行列構造の起源
この特殊なヘッセ行列構造がなぜTransformerモデルで現れるのかについても、本研究は踏み込んでいます。私たちは、この構造が大規模な行列変数間の連続的な乗算、すなわちTransformerの内部計算プロセスに由来することを突き止めました。さらに、ランダム行列理論に基づいた厳密な分析を通じて、この現象の背後にある数学的な原理を解明しました。この知見は、Transformerモデルの内部動作、ひいては大規模言語モデルの学習特性そのものに対する深い理解をもたらします。これは、単に最適化手法を改善するだけでなく、モデルの設計思想にも影響を与える可能性を秘めています。
新しい最適化手法「Adam-mini」の提案と波及効果
これらの包括的な洞察に基づいて、私たちは「Adam-mini」という新しい最適化手法を提案しました。Adam-miniは、Adamの持つ優れた訓練性能を完全に維持しつつ、そのメモリ使用量を驚くべきことに50%削減することに成功しました。これは、特にGPUメモリが限られる環境や、さらに大規模なモデルの訓練において、極めて実用的なメリットをもたらします。本研究の成果は、Adamの改善に留まらず、行列ベースの非凸最適化問題における新しい局所構造を明らかにするものです。また、Muonのような最近開発された他のニューラルネットワーク最適化手法の理解を深め、そのさらなる改善にも貢献すると期待されます。最適化技術の設計と応用において、より原理に基づいたアプローチを促す画期的な研究と言えるでしょう。
PR
文賢 →
最適化手法は、モデルの性能を左右する最重要要素です。Adamの不安定性を原理から解明し、メモリを半減させる新手法を出すのは素晴らしい。私も新しいモデルを訓練する際は、まずAdam-miniを試します。一次情報で性能を確認します。
