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

要約

Adam (Kingma & Ba、2015) などの適応オプティマイザーは、大規模な言語モデルの成功の中心となってきました。
ただし、トレーニング全体を通じて追加の移動平均状態を維持するため、メモリ要件がモデルの数倍になります。
このオーバーヘッドにより、スケーラビリティと計算効率に制約が課されます。
一方、確率的勾配降下法 (SGD) はメモリ効率の点では最適ですが、LLM トレーニングにおける能力は限られています (Zhao et al., 2024b)。
このジレンマに対処するために、LLM で Adam レベルのパフォーマンスを達成するには、SGD の前処理で十分であることを示します。
具体的には、$\mathtt{GradNorm}$ と $\mathtt{GradWhitening}$ という 2 つの単純な演算子を使用して瞬間的な確率的勾配を前処理することを提案します。
$\mathtt{GradNorm}$ は勾配分布を安定化し、$\mathtt{GradWhitening}$ は損失ランドスケープの局所的な曲率をそれぞれ打ち消します。
これにより、累積状態変数を保存する必要がなくなる確率的オプティマイザーである SWAN (SGD with Whitening And Normalization) が実現します。
経験的には、SWAN のメモリ使用量は SGD と同じであり、Adam と比較してエンドツーエンドの総メモリ量が $\約 50\%$ 削減されます。
言語モデリング タスクでは、SWAN は Adam と同等か、さらには大幅な改善を示しています。
具体的には、350M および 1.3B のパラメーターで LLaMa モデルを事前トレーニングする場合、SWAN は、半分未満のトークンで同じ評価の複雑さに到達することで 2 倍の高速化を達成します。

要約(オリジナル)

Adaptive optimizers such as Adam (Kingma & Ba, 2015) have been central to the success of large language models. However, they maintain additional moving average states throughout training, which results in memory requirements several times greater than the model. This overhead imposes constraints on scalability and computational efficiency. On the other hand, while stochastic gradient descent (SGD) is optimal in terms of memory efficiency, their capability in LLM training is limited (Zhao et al., 2024b). To address this dilemma, we show that pre-processing SGD is sufficient to reach Adam-level performance on LLMs. Specifically, we propose to preprocess the instantaneous stochastic gradients with two simple operators: $\mathtt{GradNorm}$ and $\mathtt{GradWhitening}$. $\mathtt{GradNorm}$ stabilizes gradient distributions, and $\mathtt{GradWhitening}$ counteracts the local curvature of the loss landscape, respectively. This results in SWAN (SGD with Whitening And Normalization), a stochastic optimizer that eliminates the need to store any accumulative state variables. Empirically, SWAN has the same memory footprint as SGD, achieving $\approx 50\%$ reduction on total end-to-end memory compared to Adam. In language modeling tasks, SWAN demonstrates the same or even a substantial improvement over Adam. Specifically, when pre-training the LLaMa model with 350M and 1.3B parameters, SWAN achieves a 2x speedup by reaching the same evaluation perplexity in less than half tokens seen.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google