LAMBオプティマイザ（Layer-wise Adaptive Moments for Batch training）は、ディープニューラルネットワークを訓練するための確率的最適化アルゴリズムであり、Adamオプティマイザを層ごとの正規化ステップで拡張したものである。2019年にGoogleの研究者（具体的にはYang You、Jing Li、Jonathan Hseuら）によって導入され、LAMBは、モデルの精度を低下させたり、徹底的なハイパーパラメータ調整を必要としたりすることなく、非常に大きなミニバッチサイズ（例：32,768以上）の効率的な使用を可能にするように設計された。これは、各層の重みノルムと更新ノルムの比率に基づいて層ごとの更新規模をスケーリングし、学習率を各層のパラメータのスケールから効果的に切り離すことで実現される。
LAMBは、Transformer (architecture)ベースのモデル、特に初期の大規模言語モデル（LLM）やResNetなどのビジョンアーキテクチャの訓練において特に影響力を持ってきた。TensorFlow（tf.keras.optimizers.LAMB経由）やPyTorch（NVIDIA、Hugging Face、FairScaleなどのライブラリのLAMB実装経由）などの分散訓練フレームワークでの採用により、訓練実行をスケールアップするための標準的なツールとなった。より大きなバッチを可能にすることで、LAMBは最先端モデルの訓練に必要なウォールクロック時間を短縮し、これは大規模な計算クラスターに依存するOpenAI、Anthropic、Google DeepMindなどの組織にとって重要である。
動機と背景
ディープニューラルネットワークを大きなミニバッチで訓練すると、1エポックあたりのパラメータ更新回数が減り、計算ステップが削減されるが、バッチサイズを単純にスケーリングすると、しばしば汎化性能の低下や不安定な収束につながる。これは「大バッチ訓練問題」として知られている。モメンタム付き確率的勾配降下法（SGD）やAdamなどの標準的なオプティマイザは、バッチサイズが増加する際に学習率の注意深い調整を必要とし、それでも精度が低下することが多い。LAMBオプティマイザは、バッチサイズの変化に対してオプティマイザをより堅牢にすることで、この問題に対処するために開発された。
LAMBの背後にある重要な洞察は、ディープネットワークの異なる層が、勾配と重みノルムの非常に異なるスケールを示すということである。例えば、初期の畳み込み層は小さな重みを持ち、後の完全接続層は大きな重みを持つ。Adamの単一のグローバル学習率は、一部の層には大きすぎる更新（発散を引き起こす）をもたらし、他の層には小さすぎる更新（収束を遅らせる）をもたらす可能性がある。LAMBは、層の重みと勾配ノルムの比率に基づいて更新を正規化する層ごとの適応率を導入し、各層がその規模に応じて安定したペースで動くことを保証する。
アルゴリズムの説明
LAMBは、追加の正規化ステップを備えたAdamの変種と見なすことができる。\(\theta_t\)を反復\(t\)でのパラメータ、\(g_t\)を\(\theta_t\)に関する損失の勾配とする。LAMBは、Adamと同様に、勾配の一次モーメント（\(m_t\)）と二次モーメント（\(v_t\)）を、指数減衰率\(\beta_1\)と\(\beta_2\)（通常は0.9と0.999）で維持する。バイアス補正の後、Adam更新\(\frac{m_t}{\sqrt{v_t} + \epsilon}\)を計算する。
重要な違いはトラスト比 \(\phi\)である：各層\(i\)について、\(\phi_i = \frac{||\theta_{t,i}||}{||r_{t,i}||}\)、ここで\(r_{t,i} = \frac{m_{t,i}}{\sqrt{v_{t,i}} + \epsilon}\)はその層のAdam更新（学習率なし）であり、\(||\cdot||\)はL2ノルムを示す。層\(i\)の最終更新は\(\theta_{t+1,i} = \theta_{t,i} - \eta \cdot \phi_i \cdot r_{t,i}\)であり、ここで\(\eta\)はグローバル学習率である。このトラスト比は、更新を層の重みノルムに比例してスケーリングするため、小さな層は小さな（しかし無視できない）更新を受け取り、大きな層はより大きな、しかし安定した更新を受け取る。
実際には、ゼロ除算を避けるために、分母に小さな定数（例：1e-6）が追加される。アルゴリズムには、AdamWで使用される分離された重み減衰アプローチに従って、更新に融合されたオプションの重み減衰（L2正則化）も含まれる。トラスト比がすべての層で1に設定されている場合、LAMBは標準のAdam（バイアス補正付き）に還元される。
ハイパーパラメータと調整
LAMBは、Adamからほとんどのハイパーパラメータを継承している：\(\beta_1\)（モメンタム）、\(\beta_2\)（分散減衰）、\(\epsilon\)（数値安定性）、および重み減衰率。主要な新しいハイパーパラメータはグローバル学習率 \(\eta\)であり、大バッチ訓練では0.01〜0.1の範囲で設定されることが多く、Adamで一般的な値（例：1e-3）よりも大幅に高い。著者らは、非常に大きなバッチ（例：BERTで32,768）の場合、最初の10%のステップで線形ウォームアップを伴う0.01の学習率がうまく機能し、残りのステップではコサイン減衰学習率スケジュールを使用することを推奨した。
さらに、\(\beta_2\)の選択は安定性に影響を与える可能性があり、スパースな勾配を持つモデルでは、より高い\(\beta_2\)（例：0.99）が使用されることがある。著者らはまた、バッチサイズを学習率に比例してスケーリングできる（線形スケーリング則）ことを示唆しており、これはLAMBでうまく機能するガイドラインである。例えば、バッチサイズが2倍になると、精度を失うことなく学習率も2倍にできる。
性能とベンチマーク
元の論文では、LAMBは2つの主要なタスクで評価された：ImageNetでのResNet-50の訓練（画像分類）と、マスク言語モデリングのためのBERT（transformerベースの言語モデル）の訓練である。LAMBを使用して、著者らはバッチサイズ32,768でわずか2,048回の反復でImageNetのtop-1精度76.0%を達成し、はるかに少ないエポックで小さなバッチ（例：256）で達成された最先端の精度に匹敵した。BERTでは、1,024個のTPUを使用して約3.5分で同じ精度（例：SQUADでのF1スコア1.0）にモデルを訓練し、以前の方法よりも10倍の高速化を実現した。
その後、LAMBはGoogleの内部ワークフローでBERTベースのモデルを訓練するためのデフォルトのオプティマイザとなった。論文では、LAMBがバッチサイズを1,024から65,536にスケーリングする際に、Adamとモメンタム付きSGDの両方を上回る性能を発揮したと報告されている。著者らはまた、LAMBがGradient Clipping（勾配爆発を防ぐために使用される）とうまく機能し、NVIDIA GPUやGoogle TPUなどの最新ハードウェアで使用される混合精度訓練と互換性があることを示した。
大規模訓練への応用
LAMBの主な応用は、バッチサイズが単一デバイスのメモリに収まらない分散訓練である。このような設定では、勾配はデータ並列性を使用して複数のGPUまたはTPUにわたって平均化される。例えば、OpenAIとGoogle DeepMindは、シーケンス長が数千の大きなTransformer (architecture)モデルを訓練する際に、LAMBに類似したオプティマイザを使用している。LAMB（およびその後継のLAMB2）のような新しいオプティマイザが提案されているが、LAMBは、ビジョントランスフォーマーやLLMの訓練を含む多くのオープンソースの取り組みで、AI21 LabsやSambaNovaなどの研究グループや企業によって信頼性の高い選択肢であり続けている。
機械学習の文脈では、Amazon Web Services（AWS Trainiumハードウェアを使用）で、LAMBは効率性のためにカスタムカーネルでサポートされている。同様に、IntelとAMDは自社のアクセラレータでLAMBをベンチマークしている。このオプティマイザは極端なバッチサイズを処理できるため、単一エポックのコストが高く、エポック数を減らすことが最重要である大規模データセットでのモデルの事前訓練に価値がある。
他のオプティマイザとの関係
LAMBは、SGD変種、Adam、およびAdamW（分離された重み減衰）やLARS（Layer-wise Adaptive Rate Scaling）などの後継を含む適応オプティマイザのファミリーの一部である。2017年にYouらによってCNNの大バッチ訓練用に導入されたLARSは、同様の層ごとのトラスト比を使用するが、二次モーメントを維持せず、一次モーメント（モメンタム）と勾配ノルムに依存する。LAMBは、LARSの利点（層ごとのスケーリング）とAdamの適応的なパラメータごとの学習率を組み合わせており、スパースな勾配を持つモデル（トランスフォーマーなど）に対してより堅牢である。
もう1つの密接に関連するオプティマイザは、分散削減を組み込んだNVLAMB（Nvidia製）である。ただし、LAMBはよりシンプルで広く使用されている。Sequence-to-Sequence (Seq2Seq)タスクとBeam Searchでは、LAMBは推論に直接影響しないが、訓練の収束を助け、間接的にシーケンス復号を改善する。
拡張と変種
導入以来、いくつかの変種が提案されている。LAMB2（これもGoogle製）は、勾配分散に基づく正規化係数を追加し、特定の問題での安定性を向上させる。バイアス補正付きLARSやその他の修正も一般的である。実際には、多くのフレームワークが、最初の数ステップで有益なモーメントのオプションのバイアス補正付きLAMBを実装している。PyTorchのtorch.optim.Lamb（torch_optimizerパッケージ内）などの一部の実装では、トラスト比パラメータの調整やカスタムの層ごとの学習率の使用が可能である。
新しいオプティマイザ（例えば、異なるスケーリング戦略を持つAdamW、1サイクルスケジュールなど）の出現にもかかわらず、LAMBは大バッチ訓練の強力なベースラインであり続けている。研究コミュニティは、汎化性能をさらに向上させるために、LAMBをデータ拡張やGradient Clippingと組み合わせることを探求してきた。
実用的な考慮事項と制限
LAMBは大バッチ設定で優れているが、小さなバッチサイズ（例：1,024未満）では常に最良の選択とは限らない。そのような領域では、標準のAdamやモメンタム付きSGDがよりシンプルで同等に効果的かもしれない。LAMBは層ごとのノルムを計算する計算オーバーヘッドを追加するが、これは最新のハードウェアでは無視できるが、多くの小さな層を持つモデル（例：U-Netアーキテクチャ）では無視できない場合がある。
もう1つの制限は、LAMBのトラスト比が、一部の層が非常に小さな重みノルム（ゼロに近い）を持つ場合に、不安定な訓練につながることがあることである。これは通常、分母にイプシロン項を追加し、重みがゼロにドリフトするのを防ぐ重み減衰を使用することで軽減される。さらに、LAMBは学習率とウォームアップステップの注意深い調整を必要とし、不適切なスケジュールは発散につながる可能性がある。
メモリ使用量はAdamと同様（パラメータごとに2つのモーメントベクトル）であり、よりメモリを消費するわけではない。非常に大きなモデルでは、モデルプルーニングやGradient ClippingがLAMBと併用されることがあるが、これらは直交する技術である。
結論
LAMBは、大規模深層学習のための最適化アルゴリズムのツールボックスにおいて基礎となっている。巨大なミニバッチでの効果的な訓練を可能にすることで、多くのベンチマークモデルの開発を加速し、実験コストを削減した。その層ごとの適応原理は、その後のオプティマイザ設計に影響を与え、分散訓練の課題に直面する実践者にとって、実用的でよく理解されたソリューションであり続けている。人工知能が成長し続けるにつれて、LAMBのようなオプティマイザは進化する可能性が高いが、層ごとのトラストと適応モーメントという中核的なアイデアは今後も続くであろう。