SWAN: Preprocessing SGD Enables Adam-Level Performance On LLM Training With Significant Memory Reduction

像Adam这样的自适应优化器（Kingma & Ba, 2015）在大型语言模型的成功中发挥了核心作用。然而，它们在整个训练过程中维护了额外的移动平均状态，这导致所需的内存是模型本身的几倍。这种开销对可扩展性和计算效率造成了限制。另一方面，虽然随机梯度下降（SGD）在内存效率方面是最优的，但其在大型语言模型训练中的能力有限（Zhao et al., 2024b）。 为了解决这一困境，我们证明了预处理SGD足以达到与Adam相当的性能水平。具体来说，我们提出使用两个简单的操作符对瞬时随机梯度进行预处理：$\mathtt{GradNorm}$ 和 $\mathtt{GradWhitening}$。$\mathtt{GradNorm}$ 稳定了梯度分布，而 $\mathtt{GradWhitening}$ 则抵消了损失景观的局部曲率。这导致了SWAN（带有白化和归一化的SGD）的产生，这是一种无需存储任何累积状态变量的随机优化器。实证结果显示，SWAN 的内存占用与SGD相同，相比Adam总端到端内存减少了约50%。在语言建模任务中，SWAN 表现出了与Adam相同或甚至显著优于Adam的性能。特别是在预训练包含3.5亿和13亿参数的LLaMa模型时，SWAN通过在不到一半的令牌数量下达到相同的评估困惑度，实现了2倍的速度提升。

该论文试图解决在大规模语言模型训练中，使用Adam等自适应优化器导致的内存消耗过大的问题，同时保留或提升训练性能。这是一个在深度学习社区广泛讨论的问题，尤其是在处理大型模型时。

论文提出了一种新的优化器SWAN（SGD with Whitening And Normalization），通过预处理梯度来达到与Adam相当的性能，而无需存储额外的状态变量。具体来说，SWAN使用了两个简单的操作：GradNorm和GradWhitening，分别用于稳定梯度分布和对抗损失景观的局部曲率。这一方法不仅大幅减少了内存需求，还提高了训练速度。

论文通过在LLaMa模型上的实验证明，SWAN在350M和1.3B参数的模型上达到了与Adam相同的评估困惑度，但所需的时间和计算资源显著减少。此外，SWAN的内存占用与传统的SGD相同，比Adam减少了约50%的总端到端内存。论文还提到，SWAN在预训练过程中实现了2倍的速度提升。这些结果表明，SWAN在提高训练效率和降低资源消耗方面具有显著优势。

近期在这个领域中，还有一些相关的研究，例如： - Zhao et al. (2024a) 提出了一种基于动量的优化器，旨在平衡内存消耗和训练性能。 - Liu et al. (2023) 探讨了使用二阶信息来改进SGD的方法。 - Chen et al. (2023) 研究了如何通过稀疏更新来减少优化器的内存开销。 - Kingma & Ba (2015) 的Adam优化器是自适应优化器的代表，广泛应用于各种深度学习任务。