🤖 machine learning

FOAM: Blocked State Folding for Memory-Efficient LLM Training

原著者： Ziqing Wen, Jiahuan Wang, Ping Luo, Dongsheng Li, Tao Sun

公開日 2026-05-14

📖 1 分で読めます☕ さくっと読める

原著者： Ziqing Wen, Jiahuan Wang, Ping Luo, Dongsheng Li, Tao Sun

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

巨大で超知能なロボット（大規模言語モデル）に、書き方、会話、推論を教えることを想像してみてください。そのために、あなたは膨大な量の書籍（トレーニングデータ）の図書館をそのロボットに見せます。ロボットは間違いを犯し、自分の作業を確認し、内部の「脳の設定」（パラメータ）を調整して上達することで学びます。

この調整を行う標準的な方法は、アダムという非常に慎重な会計士のようなものです。アダムはすべての脳の設定に対して詳細な帳簿を維持し、現在の間違いだけでなく、過去の間違いの履歴も追跡して、どれほど変更を加えるべきかを決定します。

問題点：会計士が重すぎる
問題は、この「会計士」（オプティマイザ）が信じられないほど重いことです。数十億もの脳の設定を持つロボットの場合、会計士の帳簿はロボット自身の脳よりも2 倍のメモリを占有します。

比喩: 家を移動させようとしていると想像してください。家具（モデル）は大きいですが、引越しトラック（オプティマイザの状態）はそれよりもさらに大きいです。もしあなたが小さなトラックしか持っていない（限られたコンピュータメモリ）場合、家がどれだけ素晴らしくても、家を全く移動させることはできません。これが、研究者がより大きく、より賢いロボットをトレーニングするのを妨げる「メモリボトルネック」です。

従来の解決策：手抜き
引越しトラックを小さくしようとした以前の試みには欠点がありました：

家の一部を凍結する: 一部の家具の移動を止め、いくつかの小さな箱だけを調整します。これはスペースを節約しますが、家の柔軟性を低下させます（性能の低下）。
ぼやけた写真を撮る: 実際の家具を動かす代わりに、スペースを節約するために低解像度の写真を撮ります。しかし、これらの写真を撮るには多くの時間とエネルギー（計算オーバーヘッド）がかかります。
帳簿の共有: 異なる部屋に同じノートブックを共有させます。これはスペースを節約しますが、調整の精度を低下させます。

新しい解決策：FOAM（スマートな折りたたみ法）
この論文は、FOAM（Folded Optimizer with Approximate Moment）を紹介しています。FOAMを想像してください。それは、重要な詳細を失うことなく、巨大な引越しトラックを小さなバンに収めるための巧妙な「折りたたみ」技術を使う、達人のパッカーです。

以下に、簡単な比喩を用いて FOAM がどのように機能するかを示します：

1. 「ブロック折りたたみ」（圧縮）

FOAM は、すべての脳の設定を個別に追跡する代わりに、それらを小さなブロック（8 または 16 の設定のグリッドなど）にグループ化します。

比喩: 100 人の人々が長い列に並んでおり、あなたの任務は彼らの平均身長を記憶することだと想像してください。100 個の個別の数字を書き留める代わりに、10 人ずつの 10 チームにグループ化します。そして、各チームの平均身長を表す1 つの数字だけを書き留めます。
結果: これにより、「帳簿」に必要なメモリが大幅に削減されます（最大 90% 削減）。

2. 「残差補正」（セーフティネット）

平均値だけを使用すれば、各人のユニークな詳細が失われます。チームの中に背の高い人が一人いて、もう一人は背が低い場合、平均値はそのことを隠してしまいます。

比喩: FOAM は賢明です。「チームの平均」を書き留めた後、実際の人物と平均値の間の差（誤差）を素早く計算します。これを「残差」と呼びます。
トリック: FOAM はこの差を永久に保持しません。それを計算し、その瞬間の更新を修正するために使用し、その後破棄します。まるで、料理人がスープを味わい、塩加減を調整した後、一度味わったスプーン一杯の正確な味を記録するのではなく、単にそれを忘れるようなものです。
重要性: これにより、メモリが圧縮されているにもかかわらず、ロボットがすべての脳の設定のユニークな詳細を学習できることが保証されます。

3. 「展開」（復元）

実際にロボットの脳を更新する時が来ると、FOAM はそれらの圧縮された「チームの平均」を取り出し、フルサイズに拡張し、その上に「残差」の補正を加えます。

結果: ロボットは、重い会計士を使った場合と全く同じように正確な更新を受け取りますが、コンピュータが運ぶ必要があったのは、折りたたまれた小さなバージョンだけでした。

論文が主張すること（結果）
著者らは、6000 万パラメータの小さなモデルから 70 億パラメータの大きなモデルまで、さまざまなロボットモデルのトレーニングにおいて FOAM をテストしました。彼らが発見したことは以下の通りです：

莫大なメモリ節約: FOAM はオプティマイザに必要なメモリを最大**90%削減します。合計すると、トレーニングのメモリフットプリントを約50%**削減します。これにより、同じハードウェアでより大きなモデルをトレーニングしたり、同じモデルをはるかに高速にトレーニングしたりすることが可能になります。
性能の低下なし: 重い圧縮にもかかわらず、FOAM でトレーニングされたモデルは、標準的な重い会計士でトレーニングされたモデルと同じくらい（場合によってはそれ以上）よく機能します。それらはより速く学習し、より高い精度に達します。
速度: 重い荷物を運ぶ必要がないため、トレーニングプロセスは高速化されます。
汎用性: これは LLaMA、Qwen、その他さまざまなタイプのロボットアーキテクチャで機能し、他のメモリ節約テクニックと組み合わせることもできます。

要約すると
FOAM は魔法のスーツケースのようなものです。それは、情報を丁寧に折りたたんで小さなスペースに詰め込み、何も失われないようにするための迅速な「補正ノート」を追加することで、膨大な量の情報（オプティマイザのメモリ）を小さなスペースに詰めることを可能にします。これにより、研究者は超高価で巨大なコンピュータを必要とすることなく、より賢く、より大きな AI モデルを構築できます。

技術的サマリー：FOAM – メモリ効率化された LLM 学習のためのブロック状態折りたたみ

1. 問題定義

大規模言語モデル（LLM）は、パラメータ数と学習データの増加を通じて卓越した性能を達成してきましたが、その規模は学習中に重大なメモリボトルネックをもたらします。事実上のオプティマイザである Adam（および AdamW）は、すべてのパラメータに対して 1 次および 2 次モーメント推定値を保存する必要があり、モデルサイズの 2 倍に相当するオプティマイザ状態のメモリオーバーヘッドを発生させます。これにより、70 億パラメータ（7B）以上の大規模モデルの事前学習や微調整はメモリ制約に縛られ、小さなバッチサイズであっても数百ギガバイトの VRAM を必要とすることが頻繁にあります。

既存のメモリ効率化ソリューションはトレードオフに直面しています：

パラメータの凍結（例：LoRA）： 低ランク部分空間への更新を制限し、モデルの表現力や性能を制限する可能性があります。
低ランク射影（例：GaLore、APOLLO）： 特異値分解（SVD）またはランダム射影を用いてオプティマイザ状態を圧縮します。効果的である一方で、SVD は多大な計算オーバーヘッドをもたらすため、ランダム射影は厳密なメモリ制約下で性能を低下させる可能性があります。
ブロック共有（例：Adam-Mini）： ストレージを削減するためにブロック間で 2 次モーメント推定値を共有しますが、1 次モーメントのオーバーヘッドを高く留めたり、メモリ削減を約 50% に制限したりすることが多いです。

2. 手法：FOAM

著者は、最適化忠実度を犠牲にすることなくオプティマイザ状態を圧縮するように設計された構造化された状態折りたたみフレームワークであるFOAM（Folded Optimizer with Approximate Moment）を提案します。FOAM は、ブロック平均化と残差補正を組み合わせた原理に基づいて動作します。

中核メカニズム

時刻ステップ $t$ における勾配行列 $G_t \in \mathbb{R}^{m \times n}$ と折りたたみレベル $l$ が与えられたとき：

ブロック平均化（Fold）： 勾配は $2^l$ 個の連続する要素からなるブロックに分割されます。折りたたみ演算子 $A^{(l)}$ が各ブロックの平均を計算し、勾配の次元を $n$ から $n/2^l$ に圧縮します。1 次および 2 次モーメント（ $\tilde{M}_t, \tilde{V}_t$ ）は、これらの圧縮された勾配に基づいて更新されます。
展開（Unfolding）： 全パラメータ空間に更新を適用するために、折りたたみ演算子 $E^{(l)}$ が圧縮されたエントリを元の次元に複製します。
残差補正： 更新を直接複製すると、ブロック内のすべての要素が同一の値を共有し、区別は初期化に依存するのみとなります。失われた情報を回復するため、FOAM は元の勾配と折りたたみ・展開版との間の誤差を表す残差項 $R_t = G_t - \tilde{G}_t E^{(l)}$ $R_{t} = G_{t} - \tilde{G}_{t} E^{(l)}$ を計算します。
- 1 次モーメントは以下のように再構成されます： $M_t = \tilde{M}_t E^{(l)} + R_t$ 。
- 2 次モーメントは以下のように再構成されます： $V_t = \tilde{V}_t E^{(l)} + R_t^2$ 。
- 重要なのは、 $R_t$ はオンザフライで計算され、直ちに解放されることであり、メモリに保存されないことです。

ハイブリッド適用

FOAM は特に MLP（多層パーセプトロン）およびアテンションモジュールに適用され、埋め込み層には標準的な Adam が使用されます。スケーリング係数 $\alpha$ は、FOAM モジュールと Adam モジュール間の学習率比を調整し、これはメモリ効率化オプティマイザにおける一般的な慣習です。

3. 主要な貢献

革新的なオプティマイザ設計： FOAM は、Full-Adam と比較してオプティマイザ状態のメモリオーバーヘッドを最大90% 削減し、総エンドツーエンドのメモリフットプリントを約50% 削減します。射影ベースの手法とは異なり、射影行列を必要とせず、全パラメータ最適化をサポートします。
理論的保証： 標準的な非凸最適化の仮定（L-滑らか性、有界分散、有界勾配）の下で、著者は FOAM がバニラ Adam と同じ $O(\frac{\log T}{\sqrt{T}})$ の収束率を達成することを証明しています。残差誤差比 $\delta_l$ は 1 以下に抑えられ、主要な収束率に悪影響を及ぼしません。
広範な評価： この手法は、モデルサイズ（60M から 7B）、アーキテクチャ（LLaMA、Qwen、GPT、DeBERTa）、タスク（C4 での事前学習、GLUE および MMLU での微調整）のすべてにおいて検証されています。

4. 実験結果

論文は以下の実証的な知見を報告しています：

事前学習性能： FOAM は、検証パープレキシティ（PPL）の点でメモリ効率化のベースライン（GaLore、APOLLO、Adam-Mini、Muon）を一貫して上回り、同時に著しく低いメモリ使用量を維持します。例えば、LLaMA-1.3B において、FOAM-3 は Adam と比較して総メモリ消費量を約 50% 削減しながら、より低い PPL を達成します。
収束速度： FOAM はより速い収束を示し、他のメモリ効率化手法と比較して目標 PPL に到達するために必要なトレーニング反復回数が少ないことがよくあります。
スケーラビリティ： FOAM は、より大規模なモデル（LLaMA-3B、7B）および長いシーケンスの学習環境（最大 2048 トークン）においても優れた性能を維持します。
微調整： GLUE および MMLU ベンチマークにおいて、FOAM は Full-Adam や他のベースラインと同等かそれ以上の性能を達成し、事前学習を超えた有用性を示しています。
ロバスト性： 除去研究（アブレーションスタディ）により、残差項が性能にとって決定的であることが確認されました。これを除去すると収束が遅くなり、最終的な PPL が高くなります。FOAM は 8 ビット量子化と組み合わせた場合や、Adam-Mini や Muon などの他のオプティマイザと統合された場合にもロバスト性を維持します。

5. 意義と主張

この論文は、FOAM が大規模言語モデル学習におけるメモリボトルネックに対する実用的で、極めて効率的かつオプティマイザに依存しない解決策を提供すると主張しています。重みの凍結や高価な射影計算の必要性を排除することで、FOAM は劇的に削減されたメモリオーバーヘッドによる全パラメータ最適化を可能にします。著者は FOAM を低ランク射影技術に対する補完的な代替手段として位置づけ、理論的には Adam の収束と一致しつつ、実証的には優れたメモリ効率と、しばしば優れた学習性能を提供する手法として提示しています。

この研究は、FOAM が既存のハードウェア上でより大規模なモデルの学習を可能にするか、または現在の最先端モデルの学習コストを削減し、高性能な LLM 学習をよりアクセスしやすくすることを示唆しています。