Adam-mini: Use Fewer Learning Rates To Gain More

要約

私たちは、50% 少ないメモリ使用量で AdamW と同等以上のパフォーマンスを実現するオプティマイザーである Adam-mini を提案します。
Adam-mini は、Adam の学習率リソース (つまり、$1/\sqrt{v}$) を削減することでメモリを削減します。
ニューラルネットのヘッセ構造を調査することにより、Adam の $v$ が期待したほど効果的に最大限に機能していない可能性があることがわかりました。
$\geq$ $v$ のこれらの学習率の 99.9% は、(1) ヘッセ行列構造に関する新しい原則に従ってパラメータを慎重にブロックに分割すれば、無害に削除できることがわかります。
(2) 単一だが良好な学習率を各パラメータ ブロックに割り当てます。
次に、優れた学習率を見つけるための 1 つの簡単な方法を提供し、Adam-mini を提案します。
経験的に、Adam-mini は、事前トレーニング、教師あり微調整、および RLHF の 39M から 13B までのさまざまなサイズの言語モデルで、AdamW と同等かそれ以上のパフォーマンスを発揮することを確認しています。
Adam-mini のメモリ使用量の削減により、GPU 間の通信オーバーヘッドも軽減され、スループットが向上します。
たとえば、$2\times$ A800-80GB GPU で Llama 2-7B を事前トレーニングする場合、Adam-mini は AdamW よりも 49.6% 高いスループットを達成し、事前トレーニングの実時間を 33% 節約します。

要約(オリジナル)

We propose Adam-mini, an optimizer that achieves on par or better performance than AdamW with 50% less memory footprint. Adam-mini reduces memory by cutting down the learning rate resources in Adam (i.e., $1/\sqrt{v}$). By investigating the Hessian structure of neural nets, we find Adam’s $v$ might not function at its full potential as effectively as we expected. We find that $\geq$ 99.9% of these learning rates in $v$ could be harmlessly removed if we (1) carefully partition the parameters into blocks following our new principle on Hessian structure; (2) assign a single but good learning rate to each parameter block. We then provide one simple way to find good learning rates and propose Adam-mini. Empirically, we verify that Adam-mini performs on par or better than AdamW on various language models sized from 39M to 13B for pre-training, supervised fine-tuning, and RLHF. The reduced memory footprint of Adam-mini also alleviates communication overheads among GPUs, thereby increasing throughput. For instance, Adam-mini achieves 49.6% higher throughput than AdamW when pre-training Llama 2-7B on $2\times$ A800-80GB GPUs, which saves 33% wall-clock time for pre-training.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google