μ스케일링으로 작은 모델을 효율적으로 확대하는 방법

초록

본 논문은 μP 이론을 기반으로 모델 폭을 확대할 때 정적·동적 동등성을 보장하는 일반화된 업스케일링 기법을 제시한다. 가중치 복제와 학습률·가중치 감쇠 등 하이퍼파라미터의 스케일링 규칙을 수학적으로 증명하고, 이를 통해 작은 모델에서 찾은 최적 하이퍼파라미터를 큰 모델에 그대로 적용할 수 있는 μTransfer 방식을 확장한다. 실험에서는 MLP, ResNet, GPT‑2 등 다양한 아키텍처와 데이터셋에 대해 제안 방법이 기존 업스케일링 대비 학습 효율과 최종 성능에서 우수함을 확인한다.

상세 분석

논문은 먼저 μP(μ‑Parameterization)와 any‑dimensional 아키텍처 개념을 활용해 “폭 확대(width upscaling)”에 대한 이론적 토대를 마련한다. 정적 동등성은 가중치를 k‑배 복제하고 1/k 스케일링함으로써 원본 모델과 동일한 함수 출력을 보장한다는 명제(Prop 2.1)로 정의된다. 여기서 k는 각 레이어별 정수 배수이며 입력·출력 레이어는 복제하지 않는다. 동적 동등성은 학습 과정 전체에서 동일한 함수 궤적을 유지하도록 학습률 γ와 가중치 감쇠 λ, 그리고 옵티마이저의 추가 파라미터 ε를 레이어별 k에 맞춰 조정하는 규칙을 제시한다(Prop 2.2, 2.4). 특히, SGD뿐 아니라 Adam·AdamW와 같은 entrywise 옵티마이저에 대해 Q_t가 m 차 동차함수라는 가정 하에 γ↑ = k^m · k^{‑1} γ, ε↑ = k^{‑1} ε, λ↑ = k^{‑1} λ(또는 k^{‑1} k^{‑m} λ) 로 스케일링하면 확대된 모델이 원본과 동일한 파라미터 진화식을 만족한다는 점을 증명한다.

이론적 결과를 바탕으로 실제 업스케일링 알고리즘을 설계한다. 먼저 작은 모델을 학습한 뒤, 위의 가중치 복제·스케일링을 적용해 넓은 모델을 초기화한다. 그러나 복제된 파라미터는 저차원 서브스페이스에 머무를 위험이 있으므로, 레이어별 작은 Gaussian 노이즈를 추가해 대칭을 깨고 새로운 자유도를 활성화한다. 이 노이즈의 크기는 μP가 제시하는 “optimal feature learning” 조건을 만족하도록 선택된다.

핵심적인 μTransfer 확장은 위의 스케일링 규칙을 그대로 적용해 작은 모델에서 튜닝한 학습률·모멘텀·ε 등을 큰 모델에 그대로 옮기는 것이다. 실험에서는 동일한 하이퍼파라미터 집합을 사용했을 때, 업스케일된 모델이 스크래치 학습 모델과 거의 동일하거나 더 빠른 수렴을 보이며, 특히 대규모 Transformer(GPT‑2)에서 FLOP 절감 효과가 두드러졌다. 또한, 무작위 초기화 대비 복제 초기화가 초기 손실을 크게 낮추고, 노이즈 주입이 없을 경우 학습이 정체되는 현상을 확인함으로써 이론과 실험이 일치함을 입증한다.

전체적으로 논문은 “폭 확대 = 함수 보존 + 하이퍼파라미터 스케일링”이라는 명확한 원칙을 제시하고, 이를 Tensor Programs 프레임워크로 무한 폭 한계까지 일반화한다. 이는 기존 경험적 히트와 달리 수학적 보장을 제공함으로써, 대규모 모델 파이프라인에서 작은 모델을 효율적인 워밍 스타트와 하이퍼파라미터 전이 수단으로 활용할 수 있는 실용적인 길을 연다.