μ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)는 신경망의 실제 학습 시간을 크게 단축시켜 학습 비용을 상당히 절감할 수 있습니다. 그러나 이들은 메타 학습 중에 접한 것보다 더 큰 네트워크를 학습할 때 특히 메타 일반화 성능이 저하되는 문제를 자주 겪습니다. 이를 해결하기 위해, 우리는 최근 제안된 최대 업데이트 매개변수화(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.