最近のOptimizerの研究について
Optimizerはロマンである。Optimizerを変えると学習が速くなったりlossがこれまでより下がったりする。すばらしい！ 本記事では今年出てきた新しいOptimizerやその関連研究を眺めてみたい。
Cautious Optimizers: Improving Training with One Line of Code
- gradientと更新方向の符号が一致している場合にのみ、そのパラメーターを更新する。なんのこっちゃと思われるかもしれないが、momentumの概念が入ると、gradientと実際の更新方向の食い違いは容易に起こりえる。
- 更新するパラメーターの数に応じて、更新量をスケールする。更新するパラメーターの数が少ない場合は、更新するパラメーターに対してはアグレッシブに更新を行う。
- AdamWとLionに適用しているが、Momentum SGDにも適用できそう。
- LLaMA-60Mから1Bまでで検証され、c-Adamは1Bモデルになって効果が出てきている（小さなモデルでは効果がない）一方で、c-Lionでは小さなモデルでも効果があるが、1Bモデルでは効果が薄れてきているように見える。もっと大きなモデル、例えば7Bモデルとかでの実験結果がぜひ見たいが、7Bの学習ともなると結構なお金がかかりそう。自分ではできない。
- 計算量は微妙に増えるが、メモリ消費量がほぼ増えない点はよい。
- timmにも最近実装され、「cautious optimizerは実験した範囲では常に効果があった」と述べられており、再現性についてある程度期待できそうである。（こちらはViTを学習している。）⇒ https://huggingface.co/rwightman/timm-optim-caution
DeMo: Decoupled Momentum Optimization
- Adamの著者であるKingmaがLast authorとして入っていることでX上でもにわかに話題となっている。実装もある。
- momentumから、fast moving componentを取り出し、fast moving componentについては即座にパラメーターの更新に利用し、fast moving componentを除去したslow moving componentのみをmomentumの新成分とし、momentumを更新する。
- fast moving componentとは、要するに、momentumの高周波成分であり、fast moving componentの取り出しには離散コサイン変換を利用する。ソースコードを参照したところ、チャンク（デフォルトではチャンクサイズは64x64）に分解し、チャンクごとにDCTにかけてtopk（デフォルトでは32）個の大きな成分を取り出す、という処理になっていた。Linear Layerだとweightは二次元の行列で、その64x64のブロックをDCTして上位32個の要素を取り出すだけなので、データ量は元の行列サイズからすると、確かに、かなり小さい。
- momentumにDCTをかけて、値の大きな要素だけをパラメーターの更新に使い、使わなかった部分だけをmomentumの更新に使う。卵の黄身と白身を使い分けるようなイメージだろうか。
- fast moving componentのみ複数ノード間で共有すればいいので、分散学習において必要な通信量が下がる、というのが売りとして挙げられている（そう、slow moving componentはクラスタ間で共有しないのである。つまり、momentumは各クラスタにおいて全然違う値になっている。この辺りはfederated learningっぽい。実際、論文でもfederate learningについても言及されている）が、どちらかというと、注目されているのは圧倒的にlossの低下が速い、という点だろう。
- gradient自体を圧縮しようという話はたくさんあるし、momentumを圧縮しつつ共有する話もある（例えばLu2022）が、momentumを圧縮した上で、その圧縮後の情報だけを本質部分だとみなしてパラメーター更新に利用する点は、類似する研究を知らない。この研究の新規性はここにあると言ってよいだろう。
- Fig.1を見ると、step単位で比較した場合のAdamWに対する収束性能の向上がすごい。ただ、OLMo-300MだとAdamWを圧倒しているが、1BになるとAdamWに少し差を縮められているように見える。モデルが7B、70Bと大きくなった場合にどうなるかは未知数である。
- 2次モーメントを保存する必要がない点はAdamに対するメリットと言える。
- 原理上は、Cautious Optimizerと組み合わせられそうである。
- 結果が良過ぎてtoo good to be trueという感想である。小さいモデルではうまく行くことが実験的に示されているが、大きなモデルだとどうなるかは未知数である。コードもあることだし、第三者の再現実験を待ちたい。
Adam-mini: Use Fewer Learning Rates To Gain More
- Adamには一次のmomentum mと二次のmomentum vがあるが、vをそれぞれのパラメーターに与えるのはredundantである、というのが著者らの主張である。
- パラメーターをブロックに分割にして、ブロック単位でvを割り当てることで、Optimizerのメモリ消費量を大幅に削減し、その分、学習の高速化が可能となる。
- 色々な考察の結果、Transformerにおいては、以下のようにブロック化するとよい、となったそうである。
- embed layerとoutput layerはtoken単位でブロック化する
- queryとkeyはhead単位でブロック化する
- valueとMLPのprojとmlpはoutput neuron単位でブロック化する
- 単純にOperator単位でブロック化すればよい、というわけではなく、アーキテクチャ毎に考察が必要となる点はやや面倒であると言わざるを得ない。そもそも、全部output neuron単位でブロック化してしまえばいいのでは、という気もする。
- 面白いことに、LoRAにおいてはAdamよりもAdam-miniの方が収束が速い、とある。パラメーターがまとめられた分だけノイズが減って学習が高速になる…のだとしたら、フルスクラッチでの学習も速くならないとおかしい気がするが、フルスクラッチでの学習は、step単位でみるとAdamとcomparableであり、特に速くなっていない。この辺は個人的には直感に反している。
ADOPT: Modified Adam Can Converge with Any $\beta_2$ with the Optimal Rate
- Adamにはハイパーパラメーターがいくつかあるが、αが学習率に相当する最も重要なパラメーターで、次に重要なパラメーターがβ2である。β2はできるだけ大きな値である方がよいが、大きくしすぎると学習が不安定になることが知られている。
- ADOPTは、Adamで$\textbf{m}/(\sqrt{\textbf{v}}+ε)$となっているところを$\textbf{m}/(\max(\sqrt{\textbf{v}},ε)$ にすることと、処理の順番を少し入れ替えることで、β2に対する安定性が向上する、とされている。
- 論文ではalgorithm 2として、ここにさらにclipping処理が追加されたものが提案されており、実用する場合はこちらの方を使うべきである。
- GPT-2のpretrainingに使った際の比較だと、Adamとほぼ同等の結果が得られていて、なおかつ、batch sizeが小さい場合にも安定して学習ができている。
- 収束が速く期待するのではなく、安定性が向上し、よりaggressiveなβ2が使いやすくなる手法である、と考えるべきであろう。
MARS: Unleashing the Power of Variance Reduction for Training Large Models
- variance reductionは最適化の分野で研究されてきたが、DNNにはあまり有効ではないとされてきた。その理由としてdropoutやbatch normalizationの存在が挙げられてきたが、LLMではこれらは使われていないので、variance reductionを再検証する価値があるので検証してみましょう、というのが論文のモチベーションらしい。
- variance reductionとは: 勾配を計算する際に、今のパラメーターを使って計算した勾配($a$)と、一つ前の時刻のパラメーターを使って計算した勾配($b$)の2つを計算し、momentumを計算する際には$m_t = a - b + m_{t-1}$とする。$b$を引くことによって、データに由来するノイズをキャンセルできる、らしい。
- 実際には全部キャンセルしてしまうのではなく、$a - β_1 b$を用いる。($β_1$は0より大きな定数。) 後の計算の都合上、論文中ではEMA（Exponential Moving Averate）として解釈できる別の表現が主として用いられる。
- MARSはvariance reduction、online mirror descent、gradient clippingの組み合わせでできている。
- Hessianが出てくるが、Hessianをそのまま計算するのは実用的ではないので、なんらかの近似手法を用いる。論文中では、対角行列で近似する手法として、AdamとLionにMARSを組み合わせたMARS-AdamとMARS-Lionが具体的なアルゴリズムとして定義されている。また、MARS-Shampooも定義されており、こちらではHessianの近似にSVDが使われるが、SVDも重たいので、sketchingやNewton iteration等で近似される。
- 読んだ限りでは、MARS-Adamは、Adamのmomentum計算部分に先程のvariance reductionとgradient clippingが入った感じになっている。
- 実験はGPT2-small(125M)からGPT2-large(770M)で行われており、より大きなモデルに適用した場合に本当に効果があるのかは今後の追試が待たれる。
- timmで行われた追試では、ViTに適用しており、そちらではAdamWと大差ない結果になっている。
SOAP: Improving and Stabilizing Shampoo using Adam
- MARSの論文読むのに疲れちゃって 全然数式が読めなくてェ…
- また今度がんばって読みます
The Road Less Scheduled
- 学習step数Tを決めずに、つまり、明示的なLR decayを行わずに、既存のLR schedulingを使った手法よりも高い性能を常に得られると主張する方法。論文名よりもSchedule Freeという名前のほうが有名だろう。
- 以下はSGDをSchedule Free化する際のやり方である。
- $x, y, z$の3種類のテンソルが出てくる。$z$が通常のSGDにおけるパラメーターであり、$x$は全ステップにおける$z$の平均である。勾配は$z$と$x$を比率$β$で混ぜて作った$y$に対して計算され、その結果は$z$に反映される。
- そういえば、昔averaged SGDってあったな…と思い出した。前前前職のtech blogに書いてた。
- averaged SGDでは「現在のパラメーター(=$z$)」に対して勾配を計算していたが、本手法では「すべてのステップでのパラメーター平均($x$)」も勾配計算の際に用いる点が違う。
- RAdamをScheduleFree化した人がすごく詳しい解説記事を書いてくれている。⇒ 全ての学習率スケジューリングを過去にするOptimizer
Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations
- 固定LRで学習しつづけて、最後に10k step程度の期間、学習率を1-sqrtの形で落としつつ学習すればcosine LR decayに匹敵する性能が得られる。事前にどれくらいのstep数で学習するかを決めずに学習を開始できるので、ハイパラが減って楽になる。
- 手法としては簡単だが、内部で何が起こっているのか、さっぱりわからない。
- 後述のINTELLECT-1の学習で既に実用されている。
Muon: An optimizer for hidden layers in neural networks | Keller Jordan blog
- NanoGPT speedrunningによって一部で有名かもしれない。論文はまだなさそう。
- Newton-Schulz iterationという手法を用いてmomentumを近似的に直交化し、それを用いてパラメーターを更新する。
- 直交化をしないよりした方がいいというのはなんとなく想像はつくのだが、それだけでこんなに収束が速くなるというのは驚きである。
DiLoCo: Distributed Low-Communication Training of Language Models
- 2023年の論文だし、OptimizerというよりはFederated Learningの話なのだが、個人的には今年読んだ中で一番面白かったといってもよいくらいの論文なので、おまけで紹介しておく。
- 分散学習する際に、サーバー群をいくつかのクラスタに分割し、クラスタ間通信は頻度を低くすることで、大規模な高速ネットワークを用いずに分散学習を可能とする。
- 大規模な高速ネットワークというとGoogle TPUクラスタの光ファイバー通信とNVLinkが想起されるが、前者は物理的なハードウェアとして購入できないし、大規模なNVLinkはとても高価である。
- inner optimizerとouter optimizerの2種類のoptimizerを用意する。inner optimizerはLLM学習によく用いられるAdamWが使われる。
- H step学習した後、学習前後でのパラメーターの差分をgradientだとみなして（なんと大胆な近似！）クラスタ間で共有し、outer optimizerを使ってパラメーターを更新する。
- この「学習前後でのパラメーターの差分をgradientだとみなしてoptimizerを適用する」という操作が面白い。こんなところに工夫を入れる余地があるとは。
- なぜか、outer optimizerとしてはNesterov SGDが圧倒的によい。
- 知る限りでは、INTELLECT-1 Release The First Globally Trained 10B Parameter Model がDiLoCoを使って学習された最も大きなモデルであり、LLaMA-7Bと学習token数はほぼcomparableで、性能の方もcomparableな感じである。10Bモデルであることを考えると7Bとcomparableという結果はやや寂しい。
- INTELLECT-1の学習では勾配がint8に量子化されていたようである。同じサイズに圧縮するならfp8にした方が絶対によいはずだが、実装の都合上int8になってしまっているのだろう。
- 勾配を量子化した理由は、inner optimizerで学習するステップ数があまり長いと収束が遅くなるためだそうだ。保守的に100step毎程度でouter optimizerを走らせるため、頻繁に勾配情報を通信する必要があり、通信時間を削減するために勾配を量子化する必要があったそう。ここだけ読むと、DeMoと相性がよさそうに思える。（勾配のデータサイズを数十分の一にできるなら量子化の必要はなく、しかもデータサイズが小さいなら頻繁に通信してもよい。）
感想
- 収束を早めるという意味で効果が期待できそうかなと思っているのは、個人的にはCautious OptimizersとDeMoの2つである。
- Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durationsは10Bモデルで使われた実績から実証済み…と言いたいところだが、大規模学習時の有効性についてはある程度検証されている、くらいまでしか言えないか。こういう、自分では思いつけそうにないが実装は簡単、という手法が好きである。
- MARSはSVDでHessian近似していけるぜ、とか書いてあるけど、実際問題、どれくらい現実的なんだろうか。勾配のサイズが今の数倍になるというのはLLMにおいてはちょっと許容できそうにない気がするのだけれど。
- ところどころ記述をサボってしまったが、丁寧に書いていると時間がかかりすぎて公開できずに塩漬けにしてしまいそうなので、ここで公開してしまうことにする。ChatGPTかなにかに「論文中からいい感じに図や数式を引用して補足を書いてくれ」と指示したら勝手に作業してくれるようになったら、そろそろAGIの実現が近いと認めます。
- 自分で言い出しておいてなんだが、Optimizerは本当にロマンなのだろうか？ 読めば読むほど「大規模なモデルで使えるのかな？」という現実的なところが気になって疲れてしまった……。最近のちょっと大きな実験は、もはや個人で気軽に追試できる規模ではなくなってしまったので、ロマンよりもむしろ寂しさを感じる。