提出MSign优化器，通过恢复权重矩阵稳定秩防止大模型训练崩溃
MSign: An Optimizer Preventing Training Instability in Large Language Models via Stable Rank Restoration
- 周期性对权重矩阵进行符号操作，恢复其稳定秩
- 在500万至30亿参数模型上完全避免训练崩溃
- 仅增加不足7%计算开销，适合大规模语言模型训练
大语言模型预训练中的训练不稳定性仍是关键挑战，常表现为突发的梯度爆炸，造成大量算力浪费。我们研究了通过μP扩展的500万参数NanoGPT模型的训练失败现象，发现崩溃前存在两个关键特征：(1) 权重矩阵稳定秩（平方Frobenius范数与平方谱范数之比）迅速下降；(2) 相邻层雅可比矩阵间对齐度上升。理论上证明，这两个条件共同导致梯度范数随网络深度呈指数增长。为打破该不稳定性机制，我们提出MSign优化器，周期性对权重矩阵施加矩阵符号操作以恢复稳定秩。在500万至30亿参数模型上的实验表明，MSign能有效防止训练失败，且计算开销低于7.0%。
原文摘要 · Abstract (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%.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。