The Road Less Scheduled
스케줄 없는 길
Aaron Defazio, Xingyu Alice Yang, Harsh Mehta, et al. (2024)
학습률 스케줄(코사인 감쇠, 선형 감쇠 등)을 완전히 제거하고, 이론적으로 최적인 평균화(averaging) 기법만으로 동등하거나 우수한 성능을 달성하는 Schedule-Free 옵티마이저를 제안하여, 총 학습 스텝 수를 사전에 지정할 필요를 없앴다.
배경
현대 딥러닝에서 학습률 스케줄은 필수적인 요소로, 코사인 감쇠(cosine annealing), 선형 워밍업+감쇠, 다항 감쇠 등 다양한 스케줄이 사용된다. 그러나 대부분의 스케줄은 총 학습 스텝 수 T를 미리 알아야 하므로, 학습을 조기 종료하거나 연장하는 것이 비효율적이다. 또한 최적의 스케줄 선택 자체가 하이퍼파라미터 탐색을 요구한다. 이론적으로 최적의 수렴률을 달성하면서도 T에 의존하지 않는 '언제든 최적(anytime optimal)' 알고리즘의 개발이 숙원 과제였다.
핵심 아이디어
Schedule-Free 옵티마이저의 핵심은 학습률 감쇠의 역할을 재해석하는 것이다. 학습률 감쇠는 본질적으로 이전 반복의 파라미터를 가중 평균하는 효과가 있다. 이 관찰에 기반하여, 학습률을 일정하게 유지하되 파라미터의 적절한 가중 평균(Primal Averaging)을 사용하면 학습률 감쇠와 동일한 수렴 성질을 달성할 수 있다. 구체적으로, 학습(evaluation이 아닌 실제 그래디언트 계산)은 빠르게 이동하는 점 y_t에서 수행하고, 평가와 추론은 느리게 이동하는 평균점 x_t에서 수행한다. y_t와 x_t 사이의 보간 관계가 스케줄의 역할을 대체한다.
방법론
알고리즘은 두 개의 파라미터 시퀀스를 유지한다. 그래디언트를 계산하는 빠른 시퀀스 z_t(SGD 또는 Adam으로 업데이트)와 평가에 사용하는 느린 시퀀스 x_t이다. 핵심 업데이트 규칙은: (1) y_t = (1-β)z_t + βx_t (보간), (2) z_{t+1} = z_t - η∇f(y_t) (옵티마이저 스텝), (3) x_{t+1} = (1-1/(t+1))x_t + (1/(t+1))z_{t+1} (Polyak 평균화). 여기서 β는 보간 계수이며, 모멘텀과 유사한 역할을 한다. SGD 버전(Schedule-Free SGD)과 Adam 버전(Schedule-Free AdamW) 모두 제공된다.
주요 결과
CIFAR-10/100에서 ResNet, ViT에 대해 코사인 스케줄의 SGD/AdamW와 동등하거나 우수한 성능을 보였으며, ImageNet에서도 유사한 결과를 달성했다. GPT-2 규모의 언어 모델 사전학습에서 코사인 스케줄의 AdamW와 동일한 검증 손실에 도달하면서, 총 학습 스텝 수를 미리 지정할 필요가 없었다. MLCommons의 AlgoPerf 벤치마크에서 8개 중 5개 워크로드에서 최고 성능을 달성하여 종합 1위를 기록했다. 학습을 조기 종료하거나 예상보다 더 학습해도 성능 저하가 없는 '언제든 최적' 특성이 확인되었다.
임팩트
Schedule-Free 옵티마이저는 딥러닝 학습 파이프라인을 단순화하는 실용적 기여와 함께, 학습률 스케줄의 이론적 역할에 대한 새로운 이해를 제공했다. 총 학습 스텝 수에 대한 사전 결정이 불필요해져, 탐색적 학습과 적응적 계산 예산 할당이 가능해졌다. PyTorch에 Schedule-Free 옵티마이저가 통합되었으며, 학계와 산업계에서 기존 스케줄 기반 학습의 간편한 대체제로 채택이 진행되고 있다.