Mesa: A Memory-saving Training Framework for Transformers

要約

高性能トランスフォーマーの設計への関心が爆発的に高まっています。
トランスフォーマーはパフォーマンスを大幅に改善しましたが、そのようなネットワークのトレーニングは、バックプロパゲーション中の勾配計算に必要なすべての中間アクティベーションを格納するため、特に長いシーケンスの場合、非常にメモリ集約的です。
この目的のために、トランスフォーマー向けのメモリ節約トレーニング フレームワークである Mesa を紹介します。
具体的には、Mesa はフォワード パス中に正確なアクティベーションを使用し、低精度バージョンのアクティベーションを保存して、トレーニング中のメモリ消費を削減します。
低精度の活性化は、逆伝播中に逆量子化され、勾配が計算されます。
さらに、マルチヘッド自己注意層の不均一な活性化分布に対処するために、ヘッドごとの活性化量子化戦略を提案します。これは、各ヘッドの統計に基づいて活性化を量子化し、近似誤差を最小限に抑えます。
トレーニングの効率をさらに高めるために、推定を実行して量子化パラメーターを学習します。
さらに重要なことは、保存されたメモリをより大きなバッチ サイズの採用またはモデル サイズのスケールアップに再投資することで、制約のある計算リソースの下でパフォーマンスをさらに改善できる可能性があることです。
ImageNet、CIFAR-100、および ADE20K での広範な実験により、Mesa がトレーニング中に柔軟なメモリ節約 (最大 50%) を達成しながら、同等またはそれ以上のパフォーマンスを達成できることが実証されました。
コードは https://github.com/ziplab/Mesa で入手できます。

要約(オリジナル)

There has been an explosion of interest in designing high-performance Transformers. While Transformers have delivered significant performance improvements, training such networks is extremely memory intensive owing to storing all intermediate activations that are needed for gradient computation during backpropagation, especially for long sequences. To this end, we present Mesa, a memory-saving training framework for Transformers. Specifically, Mesa uses exact activations during forward pass while storing a low-precision version of activations to reduce memory consumption during training. The low-precision activations are then dequantized during back-propagation to compute gradients. Besides, to address the heterogeneous activation distributions in the multi-head self-attention layers, we propose a head-wise activation quantization strategy, which quantizes activations based on the statistics of each head to minimize the approximation error. To further boost training efficiency, we learn quantization parameters by running estimates. More importantly, by re-investing the saved memory in employing a larger batch size or scaling up model size, we may further improve the performance under constrained computational resources. Extensive experiments on ImageNet, CIFAR-100 and ADE20K demonstrate that Mesa can achieve flexible memory-savings (up to 50%) during training while achieving comparable or even better performance. Code is available at https://github.com/ziplab/Mesa.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google