본문으로 건너뛰기

© 2026 Molayo

Lobste.rs헤드라인2026. 05. 22. 10:40

Aurora: 직사각형 행렬을 위한 레버리지 인지 최적화 도구 (Leverage-Aware Optimizer)

요약

Muon 최적화 도구의 행 노름 이방성 문제를 해결하기 위해 제안된 Aurora 최적화 도구에 관한 연구입니다. 행 노름 균일성과 직교성을 동시에 유지하는 최급 강하 문제를 해결하여 1.1B 모델 학습 시 100배의 데이터 효율성을 달성했습니다.

핵심 포인트

  • Muon의 행 노름 이방성으로 인한 뉴런 사멸 문제 해결
  • 행 노름 균일성과 직교성을 결합한 새로운 최적화 방식 제시
  • 1.1B 모델 학습 시 100배의 데이터 효율성 달성
  • modded-nanoGPT 스피드런에서 현재 SoTA 기록 경신
  • 기존 Muon 대비 오버헤드 6% 수준의 높은 효율성

뒤로 가기

Aurora: 직사각형 행렬을 위한 레버리지 인지 최적화 도구 (Leverage-Aware Optimizer)

요약 (TL;DR)

  • 우리는 Muon의 업데이트가 키가 큰 행렬 (tall matrices)에서 행 노름 이방성 (row-norm anisotropy)을 상속받으며, 이로 인해 MLP 레이어의 뉴런 중 상당수가 영구적으로 사멸(die)할 수 있음을 보여줍니다. 업데이트를 행 단위로 정규화 (Row normalizing)하면 이를 해결할 수 있지만, 직교성 (orthogonality)을 희생해야 합니다.
  • 우리는 행 노름 균일성 (row-norm uniformity)과 직교성 (orthogonality)의 결합 제약 조건 하에서의 최급 강하 (steepest descent) 문제를 공식화하고, 그 해결책으로 Aurora 최적화 도구를 제시합니다.
  • 우리는 Aurora를 사용하여 1.1B 모델을 학습시켰으며, 이 모델은 오픈 소스 인터넷 데이터에서 **100배의 데이터 효율성 (100x data efficiency)**을 달성하고 HellaSwag와 같은 일반적인 평가에서 더 큰 모델들을 능가했습니다.
  • 우리는 modded-nanoGPT 스피드런에 PR을 제출했으며, Aurora는 현재의 SoTA (State-of-the-Art) 기록을 능가했습니다.
  • 튜닝되지 않은 Aurora는 전통적인 Muon 대비 오버헤드가 6%에 불과하며, 즉시 교체 가능한 (drop-in replacement) 방식입니다.
  • 리만 (Riemannian) 및 바닐라 (vanilla) Aurora 모두에 대한 코드를 오픈 소스로 공개합니다: github.com/tilde-research/aurora-release.

목차

  • 서론 — Muon의 성공과 NorMuon 퍼즐
  • 행 정규화로 인한 긴장 — 키가 큰 행렬에서 행 정규화는 직교성과 충돌함
  • 정규화는 뉴런 사멸을 방지함 — Muon은 MLP 뉴런을 죽입니다; 레버리지 점수 (leverage scores)가 그 이유를 설명합니다
  • Muon 목적 함수 확장 — 결합 제약 조건으로부터 Aurora 유도
  • 결과 — nanoGPT 스피드런 SoTA, 1B 사전 학습(pretraining), 그리고 다운스트림 평가
  • 변곡점 — 이것이 중요한 이유

1: 서론

Muon 최적화 도구 [1]는 nanoGPT 스피드런 경연 대회 [2]에서의 성공으로 처음 주목을 받았습니다. 여기서 Muon은 스텝당 더 많은 계산을 필요함에도 불구하고, 수렴까지의 실제 시간 (wall-clock time) 측면에서 AdamW [3]를 능가했습니다. 그 이후로 개선된 분산 구현 [4], 더 정확한 직교화 루틴 [5] [6], 그리고 계산 및 통신 오버헤드를 줄이는 방법 [4]을 통해 Muon을 확장하려는 상당한 노력이 이루어졌습니다. 그 결과, Muon(클리핑 및 per-head 버전 포함)은 프런티어 규모의 모델 (frontier-scale models)을 학습시키기 위한 점점 더 인기 있는 선택지가 되었습니다 [7] [8] [18].

최근 오픈 연구 벤치마크 (open research benchmarks)에서 최첨단 성능 (state-of-the-art, SoTA)을 경신하는 Muon 변형 모델들이 급증하고 있습니다 [9] [10]. 한 가지 예로 NorMuon [9]이 있으며, 이는 현재 modded-nanoGPT 스피드런 (speedrun)에서 SoTA를 기록하고 있습니다. NorMuon은 Adam의 파라미터별 정규화 (per-parameter normalization)와 유사하게, 각 행을 역 RMS 노름 (inverse RMS norm)으로 스케일링하는 추가 단계를 통해 Muon을 보강합니다. 특히, 이러한 행 정규화 (row normalization) 단계는 Muon의 업데이트를 실제 직교 그래디언트 (orthogonal gradient)로부터 상당히 멀어지게 만들 수 있습니다. NorMuon이 여전히 성공적이라는 사실은 Muon의 공식화 (formulation)에 행 정규화에 의해 해결되고 있는 격차가 존재할 수 있음을 시사합니다.

우리는 행 정규화의 효과를 연구하였으며, Muon이 MLP 레이어에서 *뉴런 사멸 (neuron death)*을 초래할 수 있음을 발견했습니다. 이는 일부 뉴런이 학습 초기 단계에서 지속적으로 작은 업데이트를 받아 회복하지 못하는 현상을 말합니다. 우리는 MLP 레이어의 up 및 gate 프로젝션 (projections)에 대한 업데이트 시 행 전체에 질량 (mass)을 균등하게 재분배함으로써 이러한 실패 모드 (failure mode)를 피할 수 있음을 보여줍니다. 이러한 관찰에 착안하여, 우리는 그래디언트 직교화 (gradient orthogonalization)의 정밀도를 희생하지 않으면서 뉴런 사멸을 방지하기 위해 이 메커니즘을 사용하는 Aurora를 제안합니다. 우리는 Aurora가 1B 규모에서 Muon 및 NorMuon 모두에 비해 큰 이득을 달성하며, modded-nanoGPT 최적화 트랙의 현재 SoTA를 능가함을 보여줍니다. 나아가, Aurora의 성능 향상은 MLP 너비 (width)에 따라 확장됨을 발견하였으며, 이는 Aurora가 MLP 확장 계수 (expansion factors)가 큰 네트워크에 특히 효과적임을 믿게 합니다.

2: 행 정규화로 인한 긴장 (Tension Due to Row Normalization)

Muon의 핵심 알고리즘 구성 요소는 행렬의 극 인자 (polar factor)를 계산하기 위한 반복 알고리즘입니다.
얇은 특이값 분해 (thin Singular Value Decomposition, SVD) $A = U egin{pmatrix} ext{diag}( ext{S}) \ 0 ext{ or } 0^T ext{ (if not square)} ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ } ext{ }

[5] [6] [1]을 계산하기 위한 matmul-only (행렬 곱셈 전용) 반복 알고리즘 (iterative algorithms)의 존재는 Muon을 대규모로 실행 가능하게 만드는 핵심 요소입니다. 여기서는 두 가지 알고리즘에 집중합니다: modded-nanoGPT[21]에서 가져온 5차 반복법 (quintic iteration)과 8단계 Polar Express (PE-8)입니다. [5] 부록 A에서 polar factor (극 인자)를 계산하기 위한 다양한 알고리즘에 대해 더 철저한 개요를 제공합니다.

우리는 우리의 학습 설정(부록 B에 설명됨)에서 Muon이 더 정밀한 직교화 (orthogonalization), 즉 $Q$를 더 정확하게 근사하는 반복 알고리즘을 사용할 때 단조적으로 더 나은 다운스트림 손실 (downstream loss)을 달성한다는 것을 발견했습니다. 이러한 차이를 보여주기 위해 그림 1(Figure 1)에 우리가 테스트한 가장 정밀도가 낮은 반복법(5차 반복법)과 가장 정밀한 방법(PE-8)을 사용한 Muon 실행 결과를 도식화했습니다. 이 결과는 polar factor의 정밀도가 성능이 뛰어난 Muon 구현에 있어 중요한 구성 요소임을 확인시켜 줍니다.

그림 1. Muon은 더 정밀한 직교화 루틴을 사용할수록 단조적으로 더 낮은 손실을 달성합니다. PE-8은 1B 규모에서 5차 반복법을 실질적으로 능가합니다.

NorMuon은 polar factor의 각 행을 단위 RMS 노름 (unit RMS norm)을 갖도록 스케일링하는 추가 단계를 통해 Muon을 보완합니다.

우리는 행이 열보다 많은 행렬 (tall matrices)의 경우, 이러한 행 정규화 (row normalization)가 polar factor의 정밀도를 크게 감소시킬 수 있음을 보여주며, 이는 이론적으로나 우리의 이전 결과(그림 1)에 비추어 볼 때 바람직하지 않습니다.

주장 1. 행이 열보다 많은 행렬 (tall matrix)은 동시에 열 직교 (column-orthogonal)이면서 균일한 단위 행 노름 (uniform unit row norms)을 가질 수 없습니다. 즉, $Q$가 $Q^TQ = I$를 만족한다면, $Q^T$는 모든 행 노름이 1이 될 수 없습니다.

증명. $Q^T$의 모든 행이 노름 1을 가진다고 모순적으로 가정해 봅시다. 그러면

[IMG:1]

이러한 사실들을 종합하면 $Q^TQ
eq I$가 되며, 이는 $Q$가 tall matrix라는 우리의 가정에 모순됩니다. $Q^T$에서 행의 평균 노름은 $ rac{ ext{tr}(I)}{n} = rac{m}{n}$입니다.

따라서 NorMuon에서의 행 정규화 (row normalization)는 직교화 (orthogonalization) 루틴에 필연적으로 정밀도 결함 (precision defect)을 유발합니다. 우리는 이러한 극성 정밀도 결함 (polar precision defect)이 행 노름 (row norms)이 균일하지 않은 행렬에서 상당히 클 수 있음을 발견했습니다. 특히, 도입된 오차의 크기는 동일한 설정에서 quintic iteration 대비 PE-8의 정밀도 향상분보다 훨씬 더 클 수 있습니다 (다양한 Newton-Schulz iteration에 대한 전체 비교는 부록 A를 참조하십시오).

Row normalization precision

그림 2. (a) 행 노름 표준 편차를 {1, 2, 3, 4, 5}로 하여 n=1000개의 무작위 512×128 행렬을 샘플링하고, 서로 다른 극성 인자 (polar factor) 알고리즘 (PE-8, CANS-10, CANS-12 및 quintic)으로 직교화했습니다. 결과로 얻은 극성 인자의 정밀도는 그래프에 표시되어 있습니다 (점선). 그런 다음 단위 행 정규화 (unit row normalization)를 적용하고 업데이트된 정밀도를 플로팅했습니다 (진한 선). 행 노름 표준 편차가 3일 때, 직교 결함 (orthogonality defect)은 0.06에서 정점을 찍습니다. (b) 각 행 표준 노름 편차에 대한 평균 결함 크기를 나타냅니다.

실제로, 이러한 알고리즘들이 동일한 설정에서 NorMuon과 함께 사용될 때 PE-8과 quintic iteration 사이의 성능 격차는 사라집니다. 그러나 두 실행 모두 우리의 Muon+PE-8 베이스라인보다 우수한 성능을 보이는데, 이는 행 정규화가 독립적으로 유용할 수 있음을 시사합니다.

그림 3. NorMuon은 1B 규모에서 더 정밀한 직교화로부터 이득을 얻지 못합니다. PE-8과 quintic은 동일한 다운스트림 손실 (downstream loss)을 달성하며, 이는 행 정규화가 극성 인자 정밀도를 저하시킨다는 우리의 결과와 일치합니다.

행 정규화에 의해 발생하는 결함을 완화하기 위해, 우리는 단순히 세로로 긴 행렬 (tall matrices)의 행 노름이 1이 아닌 다른 값을 갖도록 정규화할 수 있습니다. 우리는 이 변형을 U-NorMuon이라고 부르며, 이것이 우리의 340M 설정에서 NorMuon보다 우수한 성능을 보임을 확인했습니다 (그림 4; 설정은 부록 B에 설명됨).

이제 우리의 분석을 NorMuon 하에서 행 노름(row norm) 단위 업데이트를 받는 가로가 넓은 행렬(wide matrices)과 정사각형 행렬(square matrices)로 전환하겠습니다. 좌직교(left-orthogonal) 가로형 또는 정사각형 행렬은 반드시 모든 행 노름이 1과 같으므로, 이 경우 직교성(orthogonality)에 의해 행 노름의 균일성(row-norm uniformity)이 함축됩니다.

따라서 우리는 PE-8과 같은 정밀한 직교화 루틴(orthogonalization routines) 하에서는 행 정규화(row-normalization)가 불필요하거나 심지어 해로울 수도 있다고 예상합니다. 반복적 극분해 알고리즘(Iterative polar factor algorithms)은 가공되지 않은 그래디언트(raw gradients)가 조건이 나쁜(ill-conditioned) 경향을 보이는 훈련 초기 단계에서 정밀하게 수렴하지 못할 수 있으며, 이로 인해 빠르게 노후화되는 행 노름 통계량이 축적될 수 있습니다. 이 경우, 통계량이 수정될 때까지 행 정규화는 행들을 균일하게 만들 것입니다. 우리는 340M 규모에서 키가 큰(tall) 행렬에만 행 정규화를 적용했을 때 NorMuon 및 U-NorMuon보다 작지만 무시할 수 없는 차이로 더 나은 성능을 보이는 증거를 발견했습니다 (그림 4).

그림 4. 다운 프로젝션(down projection)의 열 정규화(Column normalization)는 NorMuon 하에서 업(up) 및 게이트(gate) 행 정규화와 매우 유사하게 작동합니다. 키가 큰 행렬(업 및 게이트 프로젝션)에만 적용된 NorMuon은 모든 NorMuon 변형뿐만 아니라 U-NorMuon 및 Muon보다 우수한 성능을 보입니다.

우리는 Muon이 MLP 레이어 내 뉴런의 큰 부분집합을 효과적으로 사멸(die)하게 만들 수 있지만, 이러한 병리적 현상이 (U-)NorMuon에 의해 완화된다는 것을 보여줄 것입니다. 기존 문헌의 아이디어를 바탕으로, 우리는 이를 모든 설정에서 발견되는 Muon과 NorMuon 사이의 성능 격차에 대한 설명으로 제안합니다. 그런 다음 우리는 극분해 인자(polar factor)의 정밀도를 희생하지 않으면서 키가 큰 파라미터에 대한 업데이트를 효과적으로 행 정규화하는 Aurora를 도출합니다.

3: 정규화는 뉴런 사멸을 방지한다

우리는 *사멸한 모델 구성 요소(dead model component)*를 훈련의 초기 단계 이후 지속적으로 작은 학습 신호를 받는 모델 파라미터의 부분집합으로 정의합니다. 우리는 다음 세 가지 기준을 통해 사멸한 모델 구성 요소를 식별합니다:

  • 낮은 유효 그래디언트 노름 (Low effective gradient norm). 특정 배치(batch)에 대해 낮은 유효 그래디언트 노름 (effective gradient norm)을 받는 모델 구성 요소는 손실 함수 지형 (loss landscape)에서 상대적으로 평탄한 방향에 해당해야 합니다. 따라서 지속적으로 작은 그래디언트 노름을 갖는 구성 요소는 학습 신호를 거의 받지 못하게 됩니다. 기여도 분석 (attribution) 및 가지치기 (pruning) 문헌에서는 그래디언트 노름을 특징 중요도 (feature importance)의 대리 지표로 사용하는 강력한 전례가 있습니다 [11] [12] [13].
  • 낮은 유효 업데이트 노름 (Low effective update norm). Muon과 같은 옵티마이저 (optimizer)로 학습된 네트워크의 경우, 작은 그래디언트 노름만으로는 사멸한 네트워크 구성 요소를 식별하기에 충분하지 않습니다. 특히, 직교화 (orthogonalization) 단계는 그래디언트가 적용되기 전에 작은 크기의 방향을 증폭시키는 것을 포함하여 그래디언트를 급격하게 변화시킬 수 있습니다. 따라서 우리는 실제 파라미터 변화량을 측정하는 유효 업데이트 노름 (effective update norm) 또한 작아야 한다는 조건이 필요합니다.
  • 학습 과정에서의 지속성 (Persistence over training). 파라미터가 학습의 특정 시점에서 작은 업데이트 노름을 받더라도 네트워크 출력에 여전히 유의미하게 기여할 가능성이 있습니다. 예를 들어, 일부 구성 요소는 초기에 유용한 특징을 학습하고 해당 특징이 안정화됨에 따라 더 작은 업데이트를 받을 수 있습니다. 이러한 정적이지만 유용한 구성 요소와 진정으로 사멸한 구성 요소를 구분하기 위해, 우리는 낮은 유효 그래디언트 및 업데이트 노름이 학습 과정 전반에 걸쳐 지속적이고 예측 가능할 것을 요구합니다.

그림 5. 사멸한 뉴런을 식별하기 위한 각 기준의 시각화: 낮은 유효 그래디언트 노름, 낮은 유효 업데이트 노름, 그리고 학습 과정 전반에 걸친 두 지표의 지속성. 세 가지 기준을 모두 만족하는 뉴런을 사멸한 것으로 식별합니다.

우리는 우리가 정의한 방식의 뉴런 사멸이 Muon으로 학습된 네트워크에서 발생할 수 있으며 실제로 발생한다는 것을 보여줄 것입니다. 이는 MLP 레이어의 세로가 긴 행렬 (tall matrices)이 매우 불균일한 행 노름 (row-norms)을 가진 업데이트를 받도록 허용되기 때문입니다. 특히, Muon 업데이트는 학습 초기에 우연히 큰 업데이트를 받은 네트워크 구성 요소에 특권을 부여하여 다른 구성 요소들이 사멸하도록 만든다는 것을 보여줄 것입니다. 우리는 U-NorMuon이 이 문제를 완전히 방지함을 확인했습니다.

뉴런 사멸의 메커니즘 (The Mechanism of Neuron Death)

SVD (특이값 분해)가 희소한(thin) 행렬의 경우, 행 $i$의 *레버리지 점수 (leverage score)*를 $ au_i$라고 정의합니다. 행의 레버리지 점수는 해당 행 업데이트 스케일에 직접적으로 비례하므로, 업데이트 행렬에서 지속적으로 낮은 레버리지 점수를 가진 행들은 업데이트 질량 (update mass)의 일관되게 작은 부분만을 할당받게 됩니다. 우리는 적절한 조건(well-conditioned)을 갖춘 행렬의 직교화 (orthogonalization) 과정에서 행의 레버리지가 대부분 보존된다는 것을 보여줄 수 있으며, 따라서 Muon 환경에서는 원본 그래디언트 (raw gradient)에서 레버리지가 낮은 행이 업데이트에서도 낮은 레버리지를 갖는 경향이 있습니다.

Figure 6. 직사각형 행렬에 대한 레버리지 점수의 대화형 예시. 각 행의 레버리지 점수는 극 분해 인자 (polar factor) 하에서의 업데이트 에너지 점유율을 결정합니다. 레버리지가 작은 행은 비례적으로 작은 업데이트를 받습니다.

주장 2 (비형식적). 키가 큰 (tall) 행렬의 경우, 그래디언트 노름 (gradient norm)이 작은 행은 직교화 후에도 여전히 작은 노름을 유지하므로, Muon 업데이트는 뉴런 사멸을 방지하지 않습니다.

AI 자동 생성 콘텐츠

본 콘텐츠는 Lobste.rs AI의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.

원문 바로가기
0

댓글

0