Google、より効率的なAI学習アルゴリズムを発見するためにAutoMLを使用

Googleの研究者は、自動機械学習（AutoML）進化アルゴリズムによって発見された、ニューラルネットワークを訓練するための最適化アルゴリズムEvoLved sIgn mOmeNtum（Lion）をオープンソース化した。Lionで学習したモデルは他のオプティマイザで学習したモデルよりも、いくつかのベンチマークで高い精度を達成でき、しかも収束に要する計算サイクルは少なくて済む。

Googleはこのアルゴリズムを発見するために、シンボリックプログラムのevolutionary searchを利用した。Lionはデファクトスタンダードな最適化アルゴリズムであるAdamよりも少ないメモリと少ない命令数で動作する。他の最適化アルゴリズムとの大きな違いは、Lionは勾配の符号だけを気にしており、重みづけに同じ大きさの更新を適用することである。Lionを使用して、チームはさまざまなベンチマークタスクに対して複数のモデルを訓練した。ImageNetの分類のために訓練されたものは、最先端の精度を達成した。Lionで訓練したViTは、ImageNetの精度を2%向上させ、同時に訓練時の計算サイクルを5倍節約した。

ディープラーニングモデルは、勾配降下アルゴリズムを使用して、ネットワークの損失関数の最小値を生成するためにネットワークの重みを反復的に更新することによって訓練される。一般的に使用されている2つのアルゴリズムは、AdamとAdafactorで、ほとんどの最先端モデルの訓練に使用されてきた。これらのアルゴリズムは以前のものに比べて大幅に改善されているが、GoogleチームはAutoMLを使って改善できないかと考えた。

Lionの開発にあたりチームは検索空間をPythonに似た命令型言語で書かれた関数の集合と定義した。関数の入力はモデルの重み、勾配、学習率で、出力は重みの更新である。検索はステートメントを追加、削除、あるいは関数のコードを変更することで行われる。候補となる関数は、小規模なMLモデルの学習に使用され、学習されたモデルの精度は、候補となる学習関数の適合性を反映している。

いくつかの実験を経て、チームはシンプルさ、メモリ効率、卓越したパフォーマンスをよりどころにして最終的な関数を定めた。AdamやAdafactorと比較してLionはよりシンプルなアルゴリズムであり、ハイパーパラメータの数も少ない。学習率によってのみスケールされる固定値で重みを更新するため、Adamよりも3倍から10倍小さい学習率と、3倍から10倍大きい学習率の減衰が必要である。

ライオン擬似コードとアダム擬似コードの比較。画像ソース： https://arxiv.org/abs/2302.06675

研究チームは、Lionを使用して、ビジョン、言語、マルチモーダルなど、さまざまなタスクドメインの既発表モデルをトレーニングし、Lionでトレーニングしたモデルの精度とトレーニング計算量をAdamやAdafactorでトレーニングしたモデルと比較した。Lionのモデルは多くの場合Adamのモデルを上回ったが、ビジョンタスクやマスクされた言語モデリングのような「データセットが大量かつ高品質」なタスクでは、同様のパフォーマンスを示するケースもあった。一方、小さなバッチサイズではLionはおそらくAdamと同じ性能を発揮すると著者は指摘している。

研究チームの複数のメンバーが、Twitterで研究成果に関するユーザーの質問に回答した。あるコメントでは、なぜLionとAdafactorを比較する際に、Adamではなく一部のモデルで比較したのか、という質問があった。共同研究者のChen Liangはこう答えている。

[ほとんどの言語ベンチマーク（smaller scale LM, masked LM, finetuning, etc）において、AdamWと比較した。1B以上のパラメータを持つ言語モデル]では、ベースラインのデフォルトであり、より少ないコストでAdamWと同様の性能を発揮するAdafactorと比較しました。

PytorchやTensorflowなど、いくつかの一般的なニューラルネットワーク学習フレームワークに対するLionのPython implementationsは、GitHubで公開されている。