너비-깊이 동시 확장을 위한 µP의 스펙트럼 조건

기존의 µP는 모델의 너비 확장에만 최적화되어 있어, 현대적인 거대 모델처럼 깊이를 함께 늘릴 때 학습이 불안정해지는 한계가 있었다. 이 논문은 선형 대수 기반의 단순한 '스펙트럼 조건'을 통해 어떤 옵티마이저를 쓰더라도 너비와 깊이에 상관없이 안정적인 학습과 하이퍼파라미터 전이를 가능하게 하는 통합 공식을 제공한다.

딥러닝 모델 학습은 가중치 업데이트를 통해 데이터의 특징(feature)을 변화시키는 과정이다. 모델의 너비(width)가 커지면 각 층의 뉴런 수가 늘어나 특징의 분산이 변하고, 깊이(depth)가 깊어지면 잔차 연결을 통해 특징의 변화량이 누적되어 폭발하거나 사라지는 문제가 발생한다. 기존 µP는 너비 확장에 따른 변화만 제어했기에 깊은 모델에서는 한계가 있었다.

이 논문의 핵심은 가중치 행렬이 입력 특징을 변화시키는 정도를 나타내는 'RMS 연산자 노름'을 깊이(L)의 역수에 비례하도록 설계하는 것이다. 이는 마치 긴 파이프라인에서 물의 압력을 조절할 때, 파이프가 길어질수록 각 지점의 밸브를 조금씩 더 조여서 최종 출력단의 수압을 일정하게 유지하는 것과 같은 원리이다.

결과적으로 모델의 층이 아무리 깊어져도 출력층에서의 특징 변화량은 일정하게 유지(scale-invariant)된다. 이를 통해 연구자는 수조 개의 파라미터를 가진 거대 모델을 직접 튜닝하는 대신, 아주 작은 모델에서 최적의 학습률을 찾고 이를 수학적 공식에 따라 거대 모델에 그대로 적용할 수 있게 된다.

잔차 블록 구조를 $h_l = h_{l-1} + \alpha_l W_l^{(2)} W_l^{(1)} h_{l-1}$ 형태로 정의하고, 너비($n$)와 깊이($L$)가 무한대로 발산하는 상황에서의 특징 전파를 분석했다. 특징 학습의 안정성을 위해 초기 가중치 노름의 곱이 $\Theta(1/L)$을 만족해야 함을 수학적으로 도출했다. [입력 특징 벡터 $h_{l-1}$] → [가중치 행렬 $W_l$과의 연산 및 잔차 합산] → [출력 특징 $h_l$ 생성] → [전체 깊이 $L$에 걸쳐 누적된 특징의 크기가 발산하지 않도록 가중치 스케일을 조정]하는 메커니즘이다.

도출된 스펙트럼 제약 조건을 다양한 옵티마이저의 업데이트 규칙에 매핑했다. 예를 들어 Muon-Kimi 옵티마이저의 경우, SVD 기반의 업데이트 노름이 너비와 깊이에 따라 어떻게 변하는지 계산하고, 이를 스펙트럼 조건에 맞추기 위해 학습률($\eta$)을 $\eta \propto 1/\sqrt{n}$으로 설정하는 공식을 유도했다. [옵티마이저별 업데이트 수식 입력] → [너비/깊이 변수 대입] → [업데이트 노름의 스케일 계산] → [최종 하이퍼파라미터 스케일링 법칙 결정] 순으로 구현 전략을 수립했다.

GPT-2 스타일의 Transformer 모델을 사용하여 특징 학습 안정성을 검증했다. 표준 파라미터화(SP) 방식은 깊이가 깊어질수록 특징 노름이 급격히 증가하여 학습이 불안정해졌으나, 제안된 µP 방식은 너비($n=128$ to $4096$)와 깊이($L=4$ to $256$)의 모든 조합에서 일정한 노름을 유지하며 안정적으로 학습되었다.

하이퍼파라미터 전이 실험에서는 작은 모델($n=256, L=4$)에서 찾은 최적의 학습률이 너비와 깊이를 대폭 확장한 모델에서도 여전히 최적의 성능을 냄을 확인했다. 특히 SP 방식은 모델 크기에 따라 최적 학습률이 크게 변하여 재튜닝이 필수적이었던 반면, µP는 거의 동일한 학습률 곡선을 보여주었다.

LayerNorm이 없는 극한의 조건에서도 µP는 안정적인 학습을 유지했다. SP는 깊이가 깊어짐에 따라 그래디언트 폭주로 인해 NaN(Not a Number) 오류가 발생하며 학습이 붕괴되었으나, µP는 $L=256$의 깊은 구조에서도 견고한 성능을 기록했다.

본 연구는 RMS 연산자 노름(RMS operator norm)을 기반으로 한 통합 스펙트럼 프레임워크를 제안한다. 기존의 Tensor Programs 기반 분석이 매우 복잡했던 것과 달리, 초등 선형 대수와 확률론적 상한 추정만을 사용하여 너비-깊이 확장에 따른 µP 조건을 유도하여 이론적 접근성을 높였다.

2층 이상의 잔차 블록에서 발생하는 가중치 업데이트 간의 고차 상호작용(higher-order interaction)을 분석했다. 초기 조건과 1차 업데이트 조건만 적절히 설정하면, 2차 이상의 고차 업데이트 항들은 자동으로 $\Theta(1/L)$ 범위 내에서 제어됨을 수학적으로 증명하여 프레임워크의 견고함을 뒷받침했다.

옵티마이저별로 상이한 업데이트 스케일을 체계적으로 분류했다. SGD는 학습률을 깊이에 비례($\eta \propto L$)하게, AdamW는 깊이에 무관($\eta \propto 1$)하게, Muon-Kimi는 너비의 제곱근에 역비례($\eta \propto 1/\sqrt{n}$)하게 설정해야 한다는 구체적인 파라미터화 대응표를 제시하여 연구자들이 즉시 구현에 활용할 수 있도록 했다.

이론적 도출은 주로 선형 잔차 MLP를 기반으로 이루어졌으며, 비선형 활성화 함수와 멀티헤드 어텐션 구조에 대해서는 실험적 검증과 가정을 통해 확장성을 보였다. 또한 모멘텀 계수를 상수로 가정하여 분석에서 제외한 점이 한계로 명시되었다.