μLO: 計算効率に優れた学習済みオプティマイザのメタ汎化

μLO: Compute-Efficient Meta-Generalization of Learned Optimizers

May 31, 2024

著者: Benjamin Thérien, Charles-Étienne Joseph, Boris Knyazev, Edouard Oyallon, Irina Rish, Eugene Belilovsky

cs.AI

要旨

学習済みオプティマイザ（LOs）は、ニューラルネットワークの実時間トレーニング時間を大幅に短縮し、トレーニングコストを大幅に削減することができます。しかし、特にメタトレーニング中に見たものよりも大きなネットワークをトレーニングする場合、メタ一般化が不十分であることがよくあります。これを解決するために、最近提案されたMaximal Update Parametrization（muP）を使用します。muPは、より小さなモデルからより大きなモデルへのオプティマイザハイパーパラメータのゼロショット一般化を可能にします。我々はmuP理論を学習済みオプティマイザに拡張し、メタトレーニング問題をmuPの下で学習済みオプティマイザを見つける問題として扱います。評価の結果、muPでメタトレーニングされたLOsは、標準パラメータ化（SP）の下でトレーニングされたLOsと比較して、メタ一般化が大幅に向上することが示されました。特に、大規模な幅のモデルに適用した場合、103 GPU時間でトレーニングされた我々の最良のmuLOは、4000 TPU月の計算リソースでメタトレーニングされた公開されている最大の学習済みオプティマイザであるVeLOの性能に匹敵またはそれを上回りました。さらに、muLOsは、SPの対応するものと比較して、より深いネットワークやメタトレーニング中に見たものよりもはるかに長いトレーニング期間（25倍長い）への一般化が優れていることを示しました。

English

Learned optimizers (LOs) can significantly reduce the wall-clock training time of neural networks, substantially reducing training costs. However, they often suffer from poor meta-generalization, especially when training networks larger than those seen during meta-training. To address this, we use the recently proposed Maximal Update Parametrization (muP), which allows zero-shot generalization of optimizer hyperparameters from smaller to larger models. We extend muP theory to learned optimizers, treating the meta-training problem as finding the learned optimizer under muP. Our evaluation shows that LOs meta-trained with muP substantially improve meta-generalization as compared to LOs trained under standard parametrization (SP). Notably, when applied to large-width models, our best muLO, trained for 103 GPU-hours, matches or exceeds the performance of VeLO, the largest publicly available learned optimizer, meta-trained with 4000 TPU-months of compute. Moreover, muLOs demonstrate better generalization than their SP counterparts to deeper networks and to much longer training horizons (25 times longer) than those seen during meta-training.