Celo2: 학습된 최적화의 프리 런치를 향하여

Celo2: Towards Learned Optimization Free Lunch

학습된 옵티마이저(Learned optimizers)는 Adam과 같이 수작업으로 설계된 업데이트 규칙을 대체할 수 있는 강력한 대안이지만, 훈련 분포를 넘어선 메타 일반화(meta-generalization)에 자주 실패하고 높은 메타 훈련 비용을 초래하여 실제 적용은 제한적이었습니다. 예를 들어, 이전 연구인 VeLO는 범용 옵티마이저를 메타 훈련하기 위해 4,000 TPU 개월(GPT-3 연산량의 약 10배) 규모로 메타 훈련을 확장했지만 6억 파라미터(600M) 이상의 작업으로는 일반화하지 못했습니다. 본 연구에서는 단순한 정규화 옵티마이저 아키텍처를 설계하고 메타 훈련을 증강함으로써, VeLO 연산량의 극히 일부인 단 4.5 GPU 시간만으로도 성능이 뛰어난 범용 학습형 업데이트 규칙을 메타 훈련할 수 있다는 놀라운 발견을 제시합니다. 우리가 학습한 업데이트 규칙은 메타 훈련 분포보다 100만 배(six orders of magnitude) 더 큰 10억 파라미터 규모의 사전 훈련 작업(GPT-3 XL 1.3B)에도 안정적으로 확장됩니다. 나아가, 이는 다양한 분포 외(OOD) 작업 전반에 걸쳐 강력한 성능을 보여주며, 직교화(orthogonalization), 입출력 및 은닉 가중치에 대한 개별 업데이트 규칙, 분리된 가중치 감쇠(decoupled weight decay)를 포함하는 현대적인 최적화 환경과도 호환됩니다. 종합적으로 본 연구는 실질적으로 적용 가능한 학습형 최적화 알고리즘의 기반을 마련하며, 성능 향상을 위한 더 풍부한 메타 훈련 및 데이터 큐레이션 방법론의 탐구를 가능하게 합니다.

Learned optimizers are powerful alternatives to hand-designed update rules like Adam, yet they have seen limited practical adoption since they often fail to meta-generalize beyond their training distribution and incur high meta-training cost. For instance, prior work, VeLO, scaled meta-training to 4,000 TPU months ($\sim$10$\times$ GPT-3 compute) to meta-train a general-purpose optimizer but it failed to generalize beyond 600M parameters tasks. In this work, we present a surprising finding: by crafting a simple normalized optimizer architecture and augmenting meta-training, it becomes feasible to meta-train a performant general-purpose learned update rule on a tiny fraction of VeLO compute, 4.5 GPU hours to be precise. Our learned update rule scales stably to a billion-scale pretraining task (GPT-3 XL 1.3B) which is six orders of magnitude larger than its meta-training distribution. Furthermore, it shows strong performance across diverse out-of-distribution tasks and is compatible with modern optimization harness that includes orthogonalization, distinct update rules for input-output and hidden weights, and decoupled weight decay. In all, this work paves the way for practically applicable learnable optimization algorithms, unlocking exploration of richer meta-training and data curation recipes to further improve performance.