打破内存壁垒：对比损失的近无限批量大小扩展
计算机视觉与模式识别
2024-10-23 v1
摘要
对比损失是一种强大的表示学习方法，其中更大的批量大小通过提供更多负样本来更好地区分相似和不同数据，从而提升性能。然而，扩展批量大小受到GPU内存消耗二次增长的制约，这主要是由于相似性矩阵的完全实例化。为了解决这个问题，我们提出了一种基于分块的计算策略，该策略将对比损失计算划分为任意小的块，避免了相似性矩阵的完全物化。此外，我们引入了一种多级分块策略，以利用分布式系统的层次结构，在GPU级别使用基于环的通信来优化同步，在CUDA核心级别使用融合内核来减少I/O开销。实验结果表明，所提出的方法将批量大小扩展到前所未有的水平。例如，它使得使用8或32个A800 80GB GPU，在不牺牲任何精度的情况下，对CLIP-ViT-L/14模型进行批量大小为4M或12M的对比训练成为可能。与最先进的内存高效解决方案相比，它在保持相当速度的同时，实现了两个数量级的内存减少。代码将公开提供。
引用
@article{arxiv.2410.17243,
title = {Breaking the Memory Barrier: Near Infinite Batch Size Scaling for Contrastive Loss},
author = {Zesen Cheng and Hang Zhang and Kehan Li and Sicong Leng and Zhiqiang Hu and Fei Wu and Deli Zhao and Xin Li and Lidong Bing},
journal= {arXiv preprint arXiv:2410.17243},
year = {2024}
}