Adam-mini: Use Fewer Learning Rates To Gain More

要約

私たちは、Adam-mini を提案します。Adam-mini は、メモリ使用量を 45% ～ 50% 削減しながら、AdamW と同等以上のパフォーマンスを実現するオプティマイザです。
Adam-mini は、Adam の学習率の数を削減することでメモリを削減します。 $1/\sqrt{v}$ を使用して各パラメーターに個別の学習率を割り当てる代わりに、Adam-mini はプリセット内の $v$ の平均を使用します。
パラメータ ブロックをそのブロックの学習率として定義します。
このようなデザインは 2 つの経験的な発見に基づいています。
まず、トランスフォーマーのヘッセ行列は、さまざまなサイズの密なサブブロックを含むブロックに近い対角構造を示します。
第 2 に、これらの高密度のサブブロックのそれぞれについて、それを検索するのに十分なリソースが利用可能であれば、Adam を上回るパフォーマンスを発揮できる単一の高品質な学習率が存在します。
Adam-mini は、これらの良好な学習率を見つけて、Adam の $\geq$ 90% $v$ を管理するための 1 つの費用対効果の高い方法を提供します。
経験的に、Adam-mini は、事前トレーニング、教師付き微調整、および RLHF について、125M から 7B までのサイズのさまざまな言語モデルで AdamW と同等以上のパフォーマンスを発揮することを確認しています。
Adam-mini のメモリ使用量の削減により、GPU と CPU 間の通信オーバーヘッドも軽減され、スループットが向上します。
たとえば、2x A800-80GB GPU で Llama2-7B を事前トレーニングする場合、Adam-mini は AdamW よりも 49.6% 高いスループットを達成し、事前トレーニングの実時間を 33% 節約します。

要約(オリジナル)

We propose Adam-mini, an optimizer that achieves on-par or better performance than AdamW with 45% to 50% less memory footprint. Adam-mini reduces memory by cutting down the number of learning rates in Adam: Instead of assigning an individual learning rate for each parameter using $1/\sqrt{v}$, Adam-mini uses the average of $v$ within a pre-defined parameter block as the learning rate for that block. Such a design is inspired by two empirical findings. First, the Hessian of Transformers exhibits a near-block diagonal structure with different sizes of dense sub-blocks. Second, for each of these dense sub-blocks, there exists a single high-quality learning rate that can outperform Adam, provided that sufficient resources are available to search it out. Adam-mini provides one cost-effective way to find these good learning rates and manage to cut down $\geq$ 90% $v$ in Adam. Empirically, we verify that Adam-mini performs on par or better than AdamW on various language models sized from 125M to 7B for pre-training, supervised fine-tuning, and RLHF. The reduced memory footprint of Adam-mini also alleviates communication overheads among GPUs and CPUs, thereby increasing throughput. For instance, Adam-mini achieves 49.6% higher throughput than AdamW when pre-training Llama2-7B on 2x A800-80GB GPUs, which saves 33% wall-clock time for pre-training.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google