データセットと深層学習モデルのサイズが増大しているため、より高速でメモリ効率の高いトレーニングが重要になっています。可逆トランスフォーマーは、非常にメモリ効率の高いトレーニングのためのエキサイティングな新しい方法として最近導入されましたが、バックプロパゲーション段階での活性化再計算という追加の計算オーバーヘッドが伴います。我々は、バックプロパゲーションフェーズでの勾配計算自体を使用して、可逆トレーニングにおける追加の活性化再計算オーバーヘッドを並列化する、高速並列可逆バックプロパゲーションアルゴリズムであるPaRepropを紹介します。モデル ファミリ (ViT、MViT、Swin、RoBERTa)、データ モダリティ (Vision および NLP)、モデル サイズ (小規模から巨大まで)、トレーニング バッチ サイズにわたる広範なベンチマークを通じて、提案された PaReprop アルゴリズムの有効性を実証します。私たちの経験的結果は、PaReprop がバニラの可逆トレーニングよりも最大 20% 高いトレーニング スループットを達成し、可逆トレーニングにおけるアクティベーションの再計算による 25% 低いスループットの理論上のオーバーヘッドを大幅に軽減することを示しています。プロジェクトページ: https://tylerzhu.com/pareprop
The growing size of datasets and deep learning models has made faster and memory-efficient training crucial. Reversible transformers have recently been introduced as an exciting new method for extremely memory-efficient training, but they come with an additional computation overhead of activation re-computation in the backpropagation phase. We present PaReprop, a fast Parallelized Reversible Backpropagation algorithm that parallelizes the additional activation re-computation overhead in reversible training with the gradient computation itself in backpropagation phase. We demonstrate the effectiveness of the proposed PaReprop algorithm through extensive benchmarking across model families (ViT, MViT, Swin and RoBERTa), data modalities (Vision & NLP), model sizes (from small to giant), and training batch sizes. Our empirical results show that PaReprop achieves up to 20% higher training throughput than vanilla reversible training, largely mitigating the theoretical overhead of 25% lower throughput from activation recomputation in reversible training. Project page: https://tylerzhu.com/pareprop.