LAMB（レイヤー単位適応モーメントによるバッチ学習）は、深層ニューラルネットワークの訓練、特に大規模バッチ分散コンピューティング環境向けに設計された最適化アルゴリズムである。2019年にGoogleとトロント大学の研究者らによって導入され、数千のアクセラレータへの訓練スケーリングという課題を、モデルの精度と収束速度を維持しながら解決することを目指した。
このアルゴリズムは、Adamオプティマイザを拡張し、パラメータごとではなくレイヤーごとに適応学習率を計算する。このレイヤー単位の適応により、LAMBは異なるネットワーク層間での勾配スケールのばらつきをより効果的に処理でき、これはTransformerや残差ネットワークのような深層アーキテクチャで特に重要である。レイヤーの重みと勾配のノルムに基づいて更新を正規化することで、LAMBは非常に大きなバッチサイズでも安定かつ効率的な訓練を保証する。
背景と動機
大規模なニューラルネットワークの訓練には通常、多くのGPUやTPUに分散された膨大な計算リソースが必要である。バッチサイズを増やすことはこれらのリソースを効率的に利用する一般的な戦略だが、しばしばモデルの性能低下や収束の遅延を引き起こす。確率的勾配降下法（SGD）やAdamのような従来のオプティマイザは、レイヤー間の勾配スケールの不均一性を考慮しないグローバル学習率に依存するため、大規模バッチでは苦戦する。
LAMBオプティマイザはこれらの制限を克服するために開発された。その設計は、畳み込みネットワークの大規模バッチ訓練で以前使用されていたレイヤー単位適応レートスケーリング（LARS）アルゴリズムに触発されている。LAMBはこの概念を適応モーメント推定と組み合わせて一般化し、両アプローチの利点を統合している。
アルゴリズムの詳細
LAMBは、レイヤーの重みノルムと勾配ノルムの比率に基づいて各レイヤーの更新を計算する。ステップtにおけるパラメータテンソルの核心的な更新規則は以下の通りである：
- Adamと同様に、一次および二次モーメント推定（勾配の平均と分散）を計算する。
- 更新方向を、モーメント補正された勾配を二次モーメントの平方根に小さなイプシロンを加えたもので割って計算する。
- この方向を、レイヤーの重みノルムと更新方向ノルムの比率でスケーリングする。
- グローバル学習率を掛けて更新を適用する。
このレイヤー単位のスケーリングにより、重みノルムが大きいレイヤーは比例して大きな更新を受け、ノルムが小さいレイヤーは保守的に更新される。アルゴリズムには、勾配クリッピング技術と同様に、極端な更新を防ぐためにクリップ可能な信頼比率も組み込まれている。
著者らは、LAMBがImageNet上のResNet-50をバッチサイズ32,768で訓練し、バッチサイズ256のベースラインと同じ精度を達成しつつ、大幅に少ないステップで実現できることを実証した。これにより、数百から数千のアクセラレータにわたる分散訓練に非常に適している。
応用と影響
LAMBは、大規模言語モデルや他の深層学習モデルの訓練で広く採用されている。例えば、BERTや他のTransformerベースのモデルの大規模訓練に使用され、訓練時間を数日から数時間に短縮した。このアルゴリズムは、AWSやMicrosoft Azureのようなクラウドプラットフォーム、およびAWS TrainiumやGraphcore IPUのような専用AIハードウェアなど、ハードウェアリソースが豊富な環境で特に価値がある。
多くの後続のオプティマイザ、例えばLAMBの変種や後継は、その原理に基づいて構築されている。また、学習率スケジュールや適応最適化手法の研究にも影響を与えた。TensorFlowやPyTorchなどのフレームワークでのオープンソース実装により、広範な機械学習コミュニティが利用できるようになった。
他のオプティマイザとの比較
Adamと比較して、LAMBは大規模バッチサイズを使用する際に、通常より速い収束とより良い最終性能を達成する。Adamのグローバル学習率はしばしば慎重な調整を必要とし、大規模バッチでは不安定性を引き起こす可能性がある。LAMBのレイヤー単位の適応はこれらの問題を軽減し、より積極的なスケーリングを可能にする。
LARSと比較すると、LARSはモーメンタム付きSGD用に設計されているのに対し、LAMBは適応モーメント推定を組み込んでおり、ノイズの多い勾配やスパースな特徴に対してより堅牢である。これにより、LAMBはシーケンス間モデルやエンコーダ・デコーダフレームワークを含む幅広いアーキテクチャに対してより汎用的な選択肢となる。
制限と考慮事項
その利点にもかかわらず、LAMBには制限がないわけではない。このアルゴリズムは、信頼比率クリップ閾値やイプシロン項などの追加のハイパーパラメータを導入し、特定のタスクに応じて調整が必要になる場合がある。また、レイヤー単位のスケーリングが有益であるという前提に依存しており、高度に相関したレイヤーや特定の重み初期化スキームを使用するアーキテクチャでは常に当てはまるとは限らない。
さらに、LAMBは大規模バッチ設定で優れている一方、Adamのような単純なオプティマイザで十分な小規模バッチではその利点が減少する。研究者らはまた、アルゴリズムの性能がグローバル学習率の選択に敏感であり、最適な結果を得るには学習率ウォームアップが必要な場合があることも指摘している。
結論
LAMBは、大規模な深層学習訓練の最適化における重要な進歩を表している。レイヤー単位の適応と適応モーメントを組み合わせることで、巨大なバッチサイズでの効率的かつ安定した訓練を可能にし、生成AIと人工知能研究の時代における基盤技術となっている。その影響は元の応用を超えて広がり、後続のオプティマイザや訓練方法論の開発を形作っている。