SWAN：带有归一化与白化的 SGD 实现无状态 LLM 训练
机器学习
2025-02-24 v3 人工智能
摘要
诸如 Adam (Kingma & Ba, 2015) 等自适应优化器一直是大型语言模型成功的关键。然而，它们通常需要在整个训练过程中维护优化器状态，这可能导致内存需求是模型占用的数倍。这种开销限制了可扩展性和计算效率。相比之下，随机梯度下降（SGD）是一种无状态优化器，因为它在训练期间不跟踪状态变量。因此，它实现了最佳内存效率。然而，它在 LLM 训练中的能力有限 (Zhao et al., 2024b)。在这项工作中，我们表明以无状态方式预处理 SGD 可以在 LLM 训练中实现与 Adam 优化器相同的性能，同时大幅降低内存成本。具体而言，我们提出使用归一化和白化对瞬时随机梯度进行预处理。我们表明归一化稳定了梯度分布，而白化抵消了损失景观的局部曲率。这产生了 SWAN（带有白化和归一化的 SGD），一种消除了存储任何优化器状态需求的随机优化器。从经验上看，SWAN 具有与 SGD 相同的内存占用，与 Adam 相比，端到端总内存减少了约 50%。在语言建模任务中，SWAN 展现出与 Adam 相当甚至更好的性能：在使用 350M 和 1.3B 参数预训练 LLaMA 模型时，SWAN 通过使用一半的 token 达到相同的评估困惑度，实现了 2 倍的加速。
引用
@article{arxiv.2412.13148,
title = {SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training},
author = {Chao Ma and Wenbo Gong and Meyer Scetbon and Edward Meeds},
journal= {arXiv preprint arXiv:2412.13148},
year = {2025}
}
备注
In v2 we have revised the related work, added more comprehensive citations, and clarified our key contributions