Symbolic Discovery of Optimization Algorithms

要約

タイトル：最適化アルゴリズムのシンボリックな発見
要約：
– アルゴリズム発見をプログラムの探索として定式化し、効率的な探索技術を利用して、ディープニューラルネットワークのトレーニング用の最適化アルゴリズムを発見するための方法を提案。
– プロキシタスクとターゲットタスクの大きな汎化ギャップを埋めるために、プログラム選択および簡素化戦略を導入。
– この方法によって、簡単で効果的な最適化アルゴリズムである$Lion$($EvoLved Sign Momentum$)を発見。$Adam$と比較して、メモリ効率が高く、勾配の大きさは、各パラメーターについて符号演算を通じて計算された値と同じである。
– この方法により、画像分類、自然言語処理などのタスクにおいて、既存の最適化アルゴリズムよりも優れた結果を示し、トレーニングコンピュートを大幅に削減できることが実証された。
– $Lion$の実装は公開されている。

要約(オリジナル)

We present a method to formulate algorithm discovery as program search, and apply it to discover optimization algorithms for deep neural network training. We leverage efficient search techniques to explore an infinite and sparse program space. To bridge the large generalization gap between proxy and target tasks, we also introduce program selection and simplification strategies. Our method discovers a simple and effective optimization algorithm, $\textbf{Lion}$ ($\textit{Evo$\textbf{L}$ved S$\textbf{i}$gn M$\textbf{o}$me$\textbf{n}$tum}$). It is more memory-efficient than Adam as it only keeps track of the momentum. Different from adaptive optimizers, its update has the same magnitude for each parameter calculated through the sign operation. We compare Lion with widely used optimizers, such as Adam and Adafactor, for training a variety of models on different tasks. On image classification, Lion boosts the accuracy of ViT by up to 2% on ImageNet and saves up to 5x the pre-training compute on JFT. On vision-language contrastive learning, we achieve 88.3% $\textit{zero-shot}$ and 91.1% $\textit{fine-tuning}$ accuracy on ImageNet, surpassing the previous best results by 2% and 0.1%, respectively. On diffusion models, Lion outperforms Adam by achieving a better FID score and reducing the training compute by up to 2.3x. For autoregressive, masked language modeling, and fine-tuning, Lion exhibits a similar or better performance compared to Adam. Our analysis of Lion reveals that its performance gain grows with the training batch size. It also requires a smaller learning rate than Adam due to the larger norm of the update produced by the sign function. Additionally, we examine the limitations of Lion and identify scenarios where its improvements are small or not statistically significant. The implementation of Lion is publicly available.

arxiv情報

提供元, 利用サービス

arxiv.jp, OpenAI