μLO: Meta-Generalização Eficiente em Computação de Otimizadores Aprendidos

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

May 31, 2024

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

cs.AI

Resumo

Otimizadores aprendidos (LOs, na sigla em inglês) podem reduzir significativamente o tempo de treinamento em relógio de redes neurais, diminuindo substancialmente os custos de treinamento. No entanto, eles frequentemente sofrem com má meta-generalização, especialmente ao treinar redes maiores do que aquelas vistas durante o meta-treinamento. Para resolver isso, utilizamos a Parametrização de Atualização Máxima (muP, na sigla em inglês), recentemente proposta, que permite a generalização zero-shot de hiperparâmetros do otimizador de modelos menores para maiores. Estendemos a teoria muP para otimizadores aprendidos, tratando o problema de meta-treinamento como a busca pelo otimizador aprendido sob muP. Nossa avaliação mostra que LOs meta-treinados com muP melhoram substancialmente a meta-generalização em comparação com LOs treinados sob parametrização padrão (SP, na sigla em inglês). Notavelmente, quando aplicado a modelos de grande largura, nosso melhor muLO, treinado por 103 horas de GPU, iguala ou supera o desempenho de VeLO, o maior otimizador aprendido publicamente disponível, meta-treinado com 4000 meses-TPU de computação. Além disso, muLOs demonstram melhor generalização do que suas contrapartes SP para redes mais profundas e para horizontes de treinamento muito mais longos (25 vezes mais longos) do que aqueles vistos durante o meta-treinamento.

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.