新优化器MSign:通过稳定秩恢复机制防止大模型预训练崩溃

研究团队提出MSign优化器,周期性应用矩阵符号操作恢复权重矩阵稳定秩,理论证明可阻止梯度爆炸

AI解读:大模型预训练中最让人头疼的问题之一是训练中途突然梯度爆炸,导致大量算力白费。这篇论文发现,崩溃前有两个信号:权重矩阵的“稳定秩”快速下降,以及相邻层Jacobian越来越对齐。研究者在理论上证明这两个条件会共同导致梯度范数随网络深度指数增长。基于此,他们提出MSign优化器,做法是周期性对权重矩阵做矩阵符号操作,把稳定秩拉回来,从而打断崩溃机制。实验覆盖5M到3B参数的模型,计算开销不到7%。对做预训练的工程师来说,这意味着多了一个低成本的防护手段,但要注意论文实验主要在NanoGPT上验证,3B规模的具体训练效果和稳定性仍需更多独立测试,现阶段不必急着替换现有优化器。

arXiv论文提出名为MSign的新优化器,通过周期性矩阵符号操作恢复权重矩阵稳定秩,以防止大语言模型预训练中的训练崩溃。研究者在5M参数NanoGPT模型(经μP缩放)上观察到崩溃前的两个现象:权重矩阵稳定秩(平方Frobenius范数与平方谱范数之比)快速下降,以及相邻层Jacobian之间对齐度上升。论文理论上证明这两个条件共同导致梯度范数随网络深度呈指数增长。实验在5M至3B参数模型上进行,结果显示MSign能有效防止训练失败,计算开销低于7.0%。

信息来源