기술 요약: LLM 학습을 위한 학습률 스케일링의 비선형성에 대하여

문제 정의

대규모 언어 모델(LLM)을 효율적으로 학습시키기 위해서는 세심한 하이퍼파라미터 튜닝이 필요하지만, 타겟 규모에서 직접 학습률(learning rate)을 스윕(sweep)하는 것은 비용이 너무 많이 들어 불가능한 경우가 많습니다. 결과적으로, 실무자들은 소규모 규모의 실행 결과를 바탕으로 대규모 모델의 최적 설정을 예측하는 **하이퍼파라미터 전이(hyperparameter transfer)**에 의존합니다. 이 분야의 지배적인 가정은 최적 학습률($\eta^*$)이 데이터 규모($D$) 및 모델 크기($N$)에 대해 로그 선형(log-linear, power-law) 스케일링 관계를 따른다는 것입니다:
$$\log \eta^*(D, N) = a \log D + b \log N + c$$
그러나 기존 연구들은 이러한 로그 선형 가정이 전역적으로 유효한지 아니면 국소적으로만 유효한지를 엄격하게 평가하지 못하는 경우가 많으며, 이러한 스케일링 법칙을 적합(fit)시키는 데 필요한 계산 비용을 자주 누락합니다. 본 논문은 이 스케일링 법칙의 유효성과 서로 다른 외삽 전략(데이터 축 vs 모델 축) 및 파라미터화(학습률 vs 유효 학습률)의 효율성을 조사합니다.

방법론

저자들은 GPT-2 스타일 아키텍처(22M에서 707M 파라미터 범위)를 사용하여 FineWeb-100B 데이터셋에서 토큰 예산 5B에서 100B 사이로 수행된 체계적인 경험적 연구를 실시했습니다.

1. 실험 설계:

   • 모델은 고정된 배치 크기(0.52M 토큰)와 warmup-steady-decay(WSD) 스케줄을 사용하는 AdamW를 사용하여 학습되었습니다.
   • 학습률은 로그 그리드($\log_2 \eta \in [-16, -7]$) 상에서 스윕되었습니다.
   • 최적 학습률 ($\eta^*$): $\log \eta$에 대한 검증 손실(validation loss)에 대한 3차 다항식을 피팅하여 추정했습니다.
   • 유효 학습률 ($\eta_{eff}$): 정규화된 가중치 공간에서의 스텝 크기로 정의됩니다: $\eta_{eff}(t) = \|\hat{w}_{t+1} - \hat{w}_t\|_2$, 여기서 $\hat{w} = w/\|w\|_2$입니다. 최적 유효 학습률($\eta^*_{eff}$) 또한 유사하게 추정되었습니다.

2. 평가 지표:

   • Out-of-Distribution $R^2$ ($R^2_{OOD}$): 미지의 규모에 대한 스케일링 법칙의 예측 정확도를 측정합니다.
   • 추가 계산 비율 (Extra Compute Ratio, ECR): 동일한 검증 손실을 달키 위해 예측된 학습률을 사용할 때 발생하는 실제 최적 학습률 대비 추가적인 계산 비용을 정량화합니다.
   • 계산 제약 외삽 (Compute-Constrained Extrapolation): 동일한 계산 예산 하에서 데이터 축(고정된 큰 $N$에서 $D$를 변화시킴)과 모델 축(고정된 큰 $D$에서 $N$을 변화시킴)을 따른 외삽을 비교했습니다.

3. 이론적 분석:

   • 본 논문은 AdamW의 암묵적 역학, 특히 원시 학습률 $\eta$와 유효 학습률 $\eta_{eff}$ 사이의 관계를 가중치 노름(weight norm)의 과도적 동작(transient behavior)을 고려하여 분석합니다.
   • $\eta_{eff}$를 명시적으로 제어하는 옵티마이저인 AdamH를 사용하여 실험을 반복함으로써, 가중치 노름 역학의 영향을 분리해 냈습니다.

주요 결과

1. 최적 학습률의 비선형성

연구 결과, 최적 학습률 $\eta^*$는 전역적인 로그 선형 스케일링 법칙을 따르지 않습니다.

• 상향 곡선 (Upward Curvature): 데이터 규모 $D$와 모델 크기 $N$이 증가함에 따라, $\eta^*$는 체계적인 상향 곡선을 나타냅니다. 작은 규모에서 얻은 로그 선형 외삽은 더 큰 규모의 최적 학습률을 일관되게 **과소 예측(underpredict)**합니다.
• 유효 학습률의 선형성: 반면, 최적 유효 학습률($\eta^*_{eff}$)은 동일한 범위 내에서 실질적으로 훨씬 더 선형적인 스케일링 동작을 보이며, 더 높은 피어슨 상관 계수를 나타냅니다.

2. 데이터 축 외삽의 우월성

고정된 계산 예산 하에서 외삽 전략을 비교했을 때:

• 데이터 축 vs 모델 축: 데이터 크기 축($D$)을 따라 스케일링 법칙을 피팅하는 것이 모델 크기 축($N$)을 따를 때보다 훨씬 더 정확한 예측을 제공합니다.
• 계산 효율성: 유효 학습률을 사용하여 데이터 축 외삽을 적용하면, 실제 최적 학습률을 사용하는 것과 비교하여 약 2%의 추가 계산만 발생합니다. 이는 타겟 규모에서 전체 하이퍼파라미터 스윕을 수행하는 데 필요한 계산량의 20% 미만입니다.

3. 메커니즘적 설명

본 논문은 $\eta^*$의 비선형성을 AdamW에서의 가중치 노름의 과도적 역학 때문이라고 설명합니다:

• 두 가지 영역: $\eta$와 $\eta_{eff}$ 사이의 관계는 가중치 노름 $\|w_t\|$가 평형 상태에 도달했는지 여부에 따라 달라집니다.
  • 평형 영역 (Equilibrium Regime): $\eta_{eff} \propto \sqrt{\eta}$ ($\log$-$\log$ 공간에서 기울기 2).
  • 전-평형 영역 (Pre-equilibrium Regime): $\eta_{eff} \propto \eta$ ($\log$-$\log$ 공간에서 기울기 1).
• 규모의 영향: 더 큰 모델/데이터에 대해 최적 학습률이 감소함에 따라, 훈련 호라이즌(training horizon)은 종종 전-평형 영역(가중치 노름이 아직 완화되지 않은 상태) 내에 머물게 됩니다. 이는 유효 스텝 크기 역학을 변화시켜, 관찰된 $\eta^*$의 상향 곡선을 유발합니다.
• AdamH를 통한 검증: $\eta_{eff}$를 명시적으로 제어하는 AdamH를 사용할 때, 스케일링 법칙은 엄격하게 로그 선형이 되며, 이는 표준 AdamW에서의 비선형성이 $\eta$와 $\eta_{eff}$ 사이의 암묵적이고 시간 가변적인 매핑에서 기인함을 확인시켜 줍니다.

의의 및 주장

본 논문은 세 가지 주요 기여를 주장합니다:

1. 경험적 교정: 표준적인 학습률의 로그 선형 스케일링 법칙은 국소적인 특성일 뿐이며, 상향 곡선으로 인해 더 큰 규모에서는 실패하며 부정확한 외삽을 초래한다는 것을 입증했습니다.
2. 실용적인 전이 전략: 하이퍼파라미터 전이를 실질적으로 개선하는 두 가지 구체적인 선택지를 식별했습니다:
   • 원시 학습률($\eta$) 대신 유효 학습률($\eta_{eff}$)을 사용하여 스케일링 법칙을 파라미터화하는 것.
   • 모델 크기 축($N$) 대신 데이터 축($D$)을 따라 외삽하는 것.
3. 이론적 통찰: 관찰된 비선형성을 설명하는 메커니즘을 제공하며, 이를 가중치 노름의 과도적 동작과 연결 짓습니다. 저자들은 $\eta_{eff}$가 스케일 불변 아키텍처에서 최적화 역학을 지배하는 더 근본적인 양이라고 주장합니다.

저자들은 GPT-2 아키텍처 및 AdamW에 집중했다는 점, $\eta^*$ 추정을 위해 3차 피팅에 의존했다는 점, 그리고 분석이 100B 토큰 및 707M 파라미터 규모까지로 제한되었다는 점을 포함한 한계점을 언급합니다. 이들은 이러한 발견이 추가적인 검증 없이 모든 아키텍처나 옵티마이저에 자동으로 일반화된다고 주장하지 않습니다.