[Paper] TEON: Tensorized Orthonormalization Beyond Layer-Wise Muon for 대규모 언어 모델 사전 학습

Source: arXiv - 2601.23261v1

개요

이 논문은 **TEON (Tensorized Orthonormalization)**이라는 새로운 옵티마이저를 소개합니다. TEON은 기존의 Muon 기법을 레이어별 행렬 직교화에서 전체 네트워크 수준의 텐서 차원으로 확장하여 그래디언트를 처리합니다. 이를 통해 TEON은 더 강력한 수렴 보장을 제공하고, 60 M에서 1 B 파라미터에 이르는 대형 언어 모델(LLM) 사전 학습 시 일관되게 더 낮은 퍼플렉시티를 달성합니다.

주요 기여

• Tensor‑level orthogonalization: Muon의 레이어별 행렬 직교화를 구조화된 고차원 텐서로 일반화하여 레이어 간 그래디언트 관계를 포착합니다.
• Theoretical improvement: Muon보다 더 강력한 수렴 경계를 제공하며, 전역 직교화가 그래디언트 분산을 보다 효과적으로 감소시킴을 보여줍니다.
• Practical TEON algorithm: 근사 SVD(예: 무작위 파워 이터레이션)를 사용한 계산적으로 실행 가능한 구현을 도출하고, 표준 딥러닝 툴킷에서 작동함을 입증합니다.
• Extensive empirical validation: GPT‑style(130 M–774 M) 및 LLaMA‑style(60 M–1 B) 모델에 대한 벤치마크에서 모든 규모에서 훈련/검증 퍼플렉시티를 낮추는 성과를 달성했습니다.
• Robustness analysis: 다양한 저랭크 SVD 근사에서도 TEON이 효과를 유지함을 보여주어 대규모 분산 학습에 적합합니다.

방법론

1. Gradient Tensor Construction – 각 층의 그래디언트 행렬 $G_\ell$을 개별적으로 다루는 대신, TEON은 이를 3‑D 텐서 $\mathcal{G} \in \mathbb{R}^{L \times d_{\text{in}} \times d_{\text{out}}}$ (층, 입력 차원, 출력 차원)으로 쌓는다.
2. Tensor Orthogonalization – TEON은 $\mathcal{Q}^\top \mathcal{Q} = I$가 되도록 $\mathcal{G}$에 대한 정규 직교 기저 $\mathcal{Q}$를 찾는다. 이는 고차원 특이값 분해(HOSVD) 또는 더 저렴한 무작위 근사 방법을 적용함으로써 달성된다.
3. Update Rule – 옵티마이저는 원시 그래디언트(또는 모멘텀)를 정규 직교 기저에 투사하여 정규화된 그래디언트 $\tilde{\mathcal{G}} = \mathcal{Q}\mathcal{Q}^\top \mathcal{G}$를 얻는다. 이후 표준 Adam‑유사 단계 크기가 적용된다.
4. Approximation Strategies – 비용을 관리 가능하게 유지하기 위해, 저자들은 다음을 실험한다:
   • Randomized power iteration (few iterations) for each mode of the tensor. (텐서의 각 모드에 대해 (몇 번의 반복) 무작위 파워 이터레이션)
   • Low‑rank truncation (keeping only top‑k singular components). (저‑랭크 절단 (상위 k개의 특이 성분만 유지))
   • Layer‑wise fallback (reverting to Muon when tensor cost exceeds a threshold). (층별 폴백 (텐서 비용이 임계값을 초과하면 Muon으로 되돌림))

결과 알고리즘은 약간의 오버헤드(≈ 5–10 % 추가 연산)를 추가하면서도 일반적인 Adam‑스타일 옵티마이저의 메모리 사용량을 유지한다.

결과 및 발견

• 일관된 향상: TEON은 모든 모델 크기에서 Muon보다 약 0.5–0.8 포인트, Adam보다 약 1.0–1.5 포인트 퍼플렉시티를 개선합니다.
• 확장성: 이 이점은 모델 크기가 커짐에 따라 증가하며, 이는 층 간 그래디언트 상관관계가 큰 네트워크에서 더 두드러진다는 것을 나타냅니다.
• 견고성: 다양한 SVD 근사(rank‑k = 5, 10, 20)를 사용한 실험에서 성능 손실이 거의 없음을 보여주며, 저렴한 근사만으로도 충분함을 확인합니다.

실용적 시사점

• LLM 사전 학습에서 더 빠른 수렴: 개발자는 더 적은 학습 단계로 동일하거나 더 나은 모델 품질을 달성할 수 있어 GPU/TPU 클러스터 비용을 절감할 수 있습니다.
• 즉시 교체 가능: TEON의 API는 Adam/Muon을 그대로 모방하므로 기존 PyTorch 또는 JAX 파이프라인에 통합할 때 몇 줄의 코드만 추가하면 됩니다.
• 저정밀 환경에서 더 나은 안정성: 직교화 단계가 그래디언트 폭발/소실을 완화시켜 혼합 정밀도(FP16/BF16) 학습을 보다 신뢰할 수 있게 합니다.
• 다운스트림 파인튜닝 가능성: TEON이 더 잘 초기화된 가중치 공간을 제공하므로 도메인 특화 데이터에 대한 다운스트림 파인튜닝이 더 빠르게 수렴하고 높은 정확도에 도달할 수 있습니다.

제한 사항 및 향후 작업

• Computational overhead: 비록 적지만, 추가 5–10 % 계산량은 초대규모 실행(수십억 파라미터)에서 여전히 눈에 띌 수 있습니다.
• Memory footprint of the gradient tensor: 모든 레이어 그래디언트를 쌓으면 매우 깊은 모델에서 메모리 부담이 커질 수 있습니다; 저자들은 이를 해결하기 위해 스트리밍 또는 블록 단위 직교화를 제안합니다.
• Theoretical analysis limited to convex surrogates: 수렴 증명은 국부적으로 볼록한 근사를 전제로 합니다; 트랜스포머의 전체 비볼록 영역에 대한 보장을 확장하는 것은 아직 열려 있습니다.

향후 방향

• 오버헤드를 더 줄이기 위한 적응형 랭크 선택 탐색.
• TEON을 다른 2차 기법(예: K‑FAC)과 결합하여 더 빠른 수렴 달성.
• 텐서 직교화를 비전 트랜스포머나 디퓨전 모델과 같은 다른 분야에 적용.

저자

• Ruijie Zhang
• Yequan Zhao
• Ziyue Liu
• Zhengyang Wang
• Dongyang Li
• Yupeng Su
• Sijia Liu
• Zheng Zhang

논문 정보

• arXiv ID: 2601.23261v1
• 분류: cs.LG, cs.AI
• 발행일: 2026년 1월 30일
• PDF: Download PDF