何恺明-MoCo：资源不够，亦能玩转对比学习

“ 给我一块V100，帮你把BatchSize提到1w+”
对比学习已经成为无监督表示学习的一大范式，不研究表明，模型效果与BatchSize成正相关，大厂在训练模型时也动辄将BatchSize提到万级别（ALIGN的16384、CLIP的32768）。如何在资源有限情况下提高BatchSize，已经成为平民炼丹师的关注重点。

论文链接：
https://arxiv.org/pdf/1911.05722.pdf
开源项目：
https://github.com/facebookresearch/moco

01 背景
如前言，计算资源与BatchSize的矛盾是炼丹师使用对比学习的主要矛盾之一；既然难以提高单个step下的BatchSize，能否将历史steps产出的embeddings作为缓存使用，从而变相提高对比BatchSize呢？当然可以，但是需要解决一个问题：随着模型训练，不同steps的模型参数有一定偏移；如果偏移较大，不同steps产出的embeddings不再处于同一个特征空间，从而使对比失去意义甚至产生反效果。基于此，本文提出了MomentumContrast（MoCo）方法。

02 方法
如下图，MoCo将Q/K encoder剥离为两个结构相同、参数不共享的模型，并在训练时维护一个长为CacheSize的队列，两个模型的输入为单条文本/图片的不同增强结果。在计算对比loss时（作者使用的infoNCE），将最新产出的k_embeddings与队列中的历史k_embeddings’进行拼接，构成维度为(BatchSize + CacheSize * BatchSize, HiddenSize)的K_embeddings，从而将真实的BatchSize扩大到了原先的CacheSize+1倍。在计算完loss后，删除最早的batch、将当前batch加入队列，从而完成一次队列更新。

解决了BatchSize扩充的问题，又该如何解决不同step下可能存在的特征空间偏移问题呢？作者选择将Kencoder的优化方式进行调整，即Qencoder仍然基于BP优化，而Kencoder则是缓慢变化的、Qencoder在时间上的平滑结果，计算方式如下图，具体算法流程见附录。

03 实验
首先，在线性判别测试上（即冻结encoder模型，仅对线性判决层进行训练），MoCo和end-to-end（堆资源实现BatchSize提升）效果基本一致、优于memory bank（预先生成全量k_embeddings用于负样本采用，且若干steps更新一次k_embeddings结果）；并且从下图可以看出，判决效果和预训练时BatchSize呈正相关。
​
其次，在迁移学习上，MoCo在7项下游任务上，甚至优于基于ImageNet有监督数据集进行预训练的模型，且MoCo所使用的无监督数据集size越多，效果越好。
​

04 总结
MoCo为资源有限时进行大规模的对比学习训练提供了一种行之有效的解决方案。除了文章给出的方法以外，是否也可以尝试冻结单侧encoder并轮询训练的方式？是否也可以尝试梯度累积+间隔更新？
此外还需要注意的是，文章中应用场景为无监督场景，如果用在有监督场景或Q/K不同源的场景，由于Qencoder仅基于Q进行BP优化，需要考虑Q/K domain不匹配引起的Kencoder的OOD问题。

附录