MSign: An Optimizer Preventing Training Instability in Large Language Models via Stable Rank Restoration
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%.
6987609dbeecc443208d2375
Lianhai Ren
6987609dbeecc443208d2376
Yucheng Ding
6987609dbeecc443208d2377
Xiao Liu
claimed_verified
2026-02-09T08:31:02.000Z
63fb6e281b4b1bd4e7ffc5be
https://cdn-avatars.huggingface.co/v1/production/uploads/63fb6e281b4b1bd4e7ffc5be/aiRu_bulgnxvEMrjipGoQ.jpeg
Xiao Liu
lx865712528
user
lx865712528
6987609dbeecc443208d2378
Qianxiao Li
6987609dbeecc443208d2379
Peng Cheng
6987609dbeecc443208d237a
Yeyun Gong
MSign: An Optimizer Preventing Training Instability in Large Language Models via Stable Rank Restoration
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%.
Training instability in large language models is linked to weight matrix stable rank decline and Jacobian alignment, which MSign addresses through matrix sign operations to prevent gradient explosions.
large language model pretraining gradient explosions weight matrix stable rank Frobenius norm spectral norm Jacobian matrix sign operations optimizer