StoSignSGD：无偏结构随机性修复用于训练大语言模型的 SignSGD
机器学习
2026-04-20 v1 人工智能
最优化与控制
摘要
基于符号的优化算法，如 SignSGD，因其在分布式学习和训练大型基础模型中的卓越性能而备受关注。尽管在经验上表现优越，但已知 SignSGD 在非光滑目标上会发散，而由于 ReLU、最大池化和专家混合架构的存在，非光滑目标在现代机器学习中无处不在。为克服这一根本性局限，我们提出了 StoSignSGD 算法，该算法在保持无偏更新步骤的同时，向符号算子注入结构随机性。在（在线）凸优化领域，我们的理论分析表明，StoSignSGD 严格解决了 SignSGD 的不收敛问题，并达到了与下界相匹配的锐利收敛速率。对于更具挑战性的非凸非光滑优化，我们引入了涵盖先前定义的广义平稳测度，证明了 StoSignSGD 在已知复杂度界的基础上实现了维度因子的改进。在经验上，StoSignSGD 在多种大语言模型训练场景中表现出鲁棒的稳定性和卓越的效率。值得注意的是，在低精度 FP8 预训练中——AdamW 在此设置下会灾难性地失败——StoSignSGD 仍保持高度稳定，并相对于既定基线实现了 1.44 倍至 2.14 倍的显著加速。此外，在数学推理任务上微调 7B 大语言模型时，StoSignSGD 相较于 AdamW 和 SignSGD 均带来了显著的性能提升。最后，为剖析其成功背后的机制，我们开发了一个符号转换框架，能够将任何通用优化器转换为其无偏的、基于符号的对应版本。利用该框架，我们解构了 StoSignSGD 的核心组件，并进行了全面的消融研究，以实证验证我们的算法设计选择。
引用
@article{arxiv.2604.15416,
title = {StoSignSGD: Unbiased Structural Stochasticity Fixes SignSGD for Training Large Language Models},
author = {Dingzhi Yu and Rui Pan and Yuxing Liu and Tong Zhang},
journal= {arXiv preprint arXiv:2604.15416},
year = {2026}
}