Inf-CL: 把 Contrastive Loss 的 Batch Size 冲到100M!

主要观点总结

本文介绍了一种新的Contrastive Loss实现方式——Inf-CL，它通过分块计算策略，在单台A800机器上将batch size扩展到4M，几乎实现了Contrastive Loss batch size的无限扩展，突破了以往认为增加batch size会导致显存不足的限制。文章详细描述了Inf-CL的方法原理、实验结果的对比，包括显存节省度、速度和精度。

关键观点总结

关键观点1: 对比学习的重要性和限制

对比学习在多个领域如图文检索、图像自监督学习、文本检索中占据重要地位。但增大batch size或负样本会导致GPU显存爆炸，成为该领域的一个难题。

关键观点2: Inf-CL方法介绍

Inf-CL采用分块计算策略，通过减少显存占用实现大batch size的对比损失计算。包括前向传播和反向传播的过程以及Multi-Level Tiling策略。

关键观点3: 实验结果

实验结果显示，Inf-CL在降低显存占用的同时，只引入了极少的时间开销。并且在降低显存占用后，仍然保持了较高的训练速度和精度。

关键观点4: 相关工作与灵感来源

介绍了与本文相关的工作，如Gradient Cache、Flash Attention和Ring Attention等，这些工作为Inf-CL的灵感来源。