μLO: Rechen-effiziente Meta-Generalisierung von gelernten Optimierern

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

May 31, 2024

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

cs.AI

Zusammenfassung

Gelernte Optimierer (LOs) können die Trainingszeit von neuronalen Netzwerken signifikant reduzieren und somit die Trainingskosten erheblich senken. Allerdings leiden sie oft unter einer schlechten Meta-Generalisierung, insbesondere bei der Schulung von Netzwerken, die größer sind als die während des Meta-Trainings gesehenen. Um dies zu lösen, verwenden wir die kürzlich vorgeschlagene Maximal Update Parametrisierung (muP), die eine Generalisierung der Optimierer-Hyperparameter von kleineren auf größere Modelle ohne vorheriges Training ermöglicht. Wir erweitern die muP-Theorie auf gelernte Optimierer und behandeln das Meta-Trainingsproblem als die Suche nach dem gelernten Optimierer unter muP. Unsere Evaluation zeigt, dass LOs, die mit muP meta-trainiert wurden, die Meta-Generalisierung im Vergleich zu LOs, die unter Standardparametrisierung (SP) trainiert wurden, erheblich verbessern. Insbesondere erreicht unser bester muLO, der für 103 GPU-Stunden trainiert wurde, bei Anwendung auf Modelle mit großer Breite oder übertrifft die Leistung von VeLO, dem größten öffentlich verfügbaren gelernten Optimierer, der mit 4000 TPU-Monaten an Rechenleistung meta-trainiert wurde. Darüber hinaus zeigen muLOs eine bessere Generalisierung als ihre SP-Gegenstücke für tiefere Netzwerke und für deutlich längere Trainingszeiträume (25-mal länger) als die während des Meta-Trainings beobachteten.

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.