The Road Less Scheduled (NeurIPS 2024, Oral paper) 내용 정리

배경 및 문제 제기

일반적으로 stochastic optimization 문제, 특히 deep learning 모델 학습에서는 학습 과정을 안정적으로 수렴시키기 위해 learning rate schedule을 사용한다. 이러한 schedule은 특정 시간에 따른 learning rate 감소나 형태 변화를 명시적으로 설정하여, training 중 성능을 향상시키는 전략이다. 하지만 schedule을 사용하려면 미리 학습 종료 시점(T)을 가정하거나 전체 training 횟수(epoch)를 알고 있어야 하고, hyper-parameter tuning이 복잡해진다.

한편, 기존 이론들은 non-smooth convex setting에서 단순한 SGD를 평균화한 Polyak-Ruppert averaging(PA)나 Primal averaging(PA와 다른 형태의 averaging) 같은 averaging 기술을 통해 최적의 worst-case 수렴률을 얻을 수 있음을 보인다. 그러나 실제 deep learning 실험에서는 PA나 Primal averaging이 schedule 기반 기법에 비해 성능이 저조한 것으로 잘 알려져 있다. 즉, 이론적으로는 averaging이 완벽해 보이지만, 실무적으로는 schedule 기반 접근이 더 낫다는 모순점이 존재했다.

이 논문은 위와 같은 이론-실무 간 격차를 해소하고자 한다. 논문은 "Schedule-Free" 접근을 제안하며, schedule 없이도 어떤 시간에든 최적의 성능-시간 트레이드오프 곡선 상에 근접하는 결과를 얻을 수 있음을 보여준다. 또한 이 방법은 학습 종료 시점을 알 필요가 없고, additional hyper-parameter 없이도 momentum만으로 schedule 기반 방법과 동등하거나 더 우수한 성능을 낸다.

---

주요 아이디어

Schedule-Free 접근은 Polyak-Ruppert averaging, Primal averaging, 그리고 linear decay schedule을 일반화하는 새로운 형태의 averaging 방법을 제안한다. 이론적으로는, 특정 weight sequence를 통해 online-to-batch conversion을 할 수 있으며, 이를 적절히 조합하면 schedule을 명시적으로 사용하지 않고도 schedule과 동일한 역할을 수행할 수 있음을 보인다.

핵심 수식은 다음과 같다. 기본적으로 다음과 같은 변수를 정의한다.

• z: base parameter sequence
• x: evaluation sequence (return될 parameter의 평균적 추정)
• y: gradient 위치 점(sequence), 즉 실제 gradient를 계산하는 지점
• β: momentum parameter

Schedule-Free SGD의 형태는 다음과 같이 정리된다.

1. y_t = (1 - β) z_t + β x_t
2. z_{t+1} = z_t - γ g_t, where g_t = ∇f(y_t, ζ_t)
3. x_{t+1} = (1 - c_{t+1}) x_t + c_{t+1} z_{t+1}, 여기서 c_{t} = 1/(t+1)

이 방식은 β=0일 때 Polyak-Ruppert averaging을, β=1일 때 Primal averaging을 재현한다. β를 0과 1 사이의 값으로 두면, Polyak averaging과 Primal averaging 사이를 연속적으로 연결하는 스펙트럼 상에 놓이게 된다.

이 논문은 이를 더 일반화하여 arbitrary online algorithms에 대해서도 적용 가능한 형태로 만들고, schedule-free 방식을 통해 linear decay나 cosine schedule 등을 이론적으로 대체할 수 있음을 보인다.

---

이론적 결과

논문은 Theorem을 통해 다음을 보인다.

1. Schedule-Free SGD는 non-smooth convex 환경에서 optimal O(1/√T) 수렴률을 달성하며, 이는 momentum parameter를 어떻게 선택하든 최적을 유지한다.
2. Schedule-Free 방법은 online-to-batch 변환 정리에 기반하여 worst-case optimal성을 보장한다. 즉, 어떤 step size나 momentum을 사용하더라도 기존 schedule 없이도 이론적 관점에서 최적의 수렴률을 얻을 수 있다.
3. large learning rate 사용 조건: 실제 deep learning에서 흔히 쓰이는 큰 learning rate를 이론적으로도 정당화할 수 있는 조건을 제시한다. 즉, 특정 조건 하에서 γ = D/G와 같은 큰 step size도 optimal rate를 크게 훼손하지 않는다는 점을 보인다.

이러한 이론적 결과는 schedule을 explicitly 설정하지 않고도 training 과정 전반에 걸쳐 Pareto frontier(성능-학습시간 trade-off 상의 최적 곡선)에 근접하게 학습할 수 있음을 의미한다.

---

실험 결과

논문은 다양한 문제에서 Schedule-Free 알고리즘을 검증한다. 여기에는 logistic regression과 같은 작은 convex 문제부터, ImageNet 분류, IWSLT14, CIFAR-10, CIFAR-100, SVHN, MAE fine-tuning, DLRM, 그리고 대규모 language model training(OpenWebText GPT-2)까지 폭넓은 범위의 실험이 포함된다.

• CIFAR-10, CIFAR-100, SVHN에서는 Schedule-Free SGD가 기존 cosine schedule 대비 test accuracy면에서 탁월한 성능을 보인다.
• ImageNet ResNet-50 training에서는 Schedule-Free SGD가 기존 schedule 기반 방법과 대등하거나 약간 우세한 결과를 얻는다.
• 대규모 LLM(OpenWebText GPT-2) 실험에서도 Schedule-Free AdamW가 cosine schedule을 능가하는 성능을 보였다.
• Logistic regression 등 convex setting에서도 linear decay schedule 이상의 성능을 발휘한다.

또한 Schedule-Free 방법은 MLCommons 2024 AlgoPerf Algorithmic Efficiency Challenge의 Self-Tuning track에서 우승을 거두어, 다양한 조건 하에서 적은 hyper-parameter tuning으로도 우수한 성능을 낼 수 있음을 입증했다.

---

Implementation 및 Practical Considerations

실제로 Schedule-Free 알고리즘을 구현할 때, 기존 momentum 기반 optimizer와 큰 차이가 없다. AdamW 기반 Schedule-Free 구현 예를 들어보면, z, x sequence 두 개만 tracking하면 되고, y는 매 step에서 β를 활용해 x와 z의 선형 조합으로 얻을 수 있다. 메모리 사용량이나 계산 복잡도 측면에서 기존 방식과 거의 동일하다.

BatchNorm과 같은 running statistics를 사용하는 layer의 경우, x sequence에 맞추어 batch statistics를 업데이트하거나 재계산하는 절차가 추가로 필요할 수 있다. 이 부분은 BatchNorm의 특수성 때문에 Schedule-Free 접근법과의 결합 시 주의해야 할 사항이다.

---

결론 및 향후 과제

이 논문은 schedule을 쓰지 않고, 단순히 momentum 기반 averaging 기법만으로 schedule이 주는 모든 혜택을 얻을 수 있음을 보였다. Schedule-Free 방법은 T를 사전에 모르거나, schedule tuning에 투자하는 비용을 줄이고자 할 때 특히 유용하다.

아직 완전한 이론적 이해는 남아 있지만, quadratic 근사나 large learning rate 사용 시의 수렴성에 대한 조건이 점차 밝혀지고 있다. 향후 더 정교한 문제 설정이나 non-convex optimization scenario, adaptive preconditioning 기법과의 결합 등이 연구될 수 있다.

---

마무리

"The Road Less Scheduled" 논문은 schedule 없이도 worst-case optimal rate를 달성하는 동시에, 실제 deep learning 실험에서도 schedule기반 방법과 대등하거나 능가하는 성능을 내는 방법론을 제안한다. 이 접근은 online-to-batch conversion을 새로운 방식으로 해석한 결과이며, 기존 averaging 기법 및 schedule과의 관계를 명확히 보여줌으로써 optimization 이론과 실무 사이의 간극을 줄이는 중요한 성과라 할 수 있다.

 

덧붙여서..

이 논문은 우수성을 인정받아 NeurIPS 2024 Oral paper로 선정되었다. 포스터 세션에서 이 논문의 저자와 이야기할 기회가 있었는데, optimizer의 코드가 오픈 소스로 구현이 되어 있고, 그 구현마저 간단하다는 점을 매우 강조하였다. 아래 논문의 github를 첨부한다.

https://github.com/facebookresearch/schedule_free

GitHub - facebookresearch/schedule_free: Schedule-Free Optimization in PyTorch