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的灵感来源。

文章预览

↑ 点击 蓝字  关注极市平台 作者丨藤原豆腐皮儿@知乎（已授权） 来源丨https://zhuanlan.zhihu.com/p/1681887214 编辑丨极市平台 极市导读   本文介绍了一种新的Contrastive Loss实现方式——Inf-CL，它通过分块计算策略，在单台A800机器上将batch size扩展到4M，几乎实现了Contrastive Loss batch size的无限扩展，突破了以往认为增加batch size会导致显存不足的限制。   >> 加入极市CV技术交流群，走在计算机视觉的最前沿 TL;DR 本文提出了一种Contrastive Loss的实现方式( Inf-CL )，通过分块计算策略，我们在单台A800机器上就能把 batch size 扩展到 4M 。不严谨地说，该方案突破了以前公知的 ”contrastive loss不能scaling batch size，显存会炸“ 的前提，实现了 Contrastive Loss 的 batch size 近乎无限的扩展。 中国人不骗中国人，以后对比损失实现就用Inf-CL！！ 对比学习有多炸不用多说，在 ………………………………