[논문리뷰] Fantastic Pretraining Optimizers and Where to Find Them

링크: 논문 PDF로 바로 열기

저자: Kaiyue Wen, David Hall, Tengyu Ma, Percy Liang

핵심 연구 목표

본 논문은 언어 모델 사전 훈련에서 AdamW가 지배적인 옵티마이저임에도 불구하고, 새로운 옵티마이저들이 주장하는 1.4배에서 2배의 학습 속도 향상이 실제로는 널리 채택되지 않는 이유를 규명하고자 합니다. 저자들은 이러한 불일치가 (i) 불균등한 하이퍼파라미터 튜닝과 (ii) 제한적이거나 오해의 소지가 있는 평가 설정에서 비롯되었다고 주장하며, 이를 해결하기 위한 체계적인 연구를 수행합니다.

핵심 방법론

저자들은 0.1B부터 1.2B 매개변수에 이르는 네 가지 모델 규모와 1배에서 8배의 Chinchilla 최적 데이터-모델 비율에 걸쳐 10가지 딥러닝 옵티마이저를 체계적으로 연구합니다. 세 단계의 하이퍼파라미터 튜닝 프레임워크를 사용하여 최적의 설정을 찾고, 이를 통해 얻은 하이퍼파라미터 스케일링 법칙을 1.2B 매개변수 모델까지 외삽합니다. 주요 평가 지표는 주어진 손실에 도달하는 데 필요한 토큰 수이며, JAX 및 TPU v5 환경에서 Llama 2 아키텍처를 기반으로 실험을 진행합니다.

주요 결과

새로운 옵티마이저들이 주장하는 2배 속도 향상은 불충분하게 튜닝된 AdamW 베이스라인에서 비롯된 경우가 많으며, 잘 튜닝된 AdamW 대비 실제 속도 향상은 최대 1.4배에 불과하고 모델 규모가 커질수록 감소하여 1.2B 모델에서는 1.1배에 그칩니다. Muon, Soap, Kron과 같은 매트릭스 기반 옵티마이저는 AdamW보다 일관되게 우수하지만, 속도 향상은 모델 규모에 반비례합니다. 또한, 데이터-모델 비율이 8배 이상으로 증가하면 Muon은 Kron과 Soap에 의해 추월당하며, 학습 초기 단계의 손실 곡선만으로 옵티마이저를 평가하는 것은 오해를 유발할 수 있습니다.

AI 실무자를 위한 시사점

AI 실무자들은 새로운 옵티마이저를 평가할 때 AdamW를 포함한 모든 옵티마이저에 대해 엄격하고 공정한 하이퍼파라미터 튜닝이 필수적임을 인지해야 합니다. 작은 모델이나 낮은 데이터-모델 비율에서 나타난 성능 향상이 대규모 모델 및 데이터 조건에서는 크게 감소할 수 있으므로, 다양한 스케일링 레짐에서의 검증이 중요합니다. 또한, 학습 초기 손실 곡선에 현혹되지 않고 최종 훈련 예산에서의 성능을 기준으로 옵티마이저를 선택해야 합니다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.