MSign: Um Otimizador que Previne a Instabilidade no Treinamento de Grandes Modelos de Linguagem por meio da Restauração do Posto Estável
MSign: An Optimizer Preventing Training Instability in Large Language Models via Stable Rank Restoration
February 2, 2026
Autores: Lianhai Ren, Yucheng Ding, Xiao Liu, Qianxiao Li, Peng Cheng, Yeyun Gong
cs.AI
Resumo
A instabilidade no treinamento continua sendo um desafio crítico no pré-treinamento de grandes modelos de linguagem (LLMs), manifestando-se frequentemente como explosões súbitas de gradiente que desperdiçam recursos computacionais significativos. Estudamos falhas de treinamento num modelo NanoGPT de 5M de parâmetros dimensionado via μP, identificando dois fenómenos-chave que precedem o colapso: (1) declínio rápido no posto estável da matriz de pesos (rácio da norma de Frobenius ao quadrado para a norma espectral ao quadrado), e (2) aumento do alinhamento entre os Jacobianos de camadas adjacentes. Provamos teoricamente que estas duas condições conjuntamente causam um crescimento exponencial da norma do gradiente com a profundidade da rede. Para quebrar este mecanismo de instabilidade, propomos o MSign, um novo otimizador que aplica periodicamente operações de sinal matricial para restaurar o posto estável. Experiências em modelos de 5M a 3B de parâmetros demonstram que o MSign previne eficazmente falhas de treinamento com uma sobrecarga computacional inferior a 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%.