はじめに
深層学習の発展を支えてきた最適化手法の中でも「Adam」やその改良版「AdamW」は特に重要な役割を果たしてきました。
これらの手法は、変動の大きい勾配や疎な勾配の問題に効果的に対応し、特に Transformer のような大規模モデルのトレーニングにおいて不可欠な存在となってきました。
しかし、これらの手法にはメモリ消費が多く、計算コストが高いという課題があり、モデルが大規模化するほどその問題は顕著になります。
この状況を打開するために「SGD-SaI (Stochastic Gradient Descent with Scaling at Initialization)」という新しい最適化手法が考案されました。
この手法は、従来の適応型最適化手法に依存せず、初期の学習率スケーリングだけで十分な性能を発揮するという画期的な考え方に基づいています。
その結果、メモリ使用量を大幅に削減し、トレーニング速度も著しく向上させることに成功しています。
本記事では、SGD-SaI の背景、原理、効果、そして将来性について解説していきます。
なぜ Adam を見直す必要があるのか？
Transformer や GPT のような大規模モデルでは、学習中に各パラメータに異なる勾配が適用されるため、適応型の学習率調整が重要とされてきました。
これを実現するのが、Adamが持つ「二次モーメント項（v_t）」です。
このアプローチには主に二つの課題があります。
第一に、メモリの消費が非常に大きいことです。
Adam は各パラメータについて一次モーメント（m_t）と二次モーメント（v_t）の両方を記憶する必要があり、パラメータの3倍のメモリを消費します。
大規模モデルのトレーニングでは、このメモリ消費が深刻な問題となります。
第二に、計算コストが高いという点です。
v_t の更新やバイアス補正の処理が計算を複雑化し、トレーニングの速度を低下させます。
モデルサイズが増大するほど、これらの問題は一層深刻化します。
例えば、70億のパラメータを持つ Llama2 モデルを Adam でトレーニングする場合、50GB 以上のメモリが必要となります。
このような高コストな課題を解決するために提案されたのが、SGD-SaI です。
SGD-SaI 革新的なアイデアとは？
SGD-SaI の核心にあるのは「二次モーメントは本当に必要なのか？」という問いかけです。
従来の Adam が過去の勾配の履歴を基に学習率を調整するのに対し、SGD-SaI は最初の1回の勾配情報だけを使い、パラメータごとのスケーリング係数を固定するという革新的なアイデアに基づいています。
この手法では、各パラメータブロックごとに「勾配の信号対雑音比（g-SNR）」を計算します。
g-SNR は、勾配のL2ノルムと分散の比率を基にしており、これを用いてパラメータごとの学習率スケールを決定します。
特筆すべきは、このスケールが最初のイテレーションで計算され、その後は固定される点です。
これにより、メモリの大幅な削減と計算の効率化が実現しました。
SGD-SaI の動作の仕組み
SGD-SaI の動作はシンプルかつ効果的です。
最初のバッチで、各パラメータブロックの g-SNR を計算し、その結果に基づいてスケーリング係数を決定します。
この係数は、各パラメータの勾配更新に適用され、以降は固定され続けます。
これにより、計算負荷が大幅に軽減されます。
最初に、各パラメータブロックに対して g-SNR を計算します。
この g-SNR は、パラメータの勾配の大きさと変動の関係を数値化したもので、これを基にスケーリング係数が導き出されます。
その後、トレーニングが進むにつれて、g-SNR に基づくスケールは固定されたままになりますが、勾配は通常の SGD の手順に従って更新されます。
これにより、g-SNR の一度の計算だけで済むため、計算コストを大幅に削減しつつ、適応型の最適化の利点を得ることができます。
実験結果とその効果
SGD-SaI の有効性は、いくつかの実験で確認されています。
まず、メモリ消費の削減についてです。
GPT-2（1.5Bパラメータ）では、SGD-SaI を使用することで、メモリ使用量は 5.93GB にまで減少しました。
これは従来の Adam の半分以下の水準です。
Llama2（7Bパラメータ）では、Adam ではメモリ不足（OOM）エラーが発生しましたが、SGD-SaI では 25.15GB のメモリでトレーニングが可能でした。
トレーニング速度についても、SGD-SaI の利点が明らかです。
パラメータ更新時間は、Adam の約3分の1に短縮されました。
v_t の計算やバイアス補正が不要になったことで、トレーニングの効率が大幅に向上しました。
また、精度の向上も確認されています。
Vision Transformer（ViT）のトレーニングでは、SGD-SaI は AdamW と同等の精度を達成し、GPT-2 のトレーニングでは、SGD-SaI が Adam-mini よりも良い収束性能を示しました。
まとめ
SGD-SaI は「最初の一回がすべて」という考え方に基づき、これまでの「適応型最適化」の常識を覆しました。
メモリ使用量を大幅に削減し、トレーニング速度を飛躍的に向上させるだけでなく、精度も従来の手法と同等かそれ以上の水準に到達しました。
LLM や ViT などの大規模モデルの分野において、SGD-SaI は新たなスタンダードとなる可能性があります。
「適応型最適化は必須ではない」という事実は、研究者や実務者に新たな視点を提供することでしょう。
参考：No More Adam: Learning Rate Scaling at Initialization is All You Need
コメント