SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training

要約

Adam（Kingma＆BA、2015）などの適応オプティマイザーは、大規模な言語モデルの成功の中心となっています。
ただし、多くの場合、トレーニング全体でオプティマイザー状態を維持する必要があります。これにより、モデルフットプリントの数倍のメモリ要件が発生する可能性があります。
このオーバーヘッドは、スケーラビリティと計算効率に制約を課します。
対照的に、確率勾配降下（SGD）は、トレーニング中に状態変数を追跡しないため、ステートレスオプティマイザーです。
その結果、最適なメモリ効率を達成します。
ただし、LLMトレーニングにおけるその能力は限られています（Zhao et al。、2024b）。
この作業では、前処理SGDがステートレスの方法で行われることが、LLMトレーニングのAdam Optimizerと同じパフォーマンスを達成しながら、メモリコストを大幅に削減できることを示しています。
具体的には、正規化とホワイトニングを使用して、瞬時の確率勾配を前処理することを提案します。
正規化が勾配分布を安定させ、ホワイトニングが損失状況の局所曲率に対抗することを示します。
これにより、SWAN（ホワイトニングと正規化を備えたSGD）は、オプティマイザーの状態を保存する必要性を排除する確率的オプティマイザーです。
経験的には、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 often require to maintain optimizer states throughout training, which can result in memory requirements several times greater than the model footprint. This overhead imposes constraints on scalability and computational efficiency. Stochastic Gradient Descent (SGD), in contrast, is a stateless optimizer, as it does not track state variables during training. Consequently, it achieves optimal memory efficiency. However, its capability in LLM training is limited (Zhao et al., 2024b). In this work, we show that pre-processing SGD in a stateless manner can achieve the same performance as the Adam optimizer for LLM training, while drastically reducing the memory cost. Specifically, we propose to pre-process the instantaneous stochastic gradients using normalization and whitening. We show that normalization stabilizes gradient distributions, and whitening counteracts the local curvature of the loss landscape. This results in SWAN (SGD with Whitening And Normalization), a stochastic optimizer that eliminates the need to store any optimizer states. 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 comparable or even better performance than Adam: when pre-training the LLaMA model with 350M and 1.3B parameters, SWAN achieves a 2x speedup by reaching the same evaluation perplexity using half as many tokens.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google