MSign: Un optimizador que previene la inestabilidad del entrenamiento en modelos de lenguaje grandes mediante la restauración del rango estable
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
Resumen
La inestabilidad en el entrenamiento sigue siendo un desafío crítico en el preentrenamiento de modelos de lenguaje grandes (LLM), manifestándose a menudo como explosiones súbitas del gradiente que desperdician recursos computacionales significativos. Estudiamos las fallas de entrenamiento en un modelo NanoGPT de 5M de parámetros escalado mediante μP, identificando dos fenómenos clave que preceden al colapso: (1) la rápida disminución del rango estable de la matriz de pesos (relación entre la norma de Frobenius al cuadrado y la norma espectral al cuadrado), y (2) la creciente alineación entre los jacobianos de capas adyacentes. Demostramos teóricamente que estas dos condiciones conjuntamente causan un crecimiento exponencial de la norma del gradiente con la profundidad de la red. Para romper este mecanismo de inestabilidad, proponemos MSign, un nuevo optimizador que aplica periódicamente operaciones de signo matricial para restaurar el rango estable. Los experimentos en modelos desde 5M hasta 3B de parámetros demuestran que MSign previene efectivamente las fallas de entrenamiento con una sobrecarga computacional inferior al 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%.