Learning Rate Transfer in Normalized Transformers

"Learning Rate Transfer in Normalized Transformers" (Shigida, Hanin, Gromov 저) 논문에 대한 상세한 기술적 요약입니다.

1. 문제 제기

본 논문은 Loshchilov 등 (2025) 이 제안한 Transformer 아키텍처 변형인 **Normalized Transformers(nGPT)**에서의 하이퍼파라미터 (HP) 전이 문제를 다룹니다.

• 배경: nGPT 는 가중치와 활성화 값을 광범위하게 정규화 (LayerNorm/RMSNorm 대신 $\ell_2$-norm 사용) 함으로써 훈련 속도를 크게 향상시키고, 가중치 감쇠 (weight decay) 및 학습률 워밍업 (learning rate warmup) 의 필요성을 제거합니다.
• 문제: nGPT 는 모델 크기에 따라 하이퍼파라미터가 스케일링되지만, 저자들은 모델의 너비 (dimension), 깊이 (layer 수), 토큰 범위 (훈련 반복 횟수) 에 걸쳐 학습률 전이가 실패함을 관찰합니다.
• 기존 방법의 한계: **$\mu$P(Maximal Update Parametrization)**와 같은 표준 전이 기법은 가중치와 활성화 값의 정렬 (alignment) 에 대한 특정 이론적 가정 (특히, 특정 방식으로 완전히 정렬되거나 정렬되지 않음) 에 의존합니다. 저자들은 이러한 가정이 nGPT 에서는 성립하지 않아 확장 시 최적의 성능을 내지 못한다고 발견했습니다.

2. 방법론

저자들은 정렬 지수 (alignment exponents) (Everett 등, 2024) 에 기반한 원칙적인 이론적 프레임워크와 수치 실험을 결합한 **$\nu$GPT(nu-GPT)**라는 새로운 파라미터화를 제안합니다.

A. 이론적 프레임워크: 정렬 지수

$\mu$P 에서 사용하는 "완전 정렬"($\alpha=1$) 또는 "정렬 없음"($\alpha=1/2$) regimes 를 가정하는 대신, 저자들은 모델 차원 $d_{model}$에 대해 가중치 업데이트와 활성화 값이 어떻게 스케일링되는지를 설명하는 정렬 지수($\alpha, \omega, \nu$) 를 경험적으로 측정합니다.

• 행렬 - 벡터 곱의 변화가 어떻게 스케일링되는지에 대한 지수를 정의합니다: $\|\Delta W h\| \sim d^{\alpha} \|\Delta W\| \|h\|$.
• 주요 발견: nGPT 에서는 $\max\{\alpha, \nu\} \approx 3/4$인 "중간 정렬 (mid-alignment)" regimes 가 관찰됩니다. 이는 $\mu$P 의 가정인 $\max\{\alpha, \nu\} = 1$(완전 정렬) 또는 $1/2$(정렬 없음) 과 모순됩니다.

B. $\nu$GPT 파라미터화

"중간 정렬" 가설 ($\max\{\alpha, \nu\} = 3/4$) 과 경험적 스케일링 법칙에 기반하여, 저자들은 학습률 ($\eta$) 과 초기화 파라미터에 대한 구체적인 스케일링 규칙을 유도합니다. $m_{width}$, $m_{depth}$, $m_{data}$를 기준 모델에 대한 너비, 깊이, 토큰 수의 배수라고 합시다.

$\nu$GPT 의 주요 스케일링 규칙:

1. 전역 학습률: $m_{data}^{-1/3}$로 스케일링됩니다. 이는 일부 이전 SDE 기반 이론에서 제안된 $m_{data}^{-1/2}$와 달리 토큰 범위를 보정합니다.
2. 임베딩 가중치 ($\eta_{input}$): $m_{width}^{-1/2}$로 스케일링됩니다.
3. 히든/언임베딩 가중치 ($\eta_{hidden}, \eta_{output}$): $m_{width}^{-3/4}$로 스케일링됩니다. 이는 $\mu$P(버전에 따라 $m_{width}^{-1}$ 또는 $m_{width}^{-1/2}$ 사용) 와의 결정적인 차이점입니다.
4. 깊이 보정:
   • 학습 가능한 리스케일링 벡터 $\alpha_{A, init}$와 $\alpha_{M, init}$는 $m_{depth}^{-1}$로 스케일링됩니다.
   • "Depth-$\mu$P"(히든 학습률을 $m_{depth}^{-1/2}$로 스케일링) 와 달리, $\nu$GPT 는 nGPT 에서 히든 학습률은 깊이와 함께 스케일링해서는 안 된다고 발견했습니다. 학습 가능한 리스케일러가 자연스럽게 깊이 안정성을 처리하기 때문입니다.
5. 초기화 상수: 스케일 파라미터 ($\alpha_{scale}$) 에 대해서는 $d_{model}^{-1/2}$로 스케일링하는 대신 고정된 상수 (예: 0.03) 를 사용합니다.

3. 주요 기여

1. nGPT 에서의 $\mu$P 실패 식별: 표준 $\mu$P 의 가중치 - 활성화 정렬에 대한 가정이 정규화된 아키텍처에서는 유효하지 않으며, 이로 인해 학습률 전이가 부진함을 입증했습니다.
2. $\nu$GPT 제안: nGPT 에서 관찰된 "중간 정렬" regimes($\approx 3/4$) 를 명시적으로 타겟팅하는 새로운 파라미터화를 도입했습니다.
3. 토큰 범위 스케일링: nGPT 에서 최적 학습률이 토큰 수에 따라 $T^{-1/3}$로 스케일링됨을 확립하여, 이전 이론적 예측인 $T^{-1/2}$를 정제했습니다.
4. 깊이 전이 통찰: 학습 가능한 리스케일러를 가진 정규화된 모델의 경우, 히든 학습률에 대한 명시적인 깊이 스케일링은 불필요하며 오히려 해로울 수 있음을 보였습니다. 대신 리스케일러 초기화를 스케일링하는 것으로 충분합니다.

4. 실험 결과

저자들은 OLMo 2 토크나이저를 사용하여 FineWeb-Edu데이터셋에서 $\nu$GPT 를 검증했습니다.

• 너비 전이:
  • 기준 nGPT: 학습률 전이가 없었으며, 너비가 증가함에 따라 성능이 저하되었습니다.
  • $\mu$P 스타일 보정: 최적 학습률에서 약간의 편차를 보였습니다.
  • $\nu$GPT: 0.26B 에서 3.22B 파라미터에 이르는 너비 전반에 걸쳐 완벽한 학습률 전이를 달성했으며, 튜닝된 기준 모델 대비 검증 손실 (validation loss) 이 감소하지 않았습니다.
• 깊이 전이:
  • 기준 nGPT 는 적절한 깊이 전이를 보였으나 극단적인 깊이에서는 불안정해졌습니다.
  • $\nu$GPT 는 8 에서 128 레이어에 이르는 깊이 전반에서 안정성과 성능을 유지하며, 큰 깊이에서 학습률 선택에 대한 민감도가 낮음을 보여주었습니다.
• 토큰 범위 전이:
  • 실험은 훈련 반복 횟수가 증가함에 따라 최적 학습률에 대한 $T^{-1/3}$스케일링 법칙을 확인했습니다.
• 정렬 지수 검증:
  • 경험적 측정은 $\omega_{hidden} = \omega_{output} = 1/2$이고 $\max\{\alpha, \nu\} \approx 0.75$(중간 정렬) 임을 확인하여 $\nu$GPT 의 이론적 기반을 검증했습니다.

5. 의의

• 실용적 영향: $\nu$GPT 를 통해 연구자들은 작은 모델에서 하이퍼파라미터를 튜닝한 후 이를 대규모 규모로 직접 전이하여 대규모 정규화 Transformer 를 훈련할 수 있게 되었으며, 이는 상당한 컴퓨팅 자원을 절약합니다.
• 이론적 발전: 본 논문은 $\mu$P 의 보편성에 도전하며, 다른 아키텍처 선택 (특히 정규화 방식) 은 서로 다른 정렬 가정과 스케일링 법칙을 필요로 함을 입증했습니다.
• 효율성: 가중치 감쇠와 워밍업의 필요성을 제거하면서도 견고한 스케일링을 가능하게 함으로써, $\nu$GPT 는 차세대 기반 모델 (foundation models) 을 위한 더 간결하고 효율적인 훈련 파이프라인을 제공합니다.

요약하자면, 본 논문은 정규화 Transformer 의 스케일링 문제에 대한 엄격하고 경험적으로 근거 있는 해결책을 제시하며, "일률적 (one-size-fits-all)"인 $\mu$P 접근법을 정규화 아키텍처의 고유한 역학을 존중하는 맞춤형 파라미터화 ($\nu$GPT) 로 대체합니다.