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

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

왜 중요한가

핵심 기여

통합 스펙트럼 확장 조건 정의

잔차 네트워크에서 가중치 노름과 업데이트 크기가 너비와 깊이에 따라 어떻게 변해야 특징 학습이 안정화되는지 규정하는 수학적 조건(Condition 3.1)을 수립했다.

다양한 옵티마이저로의 µP 확장

SGD, AdamW뿐만 아니라 Muon-Kimi, Sophia, Lion 등 최신 옵티마이저들에 대해 스펙트럼 조건을 만족하는 구체적인 하이퍼파라미터 파라미터화 공식을 도출했다.

이론적 분석의 단순화

복잡한 Tensor Programs 이론 대신 초등 선형 대수와 확률론만을 사용하여 너비-깊이 확장에 따른 µP 원리를 명확하게 규명했다.

실증적 하이퍼파라미터 전이 입증

GPT-2 스타일 모델 실험을 통해 너비와 깊이를 동시에 확장해도 작은 모델에서 찾은 최적 학습률이 대형 모델에서도 최적으로 유지됨을 확인했다.

핵심 아이디어 이해하기

딥러닝 모델 학습은 가중치 업데이트를 통해 데이터의 특징(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를 기반으로 이루어졌으며, 비선형 활성화 함수와 멀티헤드 어텐션 구조에 대해서는 실험적 검증과 가정을 통해 확장성을 보였다. 또한 모멘텀 계수를 상수로 가정하여 분석에서 제외한 점이 한계로 명시되었다.

실무 활용

대규모 언어 모델(LLM)이나 생성형 AI 모델을 개발할 때 발생하는 막대한 하이퍼파라미터 튜닝 비용을 획기적으로 절감할 수 있다. 작은 모델에서의 실험 결과를 대형 모델에 즉시 적용할 수 있는 실무 가이드를 제공한다.

수조 단위 파라미터 모델 학습 시, 작은 모델(예: 100M)에서 찾은 최적 학습률을 대형 모델(예: 100B)에 즉시 적용
Muon, Sophia, Lion 등 최신 옵티마이저를 새로운 아키텍처에 도입할 때 안정적인 스케일링 법칙 수립
모델의 깊이를 추가하는 아키텍처 변경 시, 별도의 하이퍼파라미터 재검색 없이 안정적인 학습 보장

코드 공개 여부: 공개

코드 저장소 보기

키워드

µP(최대 업데이트 파라미터화)Width-Depth Scaling(너비-깊이 확장)Spectral Condition(스펙트럼 조건)Hyperparameter Transfer(하이퍼파라미터 전이)Residual Networks(잔차 네트워크)

용어 해설

최대 업데이트 파라미터화(µP): — 모델의 너비가 무한히 커지는 극한 상황에서도 특징 학습의 동역학을 일정하게 유지하도록 하이퍼파라미터를 스케일링하는 기법이다. 작은 모델에서 최적화된 학습률을 거대 모델에 재튜닝 없이 전이할 수 있게 하여 학습 비용을 획기적으로 줄여준다.
RMS 연산자 노름(RMS Operator Norm): — 행렬이 입력 벡터의 RMS(Root Mean Square) 크기를 얼마나 변화시키는지를 측정하는 지표이다. 모델의 각 층을 통과할 때 데이터의 스케일이 어떻게 변하는지 분석하는 핵심 도구로, 학습 안정성 확보를 위한 스케일링 조건 설정에 사용된다.
잔차 네트워크(Residual Network): — 입력 데이터를 다음 층으로 직접 전달하는 지름길(Skip Connection)을 추가하여 층이 깊어져도 그래디언트 소실 문제를 완화하는 아키텍처이다. 현대적인 Transformer 구조의 근간이 되며, 본 논문에서는 이 구조에서의 깊이 확장을 집중적으로 다룬다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

Adam의 좌표별 적응성이 저랭크 복원을 깨뜨리는 이유

대규모 MoE 학습률을 단계적으로 확장하는 방법

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

Adam의 좌표별 적응성이 저랭크 복원을 깨뜨리는 이유

대규모 MoE 학습률을 단계적으로 확장하는 방법