StoSignSGD: Unbiased Structural Stochasticity Fixes SignSGD for Training Large Language Models
作者: Dingzhi Yu, Rui Pan, Yuxing Liu, Tong Zhang
分类: cs.LG, cs.AI, math.OC
发布日期: 2026-04-16
💡 一句话要点
提出StoSignSGD，通过结构随机性解决SignSGD在大模型训练中非光滑目标上的发散问题
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: SignSGD 随机优化 大型语言模型 非光滑优化 低精度训练
📋 核心要点
- SignSGD在训练大模型时表现出优势，但其在非光滑目标函数上的发散问题限制了其应用。
- StoSignSGD通过引入结构随机性到符号算子中，并保持无偏更新，解决了SignSGD的收敛性问题。
- 实验表明，StoSignSGD在FP8预训练中稳定高效，并在LLM微调中优于AdamW和SignSGD。
📝 摘要（中文）
基于符号的优化算法，如SignSGD，在分布式学习和大型模型训练中表现出色。然而，SignSGD在非光滑目标上会发散，这是由于ReLU、max-pool和混合专家模型在现代机器学习中普遍存在。为了克服这一局限，我们提出了StoSignSGD，该算法将结构随机性注入到符号算子中，同时保持无偏更新。在凸优化领域，理论分析表明StoSignSGD解决了SignSGD的非收敛问题，实现了与下界匹配的收敛速度。对于更具挑战性的非凸非光滑优化，我们引入了广义平稳测度，证明StoSignSGD改进了已知的复杂度界限。实验表明，StoSignSGD在各种大型语言模型（LLM）训练中表现出鲁棒的稳定性和卓越的效率。在低精度FP8预训练中，StoSignSGD保持高度稳定，并相对于基线实现了1.44倍至2.14倍的加速。在数学推理任务上微调7B LLM时，StoSignSGD优于AdamW和SignSGD。此外，我们开发了一个符号转换框架，将任何通用优化器转换为其无偏、基于符号的对应物，从而剖析StoSignSGD成功的机制，并进行了全面的消融研究。
🔬 方法详解
问题定义：SignSGD虽然在分布式学习和训练大型模型方面表现出优势，但它在非光滑目标函数上存在发散的问题。现代机器学习中，由于ReLU、Max-Pooling和混合专家模型等非光滑激活函数的广泛使用，使得目标函数通常是非光滑的，这限制了SignSGD的应用。
核心思路：StoSignSGD的核心思路是在SignSGD的符号算子中注入结构随机性，同时保证更新步骤的无偏性。通过引入随机性，可以有效地平滑非光滑目标函数，从而避免SignSGD在非光滑区域的发散问题。保持无偏性则确保了算法的收敛性。
技术框架：StoSignSGD的整体框架是在SignSGD的基础上，对梯度符号化过程进行改进。具体来说，对于每个梯度分量，StoSignSGD不是直接取其符号，而是以一定的概率随机地选择符号。这个概率取决于梯度分量的大小，使得梯度较大的分量更有可能保持其原始符号，而梯度较小的分量则更有可能被随机翻转。这种结构化的随机性注入有效地平滑了目标函数。
关键创新：StoSignSGD的关键创新在于其结构随机性的引入方式。与简单的随机符号翻转不同，StoSignSGD的随机性是结构化的，即它依赖于梯度分量的大小。这种结构化的随机性既能有效地平滑目标函数，又能保证算法的收敛性。此外，论文还提出了一个符号转换框架，可以将任何通用优化器转换为其无偏、基于符号的对应物。
关键设计：StoSignSGD的关键设计在于随机符号翻转的概率计算方式。论文中具体给出了概率的计算公式，该公式保证了更新步骤的无偏性。此外，论文还对StoSignSGD的收敛性进行了理论分析，并给出了在凸优化和非凸非光滑优化情况下的收敛速度。在实验中，论文使用了不同的学习率和batch size等参数，并对StoSignSGD的性能进行了全面的评估。
🖼️ 关键图片
📊 实验亮点
StoSignSGD在低精度FP8预训练中表现出极高的稳定性，并且相对于基线方法实现了1.44倍到2.14倍的加速。在7B LLM的数学推理任务微调中，StoSignSGD也显著优于AdamW和SignSGD，证明了其在实际应用中的优越性。
🎯 应用场景
StoSignSGD在大型语言模型训练中具有广泛的应用前景，尤其是在低精度训练和分布式训练场景下。它可以降低训练成本，提高训练效率，并提升模型的性能。此外，该算法还可以应用于其他类型的深度学习模型，例如计算机视觉和语音识别模型。
📄 摘要（原文）
Sign-based optimization algorithms, such as SignSGD, have garnered significant attention for their remarkable performance in distributed learning and training large foundation models. Despite their empirical superiority, SignSGD is known to diverge on non-smooth objectives, which are ubiquitous in modern machine learning due to ReLUs, max-pools, and mixture-of-experts. To overcome this fundamental limitation, we propose \textbf{StoSignSGD}, an algorithm that injects structural stochasticity into the sign operator while maintaining an unbiased update step. In the regime of (online) convex optimization, our theoretical analysis shows that StoSignSGD rigorously resolves the non-convergence issues of SignSGD, achieving a sharp convergence rate matching the lower bound. For the more challenging non-convex non-smooth optimization, we introduce generalized stationary measures that encompass prior definitions, proving that StoSignSGD improves upon the best-known complexity bounds by dimensional factors. Empirically, StoSignSGD exhibits robust stability and superior efficiency across diverse large language model (LLM) training regimes. Notably, in low-precision FP8 pretraining -- a setting where AdamW fails catastrophically -- StoSignSGD remains highly stable and yields a remarkable 1.44$\times$ to 2.14$\times$ speedup relative to established baselines. Furthermore, when fine-tuning 7B LLMs on mathematical reasoning tasks, StoSignSGD delivers substantial performance gains over both AdamW and SignSGD. Finally, to dissect the mechanisms driving its success, we develop a sign conversion framework capable of transforming any general optimizer into its unbiased, sign-based counterpart. Utilizing this framework, we deconstruct the core components of StoSignSGD and present a comprehensive ablation study to empirically validate our algorithmic design choices.