TL;DR
대형 언어 모델 학습에서 가중치의 스펙트럼 관리가 학습 안정성과 일반화의 핵심 요인이다. Pion은 weight 행렬을 좌우의 orthogonal 변환으로 업데이트해 singular values를 보존하고, 학습 중 spectral norm의 변화를 차단한다. 이를 통해 µP 프레임워크와의 호환성을 유지하면서도 AdamW·Muon 대비 안정성과 효율성을 제공한다.
왜 중요한가
대형 언어 모델 학습에서 가중치의 스펙트럼 관리가 학습 안정성과 일반화의 핵심 요인이다. Pion은 weight 행렬을 좌우의 orthogonal 변환으로 업데이트해 singular values를 보존하고, 학습 중 spectral norm의 변화를 차단한다. 이를 통해 µP 프레임워크와의 호환성을 유지하면서도 AdamW·Muon 대비 안정성과 효율성을 제공한다.
핵심 기여
Spectrum-preserving bilateral update on iso-spectral manifold
Wt+1 = exp(-η Gout_t) Wt exp(-η Gin_t) 형태의 업데이트를 도입하여 W의 singular values를 보존하고 left/right subspaces를 회전시키는 핵심 규칙을 제시한다. Rt = exp(-η Gout_t), Pt = exp(-η Gin_t)는 각각 왼쪽/오른쪽 직교 행렬로 작용하며, 이로써 spectral norm은 불변하게 유지된다.
RMS-based scale-consistent updates across matrices
각 weight 매트릭스에 대해 Gin_t, Gout_t의 Frobenius 노름을 매트릭스 크기에 맞춰 정규화하고, per-weight 계수 α_t로 업데이트 크기가 일정하게 유지되도록 한다. 이로써 대형 모델에서 큰 학습률도 안정적으로 활용 가능하다.
Lie-algebra momentum designs
입력-사이드/출력-사이드에서의 모멘텀을 정확히 누적하는 Lie-algebra 모멘텀과, transported ambient-space 모멘텀, second-order 모멘텀의 네 가지 구성 중에서 Lie+Lie가 최적의 성능을 보임을 확인한다. 이는 스펙트럼 보존 업데이트의 기하학적 특성에 잘 맞는다.
Alternate update to reduce computation
입력-사이드/출력-사이드 업데이트를 교대로 수행하는 Alternate update를 제시하여 계산 비용을 낮추면서도 스펙트럼 보존 특성을 크게 해치지 않는다. Lie+Lie 구성에서 최적 성능에 근접하도록 한다.
Efficient matrix exponential approximation
exp(·)의 근사로 second-order 근사를 채택하고 Cayley 변환과의 비교를 통해 2차 근사가 스펙트럼 보존에 충분히 우수하다고 보인다. 매 업데이트는 항상 항등행렬에서 시작하므로 누적 오차가 억제된다.
µP-호환성 및 너비 전이 실험
단순 스펙트럼 정규화(Scheme I)와 Newton–Schulz 방식의 정교한 직교화(Scheme II) 두 가지 경로로 µP 조건을 만족하도록 제시한다. 너비(width) 확장에 따른 하이퍼파라미터 전이가 성공적으로 이뤄짐을 실험으로 확인한다.
실험적으로 우수한 안정성과 경쟁력 있는 성능
LLM pretraining, supervised finetuning, RLVR 등 다양한 설정에서 Pion의 안정성 및 일반화 측면이 기존 최적화기AdamW, Muon 대비 경쟁력 있음을 다수 벤치마크와 실험으로 확인한다.
핵심 아이디어 이해하기
출발점: 일반적인 가중치 최적화에서 가중치의 스펙트럼은 자주 변화하며 spectral norm의 제어가 필요하다. 기존 옵티마이저는 가중치의 스펙트럼을 임의로 조절하거나 정규화를 통해 보정하지만, 이는 학습 속도나 안정성에 한계를 준다. 논문은 W를 좌/우의 orthogonal 행렬로 곱해 스펙트럼을 보존하는 업데이트를 도입한다. 이 업데이트는 Wt의 Singular Values를 고정한 채, 행/열 공간을 회전시키는 방식으로 작동한다. 이렇게 하면 ∥Wt∥2가 상수로 유지되며, 업데이트의 크기는 회전의 강도에 의해 결정된다.
방법론
전체 접근: Wt를 Wt+1 = exp(-η Gout_t) Wt exp(-η Gin_t) 형태로 갱신한다. 여기서 Gin_t = Wt^⊤ Gt − Gt^⊤ Wt, Gout_t = Gt Wt^⊤ − Wt Gt^⊤ 이고, Gt는 ∇f(Wt)이다. 업데이트는 좌우 두 Lie-algebra 요소를 구하고, 이를 2D 평면 회전으로 분해해 각 서브공간에서 독립적으로 회전을 수행한다. Rt = exp(-η Gout_t), Pt = exp(-η Gin_t)로 생성된 두 Orthogonal 행렬의 곱으로 W를 갱신하고, W의 singular values를 보존한다.
관련 Figure
좌-우 두 Lie-algebra 업데이트의 크기 불균형 문제를 시각적으로 보여주며, 본 논문의 핵심 설계 원칙인 scale-consistent rotation의 필요성을 보조한다.
Figure 2: Inconsistent updates in Pion(좌우 측의 회전 강도가 매트릭스 간 불균형으로 나타남).
RMS 기반의 scale-consistent 업데이트의 효과를 시각화하여, Naive 업데이트에 비해 학습 안정성과 수렴 속도가 개선되는 것을 보여준다.
Figure 3: Bilateral 업데이트의 학습 곡선 및 RMS 스케일링의 효과 비교.
입력-사이드/출력-사이드 모멘텀의 다양한 구현 방식이 학습 속도에 미치는 차이를 보여주며, Lie+Lie 구성이 최적의 결과를 도출함을 뒷받침한다.
Figure 4: Lie-algebra 모멘텀 구성들의 비교(Transported ambient-space, Lie-algebra momentum 등).
Scheme I/II의 수렴 특성과 width에 따른 학습률 전이가 모두 잘 작동함을 시사하며, µP 호환성의 실용적 이점을 뒷받침한다.
Figure 14: µP 학습률 그리드에 대한 수렴 분석(Width별 일반화 성능 비교).
주요 결과
Stable LLM pretraining에서 Pion은 Pandas로 나타난 벤치마크에서 AdamW, Muon보다 평균적으로는 다소 낮은 평균 Val Loss를 보였고, Muon과의 경쟁적 성능을 유지했다. UL 큰 깊이의 네트워크 200층에서도 Pion은 가장 안정적인 학습 곡선을 보였으며, Jacobian 노름에서도 균일한 깊이 의사결정을 보여 주었다. RLVR 설정에서도 Pion이 우수한 수렴 속도를 보이며, 코드 생성 및 수학 reasoning 벤치마크에서 강한 성능을 보인다. 또한 normalization-free pretraining에서 NaN 없이 학습을 지속하는 등 스펙트럼 보존 업데이트의 안정성을 확인했다. µP 호환성 실험에서 Scheme I/II를 통해 width 확장에 따른 학습률 전이가 일관되게 유지되었고, 두 방식 모두 hyperparameter transferability를 확인했다.
실무 활용
Pion은 가중치 스펙트럼을 보존하는 최적화 기법으로, 대규모 language model의 pretraining과 finetuning에 적용 가능하다. 특히 normalization 없이도 안정적 학습이 가능하므로, 스펙트럼 제어가 필요한 프로덕션 시나리오에서 유용하다.
- 대형 언어 모델의 대규모 pretraining에서 spectral norm을 유지하며 학습 안정성 확보
- SFT/RLVR 등 파인튜닝에서 spectral structure 보존으로 성능 저하 방지
- 깊은 모델(수백 계층) 학습에서 손실 스파이크 감소와 안정적 수렴 보장
- µP 정책에 맞춘 모델 확장성 및 학습률 전이성 확보
코드 공개 여부: 공개
코드 저장소 보기키워드
AI 요약 · 북마크 · 개인 피드 설정 — 무료
출처 · 인용 안내
인용 시 "요약 출처: AI Trends (aitrends.kr)"를 표기하고, 사실 확인은 원문 보기 기준으로 진행해 주세요. 자세한 기준은 운영 정책을 참고해 주세요.