🤖 machine learning

Aurora: A Leverage-Aware Spectral Optimizer

이 논문은 Muon의 극성 인자 기하학(polar factor geometry)을 보존하면서 행 균일성(row-uniformity)을 강제하여 뉴런 정체(neuron stagnation)를 방지하고, 특히 매우 넓은 MLP 레이어를 학습시킬 때 최첨단 성능을 달성하는 레버리지 인식 스펙트럼 옵티마이저인 Aurora를 소개한다.

원저자: Alec Dewulf, Dhruv Pai, Li Yang, Ashley Zhang, Ben Keigwin

게시일 2026-06-29

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Alec Dewulf, Dhruv Pai, Li Yang, Ashley Zhang, Ben Keigwin

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 복잡한 퍼즐을 풀기 위해 거대한 작업자 팀(신경망)을 교육하고 있다고 상상해 보십시오. 이 팀에는 MLP 레이어라고 불리는 특정 부서가 있으며, 이곳에서는 업무가 개별 스테이션으로 나뉩니다. 각 스테이션은 하나의 "뉴런"이며, 최종 정답의 아주 작은 부분만을 담당하는 역할을 수행합니다.

이 논문은 이전 관리자인 Muon이 업무를 배정하던 방식의 중대한 결함을 해결하는 새로운 매니저인 Aurora를 소개합니다.

문제점: "게으른 노동자"의 루프

이전 시스템(Muon) 하에서는 관리자가 팀 전체의 성과를 살펴보고 업무의 균형을 맞추려 노력했습니다. 하지만 특정 유형의 작업(수학적으로 "tall matrices"로 표현되는)에 대해, Muon은 실수로 악순고리를 만들어냈습니다:

불균형한 분배: 어떤 작업자(뉴런)들은 기술을 엄청나게 키울 수 있는 흥미진진하고 거대한 업데이트를 받는 반면, 다른 작업자들은 거의 보이지 않을 정도로 미미한 업데이트를 받았습니다.
사멸의 소용돌이: 이렇게 작은 업데이트를 받은 작업자들은 점차 사라지기 시작했습니다. 학습을 하지 못했기 때문에 그들은 팀의 결과물에 기여하지 못하게 되었습니다. 결국, 이들은 "죽은 뉴런(dead neurons)"이 되었습니다. 여전히 급여는 받고 있지만 아무것도 하지 않는 상태가 된 것입니다.
피드백 루프: 이 작업자들이 기여를 멈추자, 관리자(Muon)는 그들을 업데이트할 필요성을 느끼지 못했고, 그 결과 그들에게 주어지는 업데이트는 더욱 작아졌습니다. 이는 스스로를 강화하는 방치 루프였습니다.

논문은 Muon으로 훈련된 대규모 모델에서 상당수의 이러한 작업자들이, 특히 네트워크의 초기 레이어에서 사실상 "일을 그만두었다"는 것을 보여줍니다.

기존의 해결책: "무차별적인 힘" 방식

일부 연구자들은 단순히 모든 작업자가 동일한 양의 에너지를 갖도록 강제하는 방식(row normalization)으로 이 문제를 해결하려 했습니다. 그들은 팀을 살펴보고 누군가 약해져 있으면 단순히 그들을 끌어올렸습니다.

하지만 이 논문은 이러한 무차별적인 접근 방식이 마치 정교한 시계를 고치기 위해 망치를 휘두르는 것과 같다고 주장합니다. 비록 작업자들이 죽는 것은 막았지만, 업무의 형태를 왜곡했습니다. 이는 팀을 자연스러운 퍼즐의 흐름에 맞지 않는 기하학적 구조로 몰아넣었습니다. 즉, 트레이드오프가 발생했습니다. 작업자들은 살려냈지만, 업무의 리듬을 깨뜨린 것입니다.

새로운 솔루션: Aurora

Aurora는 리듬을 깨뜨리지 않으면서 이 문제를 해결하는 새로운 옵티마이저(optimizer)입니다. Aurora를 두 가지 규칙을 동시에 이해하는 숙련된 안무가라고 생각하십시오:

기하학 규칙: 팀은 퍼즐을 효율적으로 풀기 위해 특정한 우아하고 직교하는 패턴(완벽하게 동기화된 춤처럼)으로 움직여야 합니다. 이것이 Muon이 잘했던 "polar factor"입니다.
공정성 규칙: 모든 단 한 명의 작업자도 동일한 양의 에너지와 관심을 받아야 합니다.

Aurora는 이 두 가지를 동시에 수행합니다. 단순히 사후에 약한 작업자를 끌어올리는 대신, 춤의 패턴과 공정성 요구를 모두 만족시키는 완벽한 움직임을 한 번에 계산해 냅니다.

실제 작동 방식:

네트워크의 "tall"한 부분(MLP 레이어의 up 및 gate projection)을 살펴봅니다.
빠른 반복 계산(몇 차례의 암산과 같은 과정)을 실행하여, 팀의 춤 동작을 완벽하게 유지하면서도 누구도 소외되지 않도록 하는 완벽한 업데이트를 찾아냅니다.
이를 통해 "죽은 뉴론" 문제를 완전히 방지합니다.

결과

저자들은 Aurora를 대규모 언어 모델(340 million 및 1.1 billion 파라미터)에 테스트하였으며, 다음과 같은 결과를 얻었습니다:

죽은 뉴런 없음: Muon과 달리, Aurora는 모든 뉴런을 활성화된 상태로 유지하며 기여하게 만들었습니다.
더 나은 성능: Aurora로 훈련된 모델은 Muon이나 "무차별적인 힘"을 사용한 모델보다 더 빠르게 학습했으며, 추론 및 지식 테스트(MMLU 등)에서 더 높은 점수를 기록했습니다.
넓을수록 좋다: 논문은 MLP 레이어의 뉴로이 더 많을수록(팀이 더 넓을수록) Aurora의 이점이 더 커진다는 것을 발견했습니다. 이는 Aurora가 죽은 작업자에게 자원을 낭비하지 않고도 매우 넓고 강력한 네트워크를 훈련하는 데 핵심적인 열쇠임을 시사합니다.

요점

Aurora는 AI 모델을 훈련하는 더 똑똑한 방법입니다. 이는 AI "뉴런"들이 조용히 사라져가는 숨겨진 결함을 해결하여, 학습 과정의 수학적 구조를 완벽하게 유지하면서도 네트워크의 모든 부분이 효과적으로 활용되도록 보장합니다. 이는 마치 절반의 직원들을 실수로 무시하는 매니저에서, 모든 사람이 완벽하게 싱크를 맞춰 춤을 출 수 있도록 보장하는 매니저로 업그레이드하는 것과 같습니다.

기술 요약: Aurora – 레버리지 인지형 스펙트럼 옵티마이저 (Leverage-Aware Spectral Optimizer)

문제 정의

본 논문은 대규모 언 모델(LLM) 학습에 널리 사용되는 스펙트럴 옵티마이저인 Muon(그라디언트 모멘텀 행렬의 폴라 인자(polar factor)를 사용하여 가중치 업데이트를 근사하는 방식)에서 발생하는 치명적인 병리 현상을 식별한다. 특히, 키가 큰 행렬 파라미터(tall matrix parameters, 예: MLP 레이어의 up 및 gate 프로젝션 행렬과 같이 행의 수 $m$ 이 열의 수 $n$ 보다 많은 경우)에 대해, Muon은 매우 불균일한 행 노름(row norms)을 나타낼 수 있다.

이러한 불균일성은 다음과 같은 자기 강화 피드백 루프를 생성한다:

작은 노름을 가진 행에 대응하는 뉴런들은 지속적으로 작은 업데이트를 받는다.
이 뉴런들은 네트워크 출력에 의미 있는 기여를 하지 못하게 되어, 사실상 "죽은(dead)" 상태가 된다.
본 논문은 이러한 현상이 **레버리지 점수(leverage scores)**와 연관되어 있음을 보여주는 기계적 증거를 제공하며, 낮은 행 노름이 모멘텀 버퍼 내의 낮은 레버리지 점수와 상관관계가 있고 Muon이 이를 교정하기보다 그대로 유지한다는 점을 밝힌다.

저자들은 기존의 행 정규화 방법(예: NorMuon)이 행을 재스케일링하여 이 문제를 완화하지만, 이는 Muon 업데이트의 기하학적 구조를 왜곡하여 진정한 폴라 인자로부터 멀어지게 만든다는 점을 관찰했다. 행 노름의 균일성을 강제하는 것과 정밀한 폴라 인자 기하학을 보존하는 것 사이의 이러한 트레이드오프가 해결하고자 하는 핵심 문제이다.

방법론

1. 정규화 및 폴라 정밀도 분석

저자들은 **폴라 인자 정밀도(polar factor precision)**와 행 정규화(row normalization) 사이의 긴장 관계를 먼저 분석한다.

폴라 근사 오차 (PAE): CANS-12와 같은 Newton-Schulz 반복법을 통해 폴라 인자를 계산하는 정밀도가 높을수록 검증 손실(validation loss)이 낮아진다는 것을 입증한다.
기하학적 충돌: 키가 큰 행렬의 경우, 열 직교성(column-orthogonality, 폴라 인자의 특성)과 단위 행 노름(unit row norms)을 동시에 만족하는 것이 수학적으로 불가능함을 증명한다. 사후 정규화(post-hoc normalization)는 업데이트를 진정한 폴라 인자로부터 멀어지게 하여 상당한 기하학적 오차를 유발한다.
U-NorMuon: 중간 단계로서, 저자들은 키가 큰 행렬에만 행 정규화를 적용하고 정규화를 위한 상태 저장형 모멘텀 버퍼를 제거한 U-NorMuon을 제안한다. 이는 통계적 지연(stale statistics)으로 인한 오버헤드 없이 행 정규화의 이점만을 분리하여, 표준 Muon 및 NorMuon보다 우수한 성능을 보임을 입증한다.

2. Aurora 옵티마이저

폴라 기하학을 보존하는 것과 균일한 행 노름을 강제하는 것 사이의 충돌을 해결하기 위해, 저자들은 Aurora를 제안한다.

정식화: Aurora는 두 가지 제약 조건을 준수하면서 그라디언트 내적의 트레이스(trace)를 최대화하는 제약 최 optimization 문제로 정식화한다:
1. 좌측 반직교성 (Left semi-orthogonality): $U^\top U = \eta^2 I$ (폴라 인자 구조 보존).
2. 동일한 행 노름: 모든 행 $i$ 에 대해 $\|U_{i,:}\|^2 = \eta^2 \frac{n}{m}$ .
Riemannian-Aurora: 저자들은 제약 조건들을 하나의 조인트 매니폴드(Stiefel 매니폴드와 동일 행 노름 매니폴드의 교집합)로 취급하는 이론적 참조 솔버인 Riemannian-Aurora를 도출한다. 이는 이론적으로는 정밀하지만, 레버리지 제약이 포함된 대규모 선형 시스템을 풀어야 하므로 계산 비용이 매우 높다.
Practical Aurora: 저자들은 다음 과정을 교대로 수행하는 반복적이고 실용적인 구현체를 제안한다:
1. 행 정규화 (Row Normalization): 노름을 균등하게 하기 위해 행을 재스케일링한다.
2. 폴라 인자 근사 (Polar Factor Approximation): 행렬을 재직교화하기 위해 Newton-Schulz 반복법을 적용한다.
  이 과정은 수렴을 보장하기 위해 파라미터 $\beta$ 에 의해 댐핑(damped)된다. 저자들은 국소적 정규성 조건 하에서 이 교대 투영(alternating projection) 방식이 두 제약 조건을 모두 만족하는 해로 선형적으로 수렴함을 증명한다.

주요 기여

정규화된 스펙트럴 옵티마이저 분석: 본 논문은 사후 단위 행 정규화가 키가 큰 행렬에 대해 폴라 인자 정밀도를 저하시킨다는 점을 입증한다. 또한, 통계적 지연 없이 행 정규화를 적용하는 무상태(stateless) 변형인 U-NorMuon을 도입하여, 행의 균일성과 정밀한 직교화가 동시에 필요함을 증명한다.
뉴런 데스(Neuron Death) 연구: 저자들은 Muon으로 학습된 모델에서의 뉴런 데스 현상을 레버리지 점수의 저하와 연결하여 기계적 설명을 제공한다. 이를 자체 학습 런 및 오픈 웨이트 모델(예: Rnj-1, Laguna-XS.2)에서 경험적으로 검증하였으며, 밀집 MLP 레이어에서 상당한 규모의 죽은 뉴런 집단을 발견했다.
Aurora 및 Riemannian-Aurora: 저자들은 Muon의 이론적 프레임워크를 엄격히 준수하면서 균일한 행 노의를 강제하는 방법으로서 Aurora를 도출한다. 또한 이중 제약 문제에 대한 고비용 참조 솔버로 Riemannian-Aurora를 제공한다.
광범위한 사전 학습 검증: 저자들은 340M 및 1.1B 파라미터 규모에서 Aurora를 검증하여, Muon, NorMuon, U-NorMuon 대비 우수한 다운스트림 성능을 입증했다.

결과

다운스트림 성능

벤치마크 점수: 1.1B 규모에서 Aurora는 Muon 베이스라인을 크게 앞질러 MMLU에서 9.1점 높은 점수를 달 기록했다. 또한 340M 및 1.1B 규모 모두에서 추론 및 지식 벤치마크 세트(ARC, HellaSwag 등) 전반에 걸쳐 가장 높은 평균 점수를 달성했다.
Speedrun: modded-nanoGPT speedrun(옵티마이저 트랙)에서 Aurora는 기존 방법들(Contra-Muon, flooring)과 결합하여 NorMuon 베이스라인보다 타겟 검증 손실에 25단계 더 빠르게 도달하며 새로운 SOTA를 경신했다.

뉴런 활용도

레버리지 점수: Muon으로 학습된 모델은 0 근처에 상당한 모드가 존재하는 양봉 분포(bimodal distribution)의 레버리지 점수를 보인다(죽은 뉴런). 반면, Aurora와 U-NorMuon은 등방성(isotropic) 레버리지 점수 분포를 생성하여 죽은 뉴런을 효과적으로 제거한다.
오픈 웨이트 검증: Muon으로 학습된 오픈 웨이트 모델(Rnj-1 및 Laguna-XS.2) 분석 결과, 초기 밀집 레이어에서 죽은 뉴런의 존재가 확인되었으며(Laguna-XS.2의 첫 번째 레이어에서 최대 약 28%), 이는 문제의 일반성을 입증한다.

효율성 및 확장성

처리량 (Throughput): 분산 환경에서 Aurora의 추가적인 계산 오버헤드는 미미하다. 추가적인 행렬 곱셈 연산을 그라디언트 통신과 오버랩할 수 있기 때문이다. 단일 GPU에서는 처리량이 약간 감소하지만, 분산 학습 시의 실제 시간(wall-clock) 오버헤드는 적절한 수준이다.
MLP 너비 확장성: Muon 대비 Aurora의 성능 이득은 MLP 확장 계수(expansion factor)에 따라 단조롭게 증가한다. MLP 레이어가 넓어질수록 Aurora의 이점이 커지며, 이는 Aurora가 매우 넓은 네트워크를 학습시키는 데 특히 효과적임을 시사한다.

의의 및 주장

본 논문은 Aurora가 현재의 대규모 언어 모델 학습에서 관찰되는 특정 병리 현상인 키가 큰 행렬 파라미터에서의 뉴런 데스를 해결한다고 주장한다. 이는 Muon 업데이트의 이론적 기하학적 구조와 균일한 행 노름의 실용적 필요성을 모두 존중하는 옵티마이저를 도출함으로써 가능하다.

저자들은 하향식(bottom-up) 설계 철학을 강조한다. 즉, 순수하게 이론적인 모델에 의존하기보다 관찰된 학습 역학(뉴런 데스, 레버리지 점수 불균형)에 기반하여 알고리즘을 도출한다. 이들은 특정 파라미터 클래스(예: 키가 큰 행렬 vs 넓은 행렬)에 맞춤화된 옵티마이저를 설계하는 **아키텍처-옵티마이저 공동 설계(architecture-optimizer co-design)**가 기존의 "원 사이즈 피츠 올(one-size-fits-all)" 전략보다 뛰어난 성과를 낼 수 있는 유망한 방향임을 주장한다.

마지막으로, 본 논문은 Aurora가 매우 넓은 MLP 레이어의 효과적인 학습을 가능하게 한다고 제안한다. 이는 현재의 Muon 기반 학습에서 죽은 뉴런의 축적으로 인해 제한될 수 있는 부분이다. 이 연구는 Aurora를 단순한 성능 향상 도구가 아니라, 대규모 언어 모델의 스펙트럴 최적화에 깔린 기하학적 가정에 대한 필수적인 교정책으로 포지셔닝한다.