MicroAdam: Accurate Adaptive Optimization with Low Space Overhead and Provable Convergence

要約

私たちは、理論的な収束保証を維持しながら、特にメモリのオーバーヘッドを最小限に抑える MicroAdam と呼ばれる Adam オプティマイザーの新しいバリアントを提案します。
これは、オプティマイザ状態に供給される前に勾配情報を圧縮することで実現され、それによってメモリ フットプリントが大幅に削減されます。
私たちは、分散最適化による古典的な \emph{エラー フィードバック} メカニズムの新しいインスタンスを介して、結果として生じる圧縮エラーを制御します。このメカニズムでは、*エラー訂正情報自体が圧縮され*、実用的なメモリの増加が可能になります。
私たちは、結果として得られたアプローチが、優れた実用的なパフォーマンスを提供しながら、AMSGrad のそれに匹敵する理論的収束保証を維持していることを証明します。
具体的には、MicroAdam が GPU 上で効率的に実装できることを示します。MicroAdam は、100 万スケール (BERT) モデルと 10 億スケール (LLaMA) モデルの両方で、非圧縮 Adam ベースラインに匹敵する実用的な収束を提供し、メモリ使用量が低く、同様の実行が可能です。
時間。
私たちのコードは https://github.com/IST-DASLab/MicroAdam で入手できます。

要約(オリジナル)

We propose a new variant of the Adam optimizer called MicroAdam that specifically minimizes memory overheads, while maintaining theoretical convergence guarantees. We achieve this by compressing the gradient information before it is fed into the optimizer state, thereby reducing its memory footprint significantly. We control the resulting compression error via a novel instance of the classical \emph{error feedback} mechanism from distributed optimization in which *the error correction information is itself compressed* to allow for practical memory gains. We prove that the resulting approach maintains theoretical convergence guarantees competitive to those of AMSGrad, while providing good practical performance. Specifically, we show that MicroAdam can be implemented efficiently on GPUs: on both million-scale (BERT) and billion-scale (LLaMA) models, MicroAdam provides practical convergence competitive to that of the uncompressed Adam baseline, with lower memory usage and similar running time. Our code is available at https://github.com/IST-DASLab/MicroAdam.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google