기술적 요약: OrScale – 계층별 신뢰비 스케일링을 통한 직교화 최적화

1. 문제 제기

Muon 옵티마이저는 뉴턴 - 슈르츠 반복을 통해 행렬 값 단계의 방향을 고정하는 직교화된 업데이트를 활용하여 대규모 훈련을 위한 실용적인 도구로 부상했습니다. 그러나 Muon 은 중요한 한계를 지니고 있습니다. 즉, 업데이트의 방향을 최적화하지만, 크기는 단일한 본질적으로 전역적인 학습률에 의해 제어된다는 점입니다.

이러한 전역적 스케일링은 비최적입니다. 왜냐하면 서로 다른 계층들 (예: 어텐션 프로젝션, MLP 업/다운 프로젝션, 임베딩) 은 매우 다른 기울기 분포와 유효 곡률을 보이기 때문입니다. Moonlight 변형은 요소별 RMS 를 균등화하기 위해 정적 형태별 인자를 도입했지만, 훈련 중 계층 크기의 동적 변화를 적응할 수는 없습니다.

SGD 와 Adam 에서 단계 크기를 전역 학습률로부터 성공적으로 분리해낸 계층별 적응형 비율 스케일링 (LARS) 과 계층별 적응형 모멘트 스케일링 (LAMB) 을 Muon 으로 이전하려는 시도는 실패했습니다. 해당 논문은 단순한 이전이 세 가지 구체적인 구조적 이유로 인해 붕괴된다고 지적합니다:

1. 퇴화 분모: 직교화된 방향의 프로베니우스 노름 ($\|Q_\ell\|_F$) 을 신뢰비 분모로 사용하는 것은 행렬의 형태 (랭크) 에 의해 결정되는 거의 일정한 값이므로 업데이트 크기를 반영하지 못해 실패합니다.
2. 단위 불일치 및 클립 포화: 원시 모멘트 노름 ($\|\tilde{M}_\ell\|_F$) 을 사용하면 기울기 단위와 가중치 단위 간의 단위 불일치가 발생하여 신뢰비가 거의 보편적으로 상단 클리핑 경계 ($r_{max}$) 에 도달하게 되어 적응이 무효화됩니다.
3. 분리된 가중치 감쇠의 비약: 보정된 분모와 분리된 가중치 감쇠를 결합하면 가중치 노름이 통제 불가능하게 증가 (10~50 배) 하는 양성 피드백 루프가 생성되어 훈련을 불안정하게 만듭니다.

2. 방법론

저자들은 OrScale을 제안합니다. 이는 단일 설계 원칙에 기반한 프레임워크입니다: 계층별 신뢰비의 분모는 옵티마이저가 매개변수에서 빼게 될 실제 업데이트 방향의 프로베니우스 노름을 측정해야 합니다.

핵심 원칙

$D_{\ell,t}$를 전역 학습률이 적용되기 전에 가중치 $W_{\ell,t}$에서 빼지는 실제 매개변수 공간 방향이라고 합시다. 신뢰비는 다음과 같이 정의됩니다:
$$r_{\ell,t} = \frac{\|W_{\ell,t}\|_F}{\|D_{\ell,t}\|_F + \epsilon}$$
중요하게도, $D_{\ell,t}$는 가중치 감쇠 항과 직교화된 방향을 모두 포함해야 합니다.

두 가지 변형

이 프레임워크는 두 가지 구체적인 구현을 산출합니다:

1. OrScale (일반/비전 기본값):

   • 업데이트 방향: $D_{\ell,t} = \lambda W_{\ell,t} + Q_{\ell,t}$ (여기서 $\lambda$는 가중치 감쇠이고 $Q$는 직교화된 방향입니다).
   • 형태 인자: $s_\ell = 1$ (정적 스케일링 없음).
   • 보정: 없음.
   • 가중치 감쇠: 신뢰비 스케일링 업데이트 단계에 직접 결합됩니다.

2. OrScale-LM (LLM 기본값):

   • 업데이트 방향: $D_{\ell,t} = \lambda W_{\ell,t} + s_\ell Q_{\ell,t}$.
   • 형태 인자: AdamW 의 요소별 RMS 와 일치시키기 위해 Moonlight 의 정적 인자 $s_\ell = 0.2\sqrt{\max(m_\ell, n_\ell)}$를 채택합니다.
   • 보정: 초기화 시 한 번 계산되는 지연 계층별 보정 상수($c_{\text{denom}, \ell}$) 를 도입합니다. 이는 초기 신뢰비를 정확히 1 ($r_{\ell,0}=1$) 로 고정합니다.
   • 보정의 중요성: 이를 통해 OrScale-LM 은 새로운 하이퍼파라미터 탐색 없이 Moonlight(및 결과적으로 AdamW) 에 대해 조정된 학습률 및 가중치 감쇠 하이퍼파라미터를 계승할 수 있습니다. 이는 초기화 시 너비 불변성을 보장합니다.

알고리즘적 세부 사항

• 결합된 가중치 감쇠: 가중치 감쇠는 별도의 되당김으로 적용되는 대신 신뢰비 계산 ($\lambda W$가 분모의 일부) 에 통합됩니다. 이는 분리된 변형에서 관찰된 비약적 피드백 루프를 방지합니다.
• 효율성: 이 방법은 매 행렬 계층당 하나의 추가 스칼라와 단계당 두 번의 프로베니우스 노름 계산만 필요로 하며, 기존 뉴턴 - 슈르츠 반복에 의해 지배되므로 벽시계 시간 오버헤드는 무시할 수준 (<1%) 입니다.

3. 주요 기여

• 알고리즘: 계층별 신뢰비 스케일링을 Muon 에 성공적으로 통합하는 통합적이고 파라미터가 적은 레시피 (알고리즘 1) 입니다. 저자들은 설계 공간에서 고유한 "작동 가능한 모서리"(실제 업데이트 방향 분모 + 결합된 WD) 를 식별합니다.
• 실패 분석: Muon 신뢰비 설계 공간을 엄격하게 매핑한 것 (표 1) 으로, 세 가지 실패 모드 (퇴화 분모, 클립 포화, 비약적 가중치 증가) 를 분석적 및 경험적으로 특징짓습니다. 이는 왜 이들이 실패하고 제안된 해결책이 성공하는지 보여줍니다.
• 이론:
  • 수렴: 핵노름 $O(1/\sqrt{T})$ 비볼록 수렴 보장을 증명합니다.
  • 적응형 이득: 측정 가능한 계층 이질성 하에서 계층 적응형 하강 상수 $\kappa_{\text{eff}} > 1$을 확립하여 Muon 의 동질적 상수보다 엄격하게 개선합니다.
  • 분리: 클립 포화 조건 하에서 MuTrust/MuScale 과의 엄격한 분리를 증명합니다.
  • 보정: OrScale-LM 이 초기화 시 $\mu$P 스타일 학습률 이전을 유지함을 보여줍니다.
• 경험적 평가: 비전 및 언어 모델 작업 전반에 걸친 포괄적인 평가.

4. 실험 결과

비전 (CIFAR-10 / DavidNet)

• 성능: OrScale 은 94.05% $\pm$ 0.08 의 검증 Top-1 정확도를 달성하여 다섯 가지 베이스라인 중 1 위를 차지했습니다.
• 개선: 이는 표준 Muon(93.70%) 대비 +0.35 포인트, Muon+Moonlight(93.75%) 대비 +0.30 포인트의 개선입니다.
• 안정성: OrScale 은 이론적 예측인 계층 적응형 이득과 일치하는 Muon 보다 더 넓은 안정적인 학습률 창을 보여주었습니다.
• 애블레이션: 식별된 실패 모드에 해당하는 변형들 (예: $\|Q_\ell\|_F$ 또는 원시 모멘트 사용) 은 예측된 증상 (예: 99.5~100% 클립 포화 또는 가중치 노름 비약) 을 보였으며 OrScale 의 성능을 달성하지 못했다는 애블레이션 연구가 확인되었습니다.

대규모 언어 모델 (FineWeb-Edu 사전 훈련)

• 스케일링: 네 가지 모델 스케일 (125M~1.1B 파라미터) 에서 평가되었습니다.
• 비교: OrScale-LM 은 네 가지 스케일 중 세 곳 (125M, 545M, 1.1B) 에서 Muon+Moonlight 를 능가했으며 모든 스케일에서 AdamW 를 능가했습니다.
• 손실 감소: 1.1B 스케일에서 OrScale-LM 은 최종 검증 교차 엔트로피를 2.6251 로 달성했으며, 이는 Muon+Moonlight 의 2.6360 과 AdamW 의 2.7304 와 비교됩니다.
• 효율성: 이 이점은 컴퓨팅 범위 전반에 걸쳐 유지되었으며, OrScale-LM 은 LLM 에게 컴퓨팅이 가장 중요한 영역에서 깔끔한 승리를 보여주었습니다.

5. 중요성 및 주장

이 논문은 OrScale 이 Muon 의 스펙트럴 기하학과 LARS/LAMB 의 계층별 적응성 사이의 간극을 성공적으로 메웠다고 주장합니다. 그 주요 중요성은 다음과 같습니다:

1. 크기 문제 해결: 구조적 실패 모드에 강건한 Muon 을 위한 최초의 시간 가변적 계층별 크기 제어를 제공합니다.
2. 실용적 이전성: OrScale-LM 변형과 지연 보정을 통해, 실무자들이 AdamW 또는 Moonlight 에 대해 조정된 학습률을 OrScale-LM 으로 재조정 없이 직접 이전할 수 있게 하여 채택 장벽을 낮춥니다.
3. 이론적 건전성: 이전의 하이브리드 시도들이 실패한 이유에 대한 엄격한 이론적 정당성을 제공하며, 제안된 방법이 이질성 하에서 엄격하게 더 나은 수렴 상수를 산출함을 증명합니다.
4. 경험적 검증: 이러한 이론적 개선이 소규모 비전 작업과 대규모 LLM 사전 훈련 모두에서 실질적인 이득으로 전환됨을 보여주어, OrScale 을 행렬 값 계층을 위한 최첨단 옵티마이저로 확립합니다.

저자들은 LLM 결과에 대해 겸손함을 유지하며, 1.1B 스케일의 이득은 더 작은 스케일에서 단일 시드 노이즈와 비교 가능하다고 지적합니다. 따라서 이 결과는 모든 영역에서 결정적인 통계적 우위라기보다는 방법의 효능에 대한 방향적 증거로 간주되어야 합니다. 그러나 스케일 전반에 걸친 일관된 개선과 실패 모드의 성공적인 애블레이션은 제안된 설계 원칙을 강력하게 지지합니다.