워밍업을 이론으로 풀다: 적응형 스케줄러와 노름 제한 최적화기

초록

본 논문은 Muon·Lion 등 노름 제한 최적화기에 대해 학습률 워밍업을 이론적으로 정당화하고, 서브옵티멀리티에 따라 지역 곡률이 감소한다는 일반화된 스무스니스 가정을 제시한다. 이를 기반으로 워밍업‑디케이 형태의 학습률 스케줄을 도출하고, 자동으로 워밍업 길이를 결정하는 실용적인 스케줄러를 설계한다. LLaMA 기반 대규모 언어 모델 사전학습 실험에서 기존 수동 튜닝 방식보다 일관되게 동등하거나 우수한 성능을 보이며, 추가 하이퍼파라미터 탐색이 필요하지 않음을 입증한다.

상세 분석

논문은 먼저 노름 제한 최적화기(LMO 기반)의 업데이트 규칙을 일반화된 형태로 정리한다. 여기서 핵심은 학습률 ηₜ와 함께 사용되는 노름 ‖·‖가 유클리드가 아니어도 적용 가능하다는 점이다. 기존 이론은 L‑스무스니스(‖∇f(x)−∇f(y)‖≤L‖x−y‖)만을 가정했으며, 이는 학습 초기에 큰 학습률을 사용할 수 없게 만든다. 저자들은 “지역 곡률이 서브옵티멀리티 갭에 비례해 감소한다”는 (ρ, K₀, K₁, K_ρ)‑스무스니스 가정을 도입한다. 구체적으로 K(x)=K₀+K₁·Δ+K_ρ·Δ^ρ (Δ=f(x)−f★) 로 정의하고, ∇f의 변동성을 이 K(x)와 연결한다. 실험적으로 Lion, Muon, normSGD 등 여러 옵티마이저에 대해 Δ와 Kₜ=‖∇f(xₜ₊₁)−∇f(xₜ)‖·‖xₜ₊₁−xₜ‖⁻¹ 사이에 2차 함수 형태의 강한 상관관계를 관측했으며, 이는 K_ρ>0 가 실제 학습 과정에서 필수임을 뒷받침한다.

이 가정 하에 Theorem 1을 증명한다. 학습률을 ηₜ=Δₜ/D·K(xₜ) 로 설정하면, Δₜ는 단조 감소하고 ηₜ는 초기에는 Δₜ가 클 때 증가(워밍업)하다가 Δₜ가 특정 임계값 Δ′= (K₀/(K_ρ(ρ−1)))^{1/ρ} 이하가 되면 감소(디케이)한다. 따라서 워밍업‑디케이 스케줄이 증명에 의해 자연스럽게 도출된다. 또한 K(xₜ) 가 시간에 따라 감소하므로, 초기 큰 ηₜ 를 사용해 빠른 수렴을 촉진하고, 후반에 작은 ηₜ 로 안정성을 확보한다.

Theorem 2에서는 가중치 디케이 λ를 포함한 업데이트 xₜ₊₁=(1−ληₜ)xₜ+ηₜ·LMO(gₜ) 를 분석한다. 여기서는 boundedness 가정이 필요 없으며, λ가 충분히 작을 경우 동일한 워밍업‑디케이 형태의 ηₜ 가 수렴률 O(1/T) 를 유지한다.

이론을 바탕으로 저자들은 실용적인 스케줄러를 설계한다. 주요 입력은 초기 학습률 상한, 최소 학습률, 그리고 최대 워밍업 스텝 수와 같은 기존 하이퍼파라미터뿐이다. 알고리즘은 현재 손실 Δₜ 를 추정하고, K(xₜ) 를 실시간으로 업데이트해 ηₜ 를 계산한다. 워밍업 종료 시점은 ηₜ 가 최대값에 도달하는 순간을 자동 감지함으로써, 별도의 손실 기반 튜닝 없이도 최적의 워밍업 길이를 찾는다.

실험에서는 LLaMA‑7B, LLaMA‑13B 등 다양한 규모의 모델을 Muon, Lion, normSGD 로 사전학습했다. 비교 대상은 (i) 고정 학습률, (ii) 선형 워밍업+코사인 디케이, (iii) 수동으로 최적화된 워밍업 길이이다. 결과는 적응형 스케줄러가 대부분의 설정에서 동일하거나 더 높은 최종 퍼플렉시티와 토큰 정확도를 달성했으며, 특히 학습 초기 손실 감소 속도가 크게 향상되었다. 또한 ablation 실험을 통해 f★ 추정 오차가 10% 수준까지 허용돼도 안정적인 동작을 보였으며, 가중치 디케이와 결합했을 때 일반화 성능이 추가로 개선됨을 확인했다.

마지막으로 저자들은 한계점도 언급한다. 현재 가정은 정확한 f★ 추정에 어느 정도 의존하며, 매우 작은 배치 사이즈나 비정규화된 손실에서는 K(x) 추정이 불안정할 수 있다. 또한 이론은 결정적 그라디언트에 기반하므로, 고차원 잡음이 큰 상황에서는 추가적인 확률적 분석이 필요하다. 그럼에도 불구하고, 워밍업을 경험적 트릭이 아닌 수학적 최적화 결과로 재해석한 본 연구는 대규모 언어 모델 훈련에서 학습률 스케줄링을 자동화하는 실용적 길을 제시한다.