[논문 리뷰] Symbolic Discovery of Optimization Algorithms

Xiangning Chen, Liang Chen외 10명|arXiv (Cornell University)|2023. 02. 13.

Machine Learning and Data Classification인용 수 167

한 줄 요약

이 논문은 최적화 알고리즘 발견을 프로그램 탐색으로 설정하고 모멘텀 추적과 부호 기반 업데이트를 사용하는 간단하고 메모리 효율적인 옵티마이저 Lion을 식별한다. 이 옵티마이저는 시각, 언어, 확산 태스크 전반에서 성능을 향상시킨다.

ABSTRACT

We present a method to formulate algorithm discovery as program search, and apply it to discover optimization algorithms for deep neural network training. We leverage efficient search techniques to explore an infinite and sparse program space. To bridge the large generalization gap between proxy and target tasks, we also introduce program selection and simplification strategies. Our method discovers a simple and effective optimization algorithm, $ extbf{Lion}$ ($ extit{Evo$ extbf{L}$ved S$ extbf{i}$gn M$ extbf{o}$me$ extbf{n}$tum}$). It is more memory-efficient than Adam as it only keeps track of the momentum. Different from adaptive optimizers, its update has the same magnitude for each parameter calculated through the sign operation. We compare Lion with widely used optimizers, such as Adam and Adafactor, for training a variety of models on different tasks. On image classification, Lion boosts the accuracy of ViT by up to 2% on ImageNet and saves up to 5x the pre-training compute on JFT. On vision-language contrastive learning, we achieve 88.3% $ extit{zero-shot}$ and 91.1% $ extit{fine-tuning}$ accuracy on ImageNet, surpassing the previous best results by 2% and 0.1%, respectively. On diffusion models, Lion outperforms Adam by achieving a better FID score and reducing the training compute by up to 2.3x. For autoregressive, masked language modeling, and fine-tuning, Lion exhibits a similar or better performance compared to Adam. Our analysis of Lion reveals that its performance gain grows with the training batch size. It also requires a smaller learning rate than Adam due to the larger norm of the update produced by the sign function. Additionally, we examine the limitations of Lion and identify scenarios where its improvements are small or not statistically significant. Lion is also successfully deployed in production systems such as Google search ads CTR model.

연구 동기 및 목표

• handcrafted한 최적화 알고리즘을 넘어서는 새로운 옵티마이저의 발견을 동기화하여 딥 뉴럴 네트워크의 학습 효율성과 일반화를 향상시키려는 목적.
• 최적화 알고리즘 발견을 실행 가능한 프로그램 탐색 문제로 공식화하여 새로운 알고리즘 설계를 탐구.
• 무한하고 희소한 탐색 공간을 탐색하고 대규모 태스크 전반에서 일반화 가능한 해를 필터링하는 기술을 개발.
• 비전, 언어, 확산 모델 전반에서 발견된 옵티마이저의 실용적 성능을 입증.

제안 방법

• 가중치, 기울기, 보조 상태에 대해 고정된 학습 시그니처를 가진 명령형 프로그램으로 옵티마이저를 표현.
• 후보 알고리즘을 인코딩하기 위해 45개의 수학 함수 집합에서 문장들의 큰 탐색 공간을 사용.
• AdamW에서 시작하는 워밍업 및 재시작으로 공간을 탐색하는 진화적 탐색.
• 잘못된/의미적으로 동등한 프로그램을 가지치고 평가 속도를 높이기 위해 추상 실행과 캐싱을 적용.
• 점진적으로 커지는 프록시 태스크에서의 퍼널 선택 및 메타-검증을 통해 대상 태스크로 일반화하는 알고리즘을 선별.
• 중복을 제거하고 업데이트를 부호 기반 모멘텀 업데이트에 맞추어 발견된 프로그램을 Lion으로 단순화.

실험 결과

연구 질문

• RQ1 무한하고 희소한 공간에서 프로그램 탐색으로 최적화 알고리즘이 효과적으로 발견될 수 있는가?
• RQ2 자동으로 발견된 옵티마이저가 프록시 태스크에서 대규모의 최첨단 학습 설정으로 일반화되는가?
• RQ3 아키텍처와 태스크 전반에 걸쳐 발견된 옵티마이저의 특성 및 실질적 한계는 무엇인가?
• RQ4 단순한 부호 기반 업데이트와 모멘텀이 실제 학습에서 AdamW 및 Adafactor와 어떻게 비교되는가?

주요 결과

• Lion은 모멘텀을 추적하고 부호 업데이트를 사용하는 단순하고 메모리 효율적인 옵티마이저로, 특정 설정에서 이전 벤치마크보다 ImageNet 제로샷 정확도가 최대 2%포인트 높고 파인튜닝 정확도는 0.1%포인트 높다.
• Lion은 JFT에서 사전 학습 컴퓨트를 최대 5배까지 줄이고 확산 모델 학습을 약 2.3배 가속할 수 있다.
• ImageNet의 다수 모델에서 Lion은 AdamW보다 우수하다고 판단되며, 더 큰 용량의 모델과 강한 배치 크기에서 더 큰 이점을 보인다.
• 비전-언어 학습(LiT/BASIC 구성)에서 Lion은 AdamW 대비 제로샷 ImageNet 정확도를 약 1.0–1.7%포인트 향상시키고 더 나은 검색 결과를 제공한다.
• Lion의 성능 우위는 배치 크기가 커질수록 커지며, 효과적인 정규화를 유지하기 위해 더 작은 학습률과 더 큰 가중치 감소가 필요하다.
• 검색 과정은 다른 옵티마이저들 중에서 Lion을 식별하고 프록시 태스크 탐색의 메타-오버피팅이 대상 태스크로 일반화되는 경향이 있음을 보여준다.