Breaking the Memory Barrier: Near Infinite Batch Size Scaling for Contrastive Loss

要約

対比損失は表現学習の強力なアプローチであり、バッチ サイズが大きくなると、より多くの負のサンプルが提供され、類似データと類似データをより適切に区別できるため、パフォーマンスが向上します。
ただし、バッチ サイズのスケーリングは、主に類似性行列の完全なインスタンス化による GPU メモリ消費の二次的な増加によって制限されます。
これに対処するために、類似性行列の完全な具体化を回避し、コントラスト損失計算を任意の小さなブロックに分割するタイルベースの計算戦略を提案します。
さらに、分散システムの階層構造を活用するマルチレベル タイル戦略を導入し、GPU レベルでリングベースの通信を採用して同期を最適化し、CUDA コア レベルで融合カーネルを使用して I/O オーバーヘッドを削減します。
実験結果は、提案された方法がバッチサイズを前例のないレベルに拡張することを示しています。
たとえば、精度を犠牲にすることなく、8 個または 32 個の A800 80GB を使用して、バッチ サイズ 4M または 12M の CLIP-ViT-L/14 モデルの対比トレーニングが可能になります。
SOTA のメモリ効率の高いソリューションと比較して、同等の速度を維持しながらメモリを 2 桁削減できます。
コードは公開されます。

要約(オリジナル)

Contrastive loss is a powerful approach for representation learning, where larger batch sizes enhance performance by providing more negative samples to better distinguish between similar and dissimilar data. However, scaling batch sizes is constrained by the quadratic growth in GPU memory consumption, primarily due to the full instantiation of the similarity matrix. To address this, we propose a tile-based computation strategy that partitions the contrastive loss calculation into arbitrary small blocks, avoiding full materialization of the similarity matrix. Furthermore, we introduce a multi-level tiling strategy to leverage the hierarchical structure of distributed systems, employing ring-based communication at the GPU level to optimize synchronization and fused kernels at the CUDA core level to reduce I/O overhead. Experimental results show that the proposed method scales batch sizes to unprecedented levels. For instance, it enables contrastive training of a CLIP-ViT-L/14 model with a batch size of 4M or 12M using 8 or 32 A800 80GB without sacrificing any accuracy. Compared to SOTA memory-efficient solutions, it achieves a two-order-of-magnitude reduction in memory while maintaining comparable speed. The code will be made publicly available.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google