Lion: Symbolic Discovery of Optimization Algorithms
わっしーアルゴリズム発見をプログラム探索として定式化する手法を提示し、ディープニューラルネットワーク学習用の最適化アルゴリズム発見に適用する。我々は、効率的な探索技術を活用して、無限かつ疎なプログラム空間を探索する。また、代理タスクとターゲットタスクの間の大きな汎化ギャップを埋めるために、プログラムの選択と簡略化戦略を導入する。(NASみたいな感じで探索して見つけた手法?)
本手法は、シンプルで効果的な最適化アルゴリズムであるLion (EvoLved Sign Momentum)を発見する。
これは、運動量のみを記録するため、Adamよりもメモリ効率が良い。
また、適応型最適化アルゴリズムとは異なり、符号演算によって計算される各パラメータの更新は同じ大きさである。
研究では、LionをAdamやAdafactorといった広く使われている最適化器と比較し、異なるタスクの様々なモデルを学習させることで、その効果を検証した。画像分類では、ImageNetにおいてViTの精度を最大2%向上させ、JFTでは事前学習計算を最大5倍削減することができた。vision-language contrastive learningでは、ImageNetにおいてゼロショットで88.3%、微調整で91.1%の精度を達成し、それぞれ従来の最高値を2%、0.1%上回ります。拡散モデルにおいては、より良いFIDスコアを達成し、学習計算量を最大2.3倍削減することで、LionはAdamを上回ります。自己回帰モデル、マスク言語モデリング、ファインチューニングでは、LionはAdamと同等以上の性能を示しました。
Lionの分析から、学習バッチサイズに応じて性能が向上することが明らかになりました。また、符号関数によって生成される更新のノルムが大きいため、Adamよりも学習速度が小さくて済む。さらに、Lionの限界を検証し、Lionの改善が小さいか、統計的に有意でないシナリオを特定する。Lionの実装は一般に公開されている。
わっしー
提案手法: AutoML-Zero (Real et al., 2020) は、トイタスクで評価しながら機械学習パイプラインのあらゆるコンポーネントを探索しようとする野心的な取り組みである。論文では、自動発見型オプティマイザの研究方向を踏襲し、特にAutoML-Zeroに触発されているが、最先端ベンチマークを改善できる効果的な最適化アルゴリズムの発見を目指すものである。
論文では、アルゴリズム発見をプログラム検索として定式化し、最適化アルゴリズムの発見に適用する方法を紹介している。主な課題は2つある。第一は、無限かつ疎なプログラム空間において、高品質のアルゴリズムを発見することである。もう一つは、小さな代理タスクから、より大規模な最先端タスクに汎化できるアルゴリズムをさらに選び出すことである。これらの課題に取り組むため、ウォームスタートとリスタートを用いた進化的探索、抽象的実行、ファネル選択、プログラムの簡略化などの技術を採用している。
論文では、シンプルで効果的な最適化アルゴリズムを発見している。Lionは、EvoLved Sign Momentumの略称である。このアルゴリズムは、様々な適応的アルゴリズムとは異なり、運動量のみを追跡し、更新を計算するために符号演算を活用することで、低いメモリオーバーヘッドと全ての次元で均一な更新の大きさを実現するものである。
Lionはそのシンプルさにもかかわらず、様々なモデル（Transformer、MLP、ResNet、U-Net、Hybrid）とタスク（画像分類、視覚-言語対照学習、拡散、言語モデリング、微調整）において優れた性能を示している。特筆すべきは、BASIC（Pham et al., 2021）のAdafactorをLionに置き換えることで、ImageNetにおいて88.3%のゼロショットと91.1%の微調整精度を達成し、これまでの最高結果をそれぞれ2%と0.1%上回ったことである。さらに、LionはJFTの事前学習計算を最大5倍削減し、拡散モデルの学習効率を2.3倍向上させ、より良いFIDスコアを達成し、言語モデリングでは最大2倍の計算削減で同等以上の性能を提供している
Lionの性質と限界を分析も行っている。符号関数を用いて計算される一様更新は、SGDや適応的手法によって生成されるものに比べて、通常、より大きなノルムをもたらすことにユーザーは注意する必要がある。したがって、Lionは、より小さな学習率lrと、有効な重み減衰強度を維持するためのより大きな非結合型重み減衰λを必要とする。詳細な指針については、セクション5を参照されたい。さらに、我々の実験では、Lionの利得はバッチサイズとともに増加し、AdamWと比較して異なるハイパーパラメータの選択に対してよりロバストであることが示された。制限としては、LionとAdamWの差は、いくつかの大規模な言語と画像-テキストデータセットでは統計的に有意ではありません。また、学習時に強い増強や小さいバッチサイズ（<64）を使用した場合、Lionの優位性は小さくなる。詳細はセクション6を参照。
わっしー課題:
バッチサイズの効果 Lionは符号演算によるノイズが加わるため、大きなバッチサイズでないと正確な方向がわからないのではないかという疑問がある。そこで、ImageNet上でViT-B/16モデルを、総トレーニングエポックを300とし、RandAugとMixupの手法を取り入れながら、様々なバッチサイズを用いて学習させたところ、図8のようになりました。図8（左）に示すように、AdamWの最適なバッチサイズは256であり、Lionの最適なバッチサイズは4,096である。このことから、Lionは確かに大きなバッチサイズを好むが、64の小さなバッチサイズでも性能は安定していることがわかる。さらに、バッチサイズを32Kに拡大し、学習ステップ数を11Kとした場合、LionはAdamWに対して2.5%の大幅な精度向上を達成し（77.9%対75.4%）、大規模バッチ学習環境での有効性を実証しています。
Lionの限界 できるだけ多くのタスクでLionを評価するように努めたが、評価は選択したタスクに限定される。視覚タスクでは、強力な補強を行った場合、LionとAdamWの性能差は縮小します。また、LionがAdamWと同等の性能を示すタスクもいくつかあり、（1）Imagen text-to-image base model、（2）文脈内学習ベンチマークより信頼できる指標である大規模内部データセットで学習した自己回帰言語モデルのパープレックス、（3）C4でのマスク言語モデリングなどがある。これらのタスクは、データセットが大規模かつ高品質であることが共通の特徴であり、その結果、オプティマイザ間の差が小さくなっています。もうひとつの潜在的な制限は、バッチサイズです。並列性を高めるためにバッチサイズを大きくすることがよくありますが、バッチサイズが小さい場合（64未満）、LionはAdamWより性能が劣る可能性があります。さらに、Lionはbfloat16での運動量追跡を必要としますが、これは巨大なモデルの学習には高価になりかねません。1つの解決策は、メモリを節約するために運動量を因数分解することです。