[论文解读] Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup

Luyu Gao, Yunyi Zhang+2 more|arXiv (Cornell University)|Jan 18, 2021

Domain Adaptation and Few-Shot Learning参考文献 16被引用 5

一句话总结

本文提出Gradient Cache，一种将编码器反向传播与批处理对比损失解耦的技术，通过缓存表示梯度，使在显存受限的GPU上实现任意大批次训练成为可能。该方法仅使用单张消费级GPU即在密集段落检索任务中达到最先进性能，其效果与使用多张高端GPU训练的模型相当，尽管运行时间多出20%，但显存使用量与批次大小无关，保持恒定。

ABSTRACT

Contrastive learning has been applied successfully to learn vector representations of text. Previous research demonstrated that learning high-quality representations benefits from batch-wise contrastive loss with a large number of negatives. In practice, the technique of in-batch negative is used, where for each example in a batch, other batch examples' positives will be taken as its negatives, avoiding encoding extra negatives. This, however, still conditions each example's loss on all batch examples and requires fitting the entire large batch into GPU memory. This paper introduces a gradient caching technique that decouples backpropagation between contrastive loss and the encoder, removing encoder backward pass data dependency along the batch dimension. As a result, gradients can be computed for one subset of the batch at a time, leading to almost constant memory usage.

研究动机与目标

• 为克服对比学习中因GPU显存限制而造成的批次大小受限问题，尤其针对硬件资源有限的研究人员。
• 实现在单张消费级GPU上训练大批次模型，这对高质量表征学习至关重要。
• 在不使用近似的情况下，保留大批次训练的精确梯度更新，确保模型性能与多GPU设置相当。
• 通过使大批次对比学习在学术界和资源受限环境中可行，实现技术普惠。

提出的方法

• 提出两阶段反向传播流程：首先计算损失到表示的梯度，然后计算表示到模型参数的梯度。
• 缓存损失关于表示的梯度（表示梯度缓存），以解耦编码器更新中的批次依赖性。
• 通过对表示的小批量进行子更新的方式执行编码器优化，实现在固定GPU显存内使用大有效批次大小。
• 通过为参数化相似度函数引入独立的距离梯度缓存，将该原理应用于深度距离函数。
• 利用缓存梯度以显存高效的方式计算精确的全批次梯度，避免存储全批次的所有激活。
• 通过缓存距离梯度和表示梯度，将该方法扩展至支持深度相似度函数Φ(·,·)，实现相似度和表示参数的联合优化。

实验结果

研究问题

• RQ1 我们能否在单张GPU上训练大批次的对比模型，且不超出显存限制？
• RQ2 将编码器反向传播与对比损失解耦，是否能实现与全批次训练等效的精确梯度更新？
• RQ3 所提出的方法能否仅使用单张消费级GPU，在密集检索基准上复现最先进性能？
• RQ4 与梯度累积和顺序训练相比，该方法在训练速度和显存使用方面的扩展性如何？
• RQ5 该技术能否推广到超越简单点积相似度的深度距离函数？

主要发现

• Gradient Cache方法使在单张RTX 2080ti上训练批次大小为512成为可能，其性能与使用8张V100 GPU训练的模型相当。
• 在增加20%运行时间的前提下，该方法在MS MARCO段落排序基准上达到最先进结果，优于标准训练和梯度累积基线。
• 该方法的峰值显存使用量与批次大小无关，保持恒定，使批次大小可突破硬件限制。
• 梯度累积因每次更新中批次内负样本较少而性能不足；顺序训练则因高方差和较低准确率而表现较差。
• 通过双缓存机制，该方法成功推广至深度距离函数，同时保持相似度和表示参数的精确梯度计算。