突破記憶障礙:對比損失的近乎無限批量大小擴展Breaking the Memory Barrier: Near Infinite Batch Size Scaling for
Contrastive Loss
對比損失是一種強大的表示學習方法,其中通過提供更多負樣本來增強性能的較大批次大小,以更好地區分相似和不相似的數據。然而,批次大小的擴展受到 GPU 記憶體消耗呈二次增長的限制,主要是由於相似性矩陣的完全實例化。為了解決這個問題,我們提出了一種基於瓦片的計算策略,將對比損失計算分為任意小的塊,避免完全實例化相似性矩陣。此外,我們引入了多級瓦片策略,利用分佈式系統的分層結構,採用 GPU 層級的環形通信來優化同步,並在 CUDA 核心級別上使用融合內核以減少 I/O 開銷。實驗結果顯示,所提出的方法可以將批次大小擴展到前所未有的水平。例如,它使得可以使用 8 或 32 個 A800 80GB 進行對比訓練 CLIP-ViT-L/14 模型,批次大小為 4M 或 12M,而不會降低任何準確性。與最先進的節省記憶體解決方案相比,它實現了記憶體減少兩個數量級,同時保持可比擬的速度。代碼將公開提供。