μ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

Аннотация

Обученные оптимизаторы (LO) могут значительно сократить время обучения нейронных сетей по настенным часам, существенно снижая затраты на обучение. Однако они часто страдают от плохой мета-обобщаемости, особенно при обучении сетей большего размера, чем те, которые виделись во время мета-обучения. Для решения этой проблемы мы используем недавно предложенную Максимальную Параметризацию Обновлений (muP), которая позволяет обобщать гиперпараметры оптимизатора из меньших моделей на большие без предварительного обучения. Мы расширяем теорию muP на обученные оптимизаторы, рассматривая проблему мета-обучения как поиск обученного оптимизатора в рамках muP. Наша оценка показывает, что LO, обученные мета-методом muP, существенно улучшают мета-обобщаемость по сравнению с LO, обученными в рамках стандартной параметризации (SP). Заметно, что, примененный к моделям большой ширины, наш лучший muLO, обученный за 103 часа на GPU, соответствует или превосходит производительность VeLO, самого крупного общедоступного обученного оптимизатора, мета-обученного на 4000 месяцев TPU-вычислений. Более того, muLO демонстрируют лучшую обобщаемость по сравнению с их 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.