今回の論文
今回取り上げるのは、Jiawei Zhao、Zhenyu Zhang、Beidi Chen、Zhangyang Wang、Anima Anandkumar、Yuandong Tian による論文「GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection」です。2024年3月6日に arXiv で公開され、その後 ICML 2024（PMLR 235）でも採択されています。研究分野は、LLM 学習最適化、メモリ効率化、低ランク近似です。公開元は PMLR で、URL は https://proceedings.mlr.press/v235/zhao24s.html です。
この論文を選んだ理由は、単なる「軽い fine-tuning 手法」ではなく、全パラメータを更新する学習を維持したまま、optimizer state のメモリを大きく削る からです。実務では、モデルを作るより「学習メモリが足りず試せない」ことのほうが先に詰まります。GaLore はそこに対して、LoRA とフルファインチューニングの中間ではなく、別の設計軸を示しています。
どんな技術か
GaLore は、学習中に得られる勾配を低ランク空間へ射影し、その圧縮表現の上で optimizer state を持つことで、学習メモリを削減する手法です。
ポイントは、LoRA のように「更新できる重みを低ランク行列に制限する」わけではないことです。GaLore はモデルの重み自体はフルサイズのまま更新しつつ、重い optimizer state だけを低ランク側へ逃がします。つまり、探索空間は full-rank のままで、メモリ消費だけを low-rank 化する発想です。
言い換えると GaLore は、重みを小さくする技術ではなく、学習時に保存し続ける情報の持ち方を変える技術です。そのため、事前学習にも fine-tuning にも使いやすく、既存の AdamW や Adafactor 系 optimizer に差し込みやすいのが特徴です。
課題
この技術が解決しようとしているのは、LLM 学習で optimizer state が巨大化し、実験の自由度を奪う問題です。
何が難しいのかというと、LLM 学習では重みそのものより、勾配や一次・二次モーメントなどの optimizer state が大きなメモリを消費するからです。とくに AdamW 系では、各パラメータに対して追加の状態を持つため、モデルサイズが増えるほど学習メモリは急速に膨らみます。学習を始める前に、GPU メモリの制約で設計案が落ちることも珍しくありません。
既存の方法では、LoRA のように更新対象を低ランク行列へ絞る方法や、ZeRO・offloading・checkpointing のように学習状態を分散・退避する方法が使われます。ただし前者は、探索空間を低ランク部分空間へ制限するため、事前学習や性能重視の fine-tuning ではフル更新に届かないことがあります。後者は有効ですが、実装や実行環境が複雑になり、単一 GPU で気軽に試す方向とは少し違います。
なぜこの課題を解く必要があるのかというと、モデル改善の試行回数そのものが、学習コストで制限されるからです。小さなチームや個人開発では、学習アルゴリズムの良し悪しより先に「そもそも回るか」が問題になります。GaLore はそこに対して、フル学習を諦めずにメモリだけ削る という実務上かなり重要な選択肢を与えます。
用語解説
- Optimizer State
- AdamW などが持つ一次モーメントや二次モーメントの内部状態です。GaLore が主に削減するのはここで、重みを圧縮するのではなく、学習状態を低ランク表現に置き換える点が重要です。
- 低ランク近似
- 大きな行列を、より小さいランクの行列積で近似する考え方です。GaLore では勾配行列が持つ主要な変化方向だけを残し、その空間で更新情報を保持します。
- LoRA
- 重みに低ランク更新行列を追加する fine-tuning 手法です。GaLore を理解する上では、LoRA が「更新先を低ランクに制限する」のに対し、GaLore は「勾配表現だけを低ランクにする」という違いが重要です。
- SVD（特異値分解）
- 行列の主要な方向を取り出す代表的な分解方法です。GaLore は勾配の重要な方向を得るために低ランク射影を使っており、その直感をつかむ上で SVD 的な見方が役立ちます。
- Per-layer Weight Update
- レイヤーごとに勾配蓄積後すぐ更新して勾配メモリを解放する実装上の工夫です。GaLore の論文本体に加えて公式実装で重要な要素で、単一 GPU 上で大きなモデルを回す現実的な鍵になっています。
技術の仕組み
GaLore の核は、勾配をそのまま optimizer に渡すのではなく、低ランク空間へ写してから更新し、最後に元の重み空間へ戻すことです。これにより、重み更新は full-rank のまま保ちつつ、optimizer state のサイズを小さくできます。
基本アイデア
通常の full fine-tuning では、各重み行列 W に対して同じ形の勾配 G が出ます。AdamW なら、この G と同じ大きさの一次モーメント m と二次モーメント v を保持するので、行列が大きいほどメモリ負担も大きくなります。
GaLore はここで、勾配 G を低ランク行列へ射影した G_low で扱います。optimizer はこの G_low に対してモーメントを更新するので、保持する状態はフルサイズではなく、ランク r に応じた小さな表現で済みます。そのうえで、得られた低ランク更新量を元の形へ戻して W を更新します。
要するに、更新先のパラメータは大きいまま、更新を計算・記憶する座標系だけを小さくする のが GaLore です。
勾配をどう圧縮するのか
論文の直感は、ニューラルネットの勾配には強い冗長性があり、毎ステップで全方向を同じ精度で追わなくても主要な更新方向を押さえれば十分、というものです。そこで各重み行列の勾配に対して低ランク射影を作り、主要な方向だけを残した圧縮表現を optimizer に渡します。
この射影は一度作って終わりではなく、学習の進行に合わせて一定間隔で更新されます。公式実装でも update_proj_gap という形で、何ステップごとに射影基底を更新するかを指定できるようになっています。つまり GaLore は、変化する勾配構造に追随する動的な low-rank 圧縮 です。
LoRA と何が違うのか
LoRA は、更新そのものを A B のような低ランク行列へ限定します。そのため optimizer state が小さくなる一方で、モデルは最初から「低ランク更新しかできない」前提で学習します。
GaLore は違います。勾配を一時的に低ランク空間へ落とし込みますが、最終的な重み更新は元の重み W に対して行います。したがって、探索するパラメータ空間は full-rank のままです。ここが、事前学習や精度重視タスクでも GaLore が比較的性能を保ちやすい理由です。
Optimizer への組み込み方
GaLore は特定の optimizer に閉じた手法ではありません。論文と公式実装では AdamW、8-bit AdamW、Adafactor 系に組み込める形になっています。README でも「2行程度の変更で既存 optimizer に差し込める」と説明されています。
実装上は、GaLore 対象パラメータに rank、scale、update_proj_gap、proj_type などを与えて専用 optimizer を使います。つまり、モデル構造を大きく変えるのではなく、optimizer 側で low-rank 勾配空間を管理する 設計です。
8-bit 化と Per-layer Update
論文では GaLore に加えて 8-bit optimizer state も組み合わせています。これにより、低ランク化した optimizer state 自体もさらに小さくでき、BF16 ベースラインに対して総学習メモリを大幅に削減できます。
さらに公式実装では、各パラメータにフックを付けて、勾配蓄積直後にレイヤー単位で optimizer step を走らせる per-layer weight updates も紹介されています。これは勾配テンソルを長く保持しないための工夫で、optimizer state だけでなく gradient memory も削減しやすくなります。論文の発想を、単一 GPU で回る形まで落とし込んだ実装上の工夫 と言えます。
どこに効くのか
GaLore が一番効くのは、巨大な線形層が多く、Adam 系 optimizer の state が支配的になる学習です。逆に言うと、activation memory が完全な支配要因である設定では、GaLore だけで全問題が解決するわけではありません。そのため実務では、activation checkpointing や mixed precision と組み合わせて使うのが現実的です。
実験と結果
論文では、GaLore が本当にフル学習に近い性能を保ちながらメモリを削減できるかを、事前学習と fine-tuning の両方で検証しています。
何を検証したのか
主な検証ポイントは次の4点です。
- optimizer state をどれだけ削減できるか
- フルパラメータ学習と比べて性能を維持できるか
- LoRA 系より有利な場面があるか
- 8-bit 化や per-layer update を組み合わせると実機でどこまで回るか
どんなデータセットや評価指標を使ったのか
事前学習では、LLaMA 1B と 7B 規模のモデルを C4 データセットで最大 19.7B トークンまで学習し、主に validation loss を比較しています。fine-tuning では RoBERTa を GLUE ベンチマークで評価し、各タスクの精度や F1 を見ています。
この構成がよいのは、単に「一つの下流タスクで効いた」ではなく、事前学習のような大規模最適化でも、一般的な fine-tuning でも効果を確認している ことです。
Optimizer State を最大 65.5% 削減
PMLR 版の要約では、GaLore は optimizer state のメモリを最大 65.5% 削減しつつ、LLaMA 1B と 7B の事前学習で性能を維持したと報告されています。学習の本丸である optimizer state に直接効いているので、単なる周辺最適化ではありません。
実務的には、この数字の意味はかなり大きいです。学習メモリが 30% しか減らないのと、半分以上減るのでは、試せるバッチサイズやモデル規模が変わります。GaLore は「GPU 1 枚では厳しい」を「まずは試せる」へ変える可能性があります。
8-bit GaLore でさらに圧縮
同じ要約では、8-bit GaLore により optimizer memory を最大 82.5%、総学習メモリを 63.3% 削減したとされています。ここで重要なのは、optimizer state だけの削減率と total memory の削減率を分けて見ている点です。
学習では、重み、勾配、optimizer state、activation が混在するため、optimizer だけ軽くしても total memory は同じ率では減りません。それでも総量で 63.3% 減るなら、かなり実用的です。
フル学習に近い性能を維持
論文の主張は「軽いけれど性能が落ちる方法」ではありません。C4 での事前学習でも、RoBERTa の GLUE fine-tuning でも、GaLore は full-rank 学習に近い性能を示しています。これは LoRA 系と違い、探索空間を低ランク部分空間に固定していないことと整合的です。
とくに事前学習でこれを示した点に価値があります。fine-tuning 向け手法は多いですが、最初からモデルを育てる段階でも使える低メモリ化 は、活用場面がかなり広いからです。
24GB コンシューマ GPU で 7B 事前学習を実証
PMLR の要約と公式 README では、GaLore と 8-bit optimizer、さらに per-layer weight updates などを組み合わせることで、24GB メモリの RTX 4090 のような consumer GPU で 7B モデルの事前学習が可能になったと報告しています。README には、activation checkpointing 付きで batch size 16、約 22.8GB 使用という実行例も記載されています。
もちろん、これは学習速度や運用容易性まで含めて万能という意味ではありません。ただし「単一 GPU で 7B の pretraining は現実的でない」という前提を崩した点は大きいです。
何に使える？
GaLore が向いているのは、フル更新を捨てたくないのに、学習メモリが足りない場面です。
ドメイン特化モデルの継続事前学習
業界文書、社内ナレッジ、コード、論文などに合わせて継続事前学習したい場合、LoRA では更新自由度が不足することがあります。GaLore なら full-parameter learning を維持しやすいので、モデル全体の表現をじわっと変えたい場面に向いています。
小規模 GPU 環境での本格 fine-tuning
スタートアップや個人開発では、A100 を何枚も前提にできないことが多いです。GaLore は、単に「LoRA で済ませる」以外の選択肢として、限られた GPU でも full fine-tuning 寄りの実験をしやすくします。
学習アルゴリズムの比較実験
optimizer、loss、データ混合比、カリキュラム学習などを比較したいとき、学習メモリが制約になると試行回数が減ります。GaLore はメモリ制約を緩めることで、プロダクトの精度改善より一段下にある「研究開発の探索速度」を上げるのに役立ちます。
低コストな社内モデル開発基盤
社内専用アシスタントやオンプレ向けモデルでは、学習専用クラスタを大きく持てないことがあります。GaLore 系の手法を前提に学習基盤を組むと、少ない GPU で回せるジョブの幅が増え、モデル開発の内製化に役立ちます。
開発や事業へのヒント
この論文から得られる一番大きなヒントは、学習効率化を「更新対象を減らす」方向だけで考えなくてよい、ということです。何を学習するかではなく、学習状態をどう持つか も大きな設計変数になります。
LoRA かフル学習かの二択を崩せる
自分で AI アプリを作るとき、これまでは「安く済ませるなら LoRA」「性能を取りにいくならフル学習」という二択になりがちでした。GaLore はその間に、「フル学習に近い自由度を保ちながらメモリだけ下げる」選択肢を置けます。これは学習戦略の設計をかなり柔らかくします。
学習基盤の差別化ポイントになる
多くのプロダクトは推論最適化に注目しがちですが、業務特化モデルや継続学習を回す会社では、学習 1 回あたりのコストがそのまま競争力になります。GaLore のような optimizer 側の工夫は、同じ GPU でも回せる実験量を増やせるため、学習基盤の差別化要素になりえます。
小さなチームほど価値がある
巨大 GPU クラスタを持つ組織なら、ZeRO や分散学習で押し切れる場面もあります。一方で小さなチームは、設定の複雑さより単純なメモリ削減のほうが価値があります。GaLore は optimizer レベルで差し込みやすく、単機運用とも相性がよいので、小規模開発ほど恩恵を受けやすいです。
今後注目すべき方向性
今後は、低ランク化が「adapter を差す話」から「勾配・optimizer・通信量をまとめて圧縮する話」へ広がる可能性があります。これは推測を含みますが、GaLore が示したのは、低ランク性は重み更新先より学習ダイナミクス側に使ったほうが得をする場面がある ということです。この考え方は、分散学習の通信圧縮や federated learning にも波及しそうです。
限界
GaLore にも明確な限界があります。
まず、削減対象の中心は optimizer state なので、activation memory が支配的な設定では効果が相対的に小さくなります。超長系列や巨大 batch の学習では、checkpointing や sequence parallel など別の対策も必要です。
次に、低ランク射影の品質は rank と射影更新間隔に依存します。rank を下げすぎると勾配情報を落としすぎる可能性がありますし、射影更新をまれにしすぎると現在の勾配構造に追随しにくくなります。したがって、GaLore は「入れれば自動で最適」ではなく、ある程度のハイパーパラメータ設計が要ります。
また、論文は full-rank 学習に近い性能を示していますが、すべてのモデル、すべてのデータ、すべての optimizer で同じとは限りません。特に very small rank や特殊なアーキテクチャでは、再検証が必要です。
実装面でも、純粋な Hugging Face Trainer 既定設定にそのまま乗せるよりは、optimizer 周りの理解が必要です。単一 GPU 実装では per-layer update など PyTorch の新しめのフック機能も関わるため、環境依存の注意が要ります。
最後に、GaLore は学習を回しやすくする技術であって、学習データの質や評価設計の悪さを解決するものではありません。データが弱ければ、低メモリで何度回せてもモデル品質は上がりません。
よくある質問
Q. GaLore は LoRA の上位互換ですか？
A. 上位互換というより、目的が違います。LoRA は少ない学習パラメータで効率よく適応するのが得意で、GaLore は full-parameter learning を保ちながら optimizer memory を減らすのが目的です。小さな差分適応なら LoRA、モデル全体をしっかり動かしたいなら GaLore が有力です。
Q. GaLore を使えば activation checkpointing は不要ですか？
A. 不要にはなりません。GaLore が主に削るのは optimizer state なので、系列長や batch size が大きいと activation memory は依然として重いです。論文周辺の実装でも checkpointing との併用が前提になる場面があります。
Q. 事前学習にも使えるのがなぜ重要なのですか？
A. fine-tuning だけなら LoRA でも足りることがありますが、継続事前学習やゼロから近い学習ではモデル全体の表現を更新したいことが多いからです。GaLore はその局面でもメモリ削減を効かせられる点が実務的です。
Q. rank はどう決めればよいですか？
A. 論文と実装では 128 や 1024 などの設定例がありますが、最適値はモデルサイズや対象層で変わります。まずは論文の近い設定を起点にし、validation loss と使用メモリの両方を見ながら調整するのが現実的です。
Q. 小規模プロダクトでも導入する価値はありますか？
A. あります。とくに GPU が少なく、LoRA では性能が足りないが本格的な分散学習基盤を組む余力もない場合に価値があります。GaLore はその中間を埋める技術として考えるとわかりやすいです。
今日の学び
この論文は、LLM 学習で optimizer state が巨大化し、フル学習がメモリ制約で難しくなる課題を扱いました。そこに対して、勾配を低ランク空間へ射影して optimizer state を圧縮しつつ、重み更新自体は full-rank のまま保つ GaLore という技術で解こうとしました。
ここから得られるヒントは明確です。AI 開発では、更新対象を減らすだけが効率化ではありません。学習状態の持ち方を変えるだけでも、試せるモデル規模と実験回数は大きく変わる ということです。