딥러닝 손실과 학습률의 스케일링 법칙: 볼록성 지배 현상

초록

본 논문은 딥러닝 훈련 초기에 손실 함수가 약한 볼록성(weak convexity)을 보이며, 이때 손실의 최종값을 학습률 스케줄에 대한 명시적 상한식으로 예측할 수 있음을 보인다. 이를 기반으로 학습률 피크값을 (O(1/\sqrt{T})) 로 스케일링하면 손실이 (O(1/\sqrt{T})) 수렴하고, 다양한 모델·옵티마이저·스케줄에 대해 80배·70배까지 외삽이 가능함을 실증한다.

상세 분석

논문은 먼저 기존 연구에서 관찰된 “볼록‑유사” 현상을 정리하고, 이를 수학적으로 다루기 위해 두 가지 핵심 가정을 제시한다. 첫 번째는 손실 함수가 초기 학습 단계 이후 약한 볼록성(weak convexity)을 만족한다는 가정이며, 이는 Hessian 스펙트럼이 양의 고유값에 의해 지배되는 현상과 일치한다. 두 번째는 미니배치 그래디언트의 2‑노름이 유계 (G) 로 제한된다는 전제다. 이 두 가정 하에서 저자들은 전통적인 convex 최적화 이론(특히 Defazio et al., 2023의 Corollary 12)을 재구성하여, 임의의 학습률 시퀀스 ({\eta_t}) 에 대해 최종 손실 (L_T) 의 상한을 식 (2.4) 로 도출한다.

식 (2.4)는 평균 학습률 (\sum_{t=1}^T \eta_t) 과 제곱 평균 (\sum_{t=1}^T \eta_t^2) 을 포함하며, 여기서 피크 학습률 (\eta_{\text{peak}}) 을 (1/\sqrt{T}) 로 스케일링하면 (O(1/\sqrt{T})) 수렴을 보장한다. 특히, “Qualified Schedule”이라는 개념을 도입해, 학습률 스케줄 (s_t(T)) 가 연속적 적분 형태의 조건 (2.5)를 만족하면 피크 학습률을 (1/\sqrt{T}) 로 두어도 최적 수렴률을 얻을 수 있음을 증명한다. 이론적으로는 선형 감소, 코사인 감소, Warm‑up‑Stable‑Decay (WSD) 스케줄이 조건을 만족하고, 상수 혹은 제곱근 역수 스케줄은 실패한다는 결과를 제시한다.

실험 부분에서는 ResNet, ViT, GPT‑2 등 다양한 아키텍처와 ImageNet, OpenWebText, Cauldron 등 여러 데이터셋에 대해 SGD, AdamW, Muon, LoRA 등 다양한 옵티마이저를 적용하였다. 각 실험에서 손실 곡선과 학습률 곡선을 동시에 시각화하고, 식 (3.1) (2.4의 일반화 형태) 로 예측한 손실과 실제 손실의 R²가 0.95 이상임을 확인했다. 특히, 학습 초기에 약 5~10% 정도의 “warm‑up” 구간을 제외하면 예측 정확도가 급격히 상승한다는 점이 강조된다.

스케일링 법칙 측면에서는 피크 학습률 (\eta_{\text{peak}} = \eta_{\text{ref}}/\sqrt{T}) 와 손실 상한 (L_{\text{SGD-last}}(T) - L^* \approx Q(\eta_{\text{ref}})/\sqrt{T}) 를 데이터‑드리븐 방식으로 피팅하였다. 여기서 (Q(\eta_{\text{ref}})=q_2/\eta_{\text{ref}} + \eta_{\text{ref}} q_1) 는 두 상수 (q_1, q_2) 에 의해 결정되며, 최적 (\eta_{\text{ref}} = \sqrt{q_1/q_2}) 를 선택하면 손실 상한이 최소화된다. 이 모델을 이용해 훈련 시간 (T) 와 모델 파라미터 수 (N) 에 대해 각각 80배·70배까지 외삽했을 때도 실제 손실과 매우 근접함을 보였다.

결론적으로, 논문은 딥러닝 최적화가 전통적인 비볼록 문제와는 달리 초기 몇 에포크 이후 “볼록 지배” 구간에 진입한다는 새로운 관점을 제시하고, 이를 통해 학습률 스케줄 설계와 스케일링 법칙을 이론‑실험적으로 일관되게 설명한다. 또한, 학습률을 (1/\sqrt{T}) 스케일링하는 것이 거의 보편적인 최적 전략임을 강조함으로써, 대규모 모델 훈련 시 실용적인 가이드라인을 제공한다.