StoSignSGD：无偏结构随机性修复 SignSGD 用于训练 大语言模型

基于符号的优化算法（例如 SignSGD）因其在分布式学习和训练大型基础模型中的性能而受到关注。尽管 SignSGD 具有经验优势，但众所周知，它在非平滑目标上存在分歧，由于 ReLU、最大池和专家混合，非平滑目标无处不在。为了克服这一限制，我们提出了 StoSignSGD，一种将结构随机性注入符号运算符，同时保持无偏更新步骤的算法。在（在线）凸优化机制中，StoSignSGD 严格解决了 SignSGD 的不收敛问题，实现了与下限匹配的急剧收敛速度。对于更具挑战性的非凸非平滑优化，我们引入了包含先验定义的广义固定度量，证明 StoSignSGD 通过维度因素改进了最著名的复杂性界限。根据经验，StoSignSGD 在不同的 大语言模型 (LLM) 训练体系中都是稳定且高效的。在积极的低精度预训练中，涵盖 FP8 和要求更高的 FP4 体系（其中 AdamW 发生灾难性失败），StoSignSGD 保持稳定并始终表现最佳。与 FP8 下的既定基线相比，它实现了 1.44 倍至 2.14 倍的加速。在 4 位精度下，最大 OLMo2-370M 模型的下游精度比最强稳定基线提高了 1.13 个点，并且这种优势随着模型大小和数据规模的扩大而增长。当 微调 7B LLM 执行数学推理任务时，StoSignSGD 也比 AdamW 和 SignSGD 具有明显的优势。最后，为了解释它的工作原理，我们开发了一个符号转换框架，可以将任何通用优化器转换为无偏的、基于符号的对应优化器。使用此框架，我们分解了 StoSignSGD 的核心组件，并运行全面的消融研究来验证我们的设计选择。