기술 요약: 무작위 특징 모델에 대한 최적 학습률 스케줄 및 스케일링 법칙 이론

문제 제기

딥러닝 모델을 훈련하려면 학습률 (LR), 배치 크기, 훈련 기간과 같은 하이퍼파라미터를 선택해야 합니다. 경험적인 '시행착오'가 일반적이지만, 최적 하이퍼파라미터가 훈련 기간 $T$에 따라 어떻게 스케일링되는지에 대한 이론적 이해는 부족합니다. 모델 크기를 변경할 때 하이퍼파라미터를 전이하기 위한 기존 전략들은 종종 서로 다른 훈련 기간 간에는 전이가 실패합니다. 또한, 다양한 LR 스케줄 (예: 선형 감쇠, 워밍업 - 안정 - 감쇠, 코사인 어닐링) 이 실제로 사용되고 있지만, 주어진 작업에 대해 어떤 스케줄이 최적인지를 결정하는 이론적 조건은 여전히 불분명합니다. 본 논문은 훈련 기간과 학습 작업의 특정 구조에 의존하는 최적 LR 스케줄을 유도하기 위한 원칙적인 프레임워크의 필요성을 해결합니다.

방법론

저자들은 최적 하이퍼파라미터 스케줄을 유도하기 위해 확률적 경사 하강법 (SGD) 으로 훈련된 멱법칙 무작위 특징 모델을 사용합니다. 이 모델은 $y(x) = w^* \cdot \psi(x) + \sigma_0 z$인 타겟 함수를 갖는 입력 데이터 $x$를 가정하며, 여기서 특징 $\psi(x)$는 멱법칙 스펙트럼 감쇠 ($\lambda_k \sim k^{-b}$) 를 보이고, 교사 가중치는 $(w^*_k)^2 \lambda_k \sim k^{-a}$로 스케일링됩니다.

핵심 방법론은 최적 제어 이론을 포함합니다:

1. 동역학 공식화: 기대 테스트 손실 $L_t$는 학습률 스케줄 $\{\eta_t\}$과 배치 크기 스케줄 $\{m_t\}$의 함수로 유도됩니다. 손실 동역학은 모드 오차 $c_{t,k}$에 대한 선형 재귀 관계에 의해 지배됩니다.
2. 최적화 목표: 고정된 기간 $T$에서 최종 테스트 손실 $L_T$를 최소화하는 것입니다. 저자들은 세 가지 최적화 문제를 공식화합니다:
   • 최적 상수 LR.
   • 최적 LR 스케줄 (시간에 따라 변하는 $\eta_t$).
   • 결합된 최적 LR 및 배치 크기 스케줄.
3. 해법 접근:
   • 수치적: 유한한 $T$에 대한 이산 시간 최적 제어 문제를 해결하기 위해 단일 슈팅 방법 (CasADi 에서 구현됨) 을 사용합니다.
   • 해석적: 노에테르 정리와 손실 함수의 점근 분석을 사용하여 큰 $T$에 대한 연속 시간 근사를 유도합니다. 여기에는 통합 학습률 $\chi(t) = \int_t^T \eta(s) ds$의 함수를 최소화하는 작업이 포함됩니다.

주요 기여 및 결과

1. 두 가지 구별되는 위상의 식별

분석은 스펙트럼 지수 $a$(작업 난이도) 와 $b$(특징 복잡도) 간의 관계에 의해 결정되는 최적 학습률 스케줄을 위한 두 가지 구별되는 체계를 보여줍니다:

• 쉬운 위상 ($a > b$):

  • 최적 스케줄은 다항식 감쇠 형태를 따릅니다: $\eta^*_T(t) \sim T^{b/a - 1} f(t/T)$, 여기서 $f(z) \propto (1-z)^{2b-1}$입니다.
  • 학습률은 훈련 궤적 전반에 걸쳐 편향과 분산 감소를 균형 있게 맞추기 위해 점차 감쇠합니다.
  • 손실 스케일링 지수는 고정된 배치 크기에 대한 미니맥스 최적 속도와 일치하는 $\zeta_T = (a-1)/a$입니다.

• 어려운 위상 ($a < b$):

  • 최적 스케줄은 **워밍업 - 안정 - 감쇠 (WSD)**와 유사합니다. 대부분의 훈련 단계 동안 일정한 최대 학습률 ($\eta_{max}$) 을 유지한 후, 전체 시간의 소멸하는 비율 ($1 - t_s/T \sim T^{-(b-a)/(2b-1)}$) 에 걸쳐 급격한 어닐링 단계를 거칩니다.
  • 이 전략은 최대 상수 LR 을 통한 편향 억제를 우선시하며, 분산 감소를 최종 단계로 압축합니다.
  • 손실 스케일링 지수는 $\zeta_T = (a-1)/b$로 개선되어 상수 LR 전략보다 우수한 성능을 보입니다.

2. 벤치마크 대비 우월성

유도된 최적 스케줄은 표준 벤치마크를 엄격하게 능가합니다:

• 상수 LR: 최적 스케줄은 두 위상 모두에서 더 나은 스케일링 지수를 달성합니다.
• 멱법칙 스케줄: 멱법칙 스케줄의 지수 ($\eta \sim T^{-\xi} t^{-\delta}$) 를 최적화하더라도, 점근적으로 최적 상수 LR 기준과 동일한 스케일링 행동을 회복하여 제어 이론을 통해 유도된 기간 의존적 최적 스케줄의 성능을 따라가지 못합니다.

3. 학습률 및 배치 크기의 결합 최적화

저자들은 고정된 데이터 예산 $B_{tot}$ 하에서 $\eta(t)$와 $m(t)$를 결합하여 최적화하는 분석을 확장합니다:

• 쉬운 위상: "배치 램프"(시간에 따라 증가하는 배치 크기) 를 일정한 최대 학습률과 결합하면 월클록 시간을 단축합니다. 손실은 $T^{-(a-1)/b}$로 스케일링되어 고정된 배치 크기에 비해 월클록 지수를 개선하지만, 샘플 예산 지수는 $(a-1)/a$로 유지됩니다.
• 어려운 위상: 결합 최적화는 고정 배치 최적과 동일한 스케일링 지수를 산출하며, 배치 크기는 유계로 유지됩니다.

4. 최적 모멘텀 스케줄

모멘텀이 필요한 작업의 경우, 저자들은 $\eta(t)$와 함께 모멘텀 매개변수 $\beta(t)$를 최적화합니다.

• 쉬운 위상에서 최적 모멘텀 스케줄은 최적 SGD 와 유사한 스케일링 법칙을 산출합니다.
• 어려운 위상에서 모멘텀 스케줄을 최적화하면 SGD 단독으로 달성 가능한 것 이상의 스케일링 지수 개선을 제공합니다. 최적 모멘텀 동역학은 초기 과도기 감소 후 $t/T \to 1$로 접근함에 따라 증가하는 모습을 보입니다.

5. 심층 네트워크에서의 경험적 검증

이론적 위상은 딥러닝 벤치마크에서 검증되었습니다:

• CIFAR-5M 의 ResNet: 어려운 위상 행동을 보입니다. 어닐링 스케줄에 대한 최적 기본 학습률은 기간 전반에 걸쳐 거의 일정하며, 이는 WSD 프로필과 일치합니다.
• C4 의 GPT-2 스타일 트랜스포머: 쉬운 위상 행동을 보입니다. 최적 기본 학습률은 어닐링 하에서도 기간 $T$에 따라 이동 (감소) 하며, 이는 다항식 감쇠 프로필과 일치합니다.

중요성 및 주장

본 논문은 다양한 작업과 훈련 기간에 걸쳐 학습률 스케줄이 왜 변하는지 이해하기 위한 이론적 기초를 제공한다고 주장합니다. 해결 가능한 무작위 특징 모델에 최적 제어 이론을 적용함으로써, 저자들은 다음을 입증합니다:

1. 기간 의존성: 최적 학습률은 훈련 기간 간에 자동으로 전이되지 않습니다. 스케줄은 $T$에 기반하여 명시적으로 설계되어야 합니다.
2. 작업 구조: 작업의 "난이도"(스펙트럼 지수 $a$와 $b$로 특징화됨) 는 작업이 쉬운 위상인지 어려운 위상인지 결정하며, 이는 최적 스케줄 형태 (다항식 감쇠 대 WSD) 를 결정합니다.
3. 효율성: 기간 의존적 스케줄과 배치 크기의 결합 최적화는 표준 휴리스틱에 비해 스케일링 법칙을 크게 개선하고 월클록 훈련 시간을 단축할 수 있습니다.

저자들은 한계를 인정하며, 현재 이론은 느린/커널 영역의 선형 모델에 적용되며, 어려운 위상에서 모멘텀에 대한 구체적인 개선된 지수를 해석적으로 유도하거나 Adam 과 같은 정교한 옵티마이저나 전치 조건 방법을 고려하지는 않는다고 명시합니다. 그러나 이 작업은 WSD 와 같은 스케줄의 효과성에 대한 원칙적인 설명을 제공하며, 기간 전반에 걸친 기본 학습률의 경험적 튜닝이 모델과 작업의 근본적인 스펙트럼 특성을 고려해야 함을 시사합니다.