MSign: 안정적 랭크 복원을 통한 대규모 언어 모델의 학습 불안정성 방지 최적화 알고리즘

초록

대규모 언어 모델(LLM) 사전 학습에서 훈련 불안정성은 여전히 중요한 과제로 남아 있으며, 종종 갑작스러운 그래디언트 폭발로 나타나 상당한 계산 자원을 낭비하게 됩니다. 본 연구에서는 μP를 통해 확장된 5M 매개변수 NanoGPT 모델에서의 훈련 실패를 분석하며, 붕괴에 선행하는 두 가지 주요 현상을 확인했습니다: (1) 가중치 행렬의 안정성 계수(프로베니우스 놈의 제곱과 스펙트럴 놈의 제곱의 비율)의 급격한 감소, (2) 인접 계층 야코비안 간의 정렬도 증가. 우리는 이 두 조건이 함께 네트워크 깊이에 따른 기하급수적 그래디언트 놈 성장을 유발한다는 것을 이론적으로 증명합니다. 이러한 불안정성 메커니즘을 차단하기 위해, 안정성 계수를 복원하기 위해 주기적으로 행렬 부호 연산을 적용하는 새로운 옵티마이저인 MSign을 제안합니다. 5M에서 3B 매개변수에 이르는 모델을 대상으로 한 실험에서 MSign이 7.0% 미만의 계산 오버헤드로 훈련 실패를 효과적으로 방지함을 입증했습니다.

English

Training instability remains a critical challenge in large language model (LLM) pretraining, often manifesting as sudden gradient explosions that waste significant computational resources. We study training failures in a 5M-parameter NanoGPT model scaled via μP, identifying two key phenomena preceding collapse: (1) rapid decline in weight matrix stable rank (ratio of squared Frobenius norm to squared spectral norm), and (2) increasing alignment between adjacent layer Jacobians. We prove theoretically that these two conditions jointly cause exponential gradient norm growth with network depth. To break this instability mechanism, we propose MSign, a new optimizer that periodically applies matrix sign operations to restore stable rank. Experiments on models from 5M to 3B parameters demonstrate that MSign effectively prevents training failures with a computational overhead of less than 7.0%.

MSign: 안정적 랭크 복원을 통한 대규모 언어 모델의 학습 불안정성 방지 최적화 알고리즘

MSign: An Optimizer Preventing Training Instability in Large Language Models via Stable Rank Restoration

초록

Support