기술 요약: 가중치 벡터의 크기와 방향 분리를 통한 신경망 학습 개선

1. 문제 정의

현대적인 신경망 학습은 가중치 행렬을 단일한 모놀리식(monolithic) 객체로 취급하는 Adam이나 Muon과 같은 옵티마이저에 의존합니다. 그러나 모든 가중치 행렬 $W$는 본질적으로 두 가지 구별되는 양, 즉 크기(magnitude, $\|W\|$)와 방향($\hat{W} = W/\|W\|$)을 내포하고 있습니다. 표준 최적화 과정에서 이 두 양은 서로 결합되어 있어 불안정하고 예측 불가능한 학습 역학을 초래합니다:

• 방향 변화의 의존성: 업데이트로 인해 발생하는 각도 변화는 현재 크기에 반비례하여 결정됩니다. 결과적으로 학습률(LR)이 방향 변화율을 직접적으로 제어하지 못하며, 이 변화율은 레이어마다, 그리고 시간에 따라 달라질 수 있습니다.
• 의도하지 않은 크기 드리프트(Magnitude Drift): 업데이트는 (스케일 불변성이나 노이즈로 인해) 현재 가중치와 수직인 방향으로 발생하는 경향이 있으며, 이는 방사형 그래디언트(radial gradient)가 없음에도 불구하고 크기를 증가시킵니다. 이러한 "통제되지 않는 성장"은 실무자들이 학습을 안정화하기 위해 **가중치 감쇠(weight decay)**와 같은 간접적인 메커니즘에 의존하거나, 초기 불안정성을 방지하기 위해 **웜업(warmup)**에 의존하게 만듭니다.
• 스케일링의 한계: 현재의 스케일링 방법론(예: $\mu$P)은 모델의 너비나 깊이를 변경할 때 하이퍼파라미터 재조정(re-tuning)을 요구하는 경우가 많습니다. 또한, 가중치를 고정된 노름(fixed-norm) 구체 위에 제한하는 최근의 접근 방식들은 학습 가능한 크기 파라미터가 부족하여, 스케일 제어를 전적으로 정규화 레이어에 맡기는 한계가 있습니다.

2. 방법론: 크기-방향(MD) 분리

저자들은 각 가중치 행렬을 고정된 노름의 방향과 학습 가능한 크기 이득(gain)으로 분해하는 옵티마이저 수정 방식인 크기-방향(Magnitude–Direction, MD) 분리를 제안합니다.

핵심 분해

각 가중치 행렬 $W$는 다음과 같이 재매개변수화됩니다:
$$W = \text{diag}(\gamma_{\text{row}}) \, \hat{W} \, \text{diag}(\gamma_{\text{col}})$$
여기서:

• $\hat{W}$는 방향이며, 고정된 초구(hypersphere, 예: 프로베니우스 노름) 위에 존재하도록 제한됩니다.
• $\gamma_{\text{row}}$와 $\gamma_{\text{col}}$은 각각 행과 열에 작용하는 벡터인 학습 가능한 크기 이득입니다.

최적화 절차

이 방법은 모델의 포워드 패스(forward pass)를 변경하지 않고 옵티마이저 단계 내에서 완전히 작동합니다(모델은 하나의 융합된 텐서 $W$를 보게 됩니다):

1. 복구(Recovery): 융합된 가중치 $W$로부터 $\hat{W}$와 이득 $\gamma$를 복구합니다.
2. 그래디언트 분할(Gradient Splitting): 그래디언트 $G = \partial L / \partial W$를 방향 그래디언트($\partial L / \partial \hat{W}$)와 각 축에 대해 합산된 이득 그래디언트로 분할합니다.
3. 방향 업데이트: 기본 옵티마이저(예: Adam, Muon)를 사용하여 학습률 $\eta_W$로 $\hat{W}$를 업데이트한 후, 고정된 노름을 유지하도록 구체 위로 투영(project)합니다.
4. 이득 업데이트: 별도의 학습률 $\eta_\gamma$를 사용하여 이득 $\gamma$를 업데이트합니다(양수성을 보장하기 위해 softplus와 같은 매끄러운 맵을 선택적으로 사용).
5. 재조립(Reassembly): 다음 포워드 패스를 위해 $W$를 재구성합니다.

주요 특성

• 옵티마이저 불가지론적(Optimizer Agnostic): 이 방법은 모든 기본 옵티마이저(Adam, Muon, Shampoo 등)를 감싸는 래퍼(wrapper) 역할을 합니다.
• 직접적인 LR 제어: 방향을 구체 위에 고정함으로써, 학습률이 상대적인 각도 업데이트를 직접 설정하게 하여 현재 가중치 크기에 대한 의존성을 제거합니다.
• 휴리스틱 제거: 가중치가 이미 노름이 제한되어 있으므로 가중치 감쇠가 필요 없으며, 초기 단계의 불안정한 대규모 업데이트가 조절되므로 웜업이 필요하지 않습니다.
• 세밀한 스케일 제어: 정규화 레이어의 이득에만 의존하는 방식과 달리, MD 분리는 개별 행과 열의 스케일을 명시적으로 학습하고 제어할 수 있게 하여 모델이 특정 특징을 증폭하거나 감쇄할 수 있도록 합니다.

3. 주요 기여 및 연구 결과

실험적 결과

저자들은 밀집 모델(181M ~ 1.29B 파라미터)과 대규모 혼합 전문가(MoE) 모델(최대 6.7B 총 파라미터) 전반에 걸쳐 MD 분리를 검증했습니다.

• 성능 향상: MD 분리는 잘 튜닝된 베이스라인(AdamW 및 Muon)보다 우수한 성능을 보입니다. 대규모 MoE 모델에서 MuonMD(Muon에 MD 분리를 적용)는 AdamW와 동일한 손실 값에 도달하는 데 약 2배 적은 연산량을 사용합니다.
• 학습률 전이(Learning Rate Transfer): 이 방법은 모델 너비 간의 제로샷 학습률 전이를 가능하게 합니다. 작은 모델에서 발견된 최적의 행렬 학습률은 재조정 없이도 훨씬 더 큰 모델에서도 최적으로 유지됩니다. 이는 너비와 깊이 스케일링 모두에 적용됩니다.
• 웜업 없는 학습: 실험 결과, 웜업을 제거하는 것은 안정성을 유지할 뿐만 아니라 최종 손실을 개선합니다. 이는 초기 학습 단계가 감소된 학습률로 인해 낭비되지 않기 때문입니다.
• Ablation 통찰:
  • 정규화 축: 방향을 프로베니우스 노름(전체 행렬)을 통해 제한하는 것이 가장 유연하며, 행/열 제약과 유사한 성능을 보임을 확인했습니다.
  • 이득 매개변수화: 결합된 행 및 열 이득($\gamma_{\text{row}} + \gamma_{\text{col}}$)이 스칼라 또는 단일 축 이득보다 우수합니다. 이득을 위한 softplus 매개변수화는 직접 또는 지수적 업데이트보다 약간 더 일관된 우위를 제공합니다.
  • 임베딩: 임베딩 벡터를 단위 L2 노름으로 제한하는 것이 제약이 없는 경우보다 약간 더 나은 성능을 보입니다.

관련 연구와의 비교

본 논문은 MD 분리를 구형 학습(spherical training) 및 재매개변수화에 관한 점증하는 연구 흐름 속에 위치시킵니다.

• nGPT와 달리, nGPT는 특정 아키텍처 변경(예: L2 정규화, 보간된 잔차)과 구형 제약을 결로 묶지만, MD 분리는 표준 아키텍처에 적용 가능한 옵티마이저 측면의 수정 사항입니다. nGPT 아키텍처에 적용했을 때 MD 분하는 성능을 더욱 향상시킵니다.
• Muown이나 HyperP와 달리, 이들은 학습 가능한 크기 없이 노름을 고정하거나 해당 방법을 특정 옵티마이저에 종속시키지만, MD 분리는 명시적으로 크기를 학습하며 옵티마이저에 구애받지 않습니다.
• MD 분리는 이전 연구(예: Kosson et al., 2024a)에서 식별된 "크기-방향 간섭" 문제를 두 양의 역학을 분리함으로써 해결합니다.

4. 의의 및 주장

본 논문은 크기와 방향을 별도로 제어되는 양으로 취급하는 것이 더 예측 가능한 학습 역학을 제공하며, 현대적인 옵티마이저에 대한 단순하고 폭넓게 적용 가능한 개선책이라고 주장합니다.

• 단순성: 아키텍처 변경이 필요 없으며, 옵티마이저 단계에서 최소한의 계산 오버헤드(요소별 연산)만을 추가합니다.
• 강건성: 가중치 감쇠와 웜업의 필요성을 제거하여 학습 레시피를 단순화하고, 하이퍼파라미터 탐색 공간을 줄여줍니다.
• 확장성: 모델 너비 간의 학습률 전이를 위한 원칙적인 메커니즘을 제공하여, 대규모 언어 모델 스케일링의 주요 병목 현상을 해결합니다.

저자들은 옵티마이저 단계의 추가 로직이 약간의 오버헤드를 발생시키지만, 학습 안정성, 연산 효율성, 하이퍼파라미터 전이 측면에서의 이점이 MD 분리를 대규모 신경망 학습을 위한 중요한 진보로 만든다고 결론짓습니다.