moco系列学习总结

MoCo系列学习总结

刚开始博客，写的不好请多多包含。如有错误，请指出，感谢。

• 原文链接
  • MoCo-v1:https://arxiv.org/abs/1911.05722
  • MoCo-v2:https://arxiv.org/abs/2003.04297
  • MoCo-v3:https://arxiv.org/abs/2104.02057
• 源码链接
  • MoCo-v1:https://github.com/facebookresearch/moco
  • MoCo-v2:https://github.com/facebookresearch/moco
  • MoCo-v3:https://github.com/facebookresearch/moco-v3

一、研究背景

• 当前存在两种自监督的训练策略，两者都存在较为明显的缺陷。
  • 端到端训练策略
    • 方法：每次对一个batch中的图像进行对比学习。
    • 优点：端到端的训练速度快，正负样本的编码都随着模型的更新而更新，特征一致性好。
    • 缺陷：每个step训练样本只能与同一个batch中的样本进行对比，这就非常依赖大的batch_size，更大的batch_size在计算loss时，能够使其与整个数据集上的loss更加接近，从而带来更好的训练效果。
  • memory bank训练策略
    • 方法：将数据集的特征提前使用预训练权重的模型提取好存下，作为memory bank。在训练时随机在memory bank中采样，该step结束后更新memory bank中对应的特征。
    • 优点：memory bank可以直接存下整个imageNet的特征向量用于对比，能够在训练时拥有较好的样本多样性，解决了对大batch_size的依赖。
    • 缺陷：每个step从memory bank中随机采样进行模型更新。但随着模型不断更新，memory bank中的特征更新速度比较缓慢，导致有一些特征要过很久才能被第一次采样到，完成更新。memory bank中许多特征都是被不同时期的模型更新的，其特征一致性相对较差。
      
  • MoCo-v1，MoCo采用动量更新的模型和一个队列作为特征的字典解决该问题。
    • 动量更新的模型，尽可能使图像特征保持一致性。
    • 队列，先进先出的特性，总是优先提出模型早期提取的特征，同样是保持特征一致性。
    • 每个step从队列中随机采样，类似于memory bank，能够解决对大batch_size的依赖。

二、MoCo-v1

• 模型原理
  

  • encoder：backbone,表示常规更新的backbone。
  • momentum encoder:其结构与encoder相同，其参数通过动量加权更新，相较于encoder更新速度更慢。
  • $x^{query}$：$arg1(x)$，数据增广1。
  • $x_{0,1,...,n}^{key}$:$arg2(x_{0,1,...,n})$，数据增广2。
  • queue:使用momentum encoder对$x_{0,1,...,n}^{key}$编码后的特征表征$k_0,k_1,...,k_n$。
    
  • 设batch_size = $N$
  • l_pos：将一个batch中的q和对应的k相乘，得到一个$N\times 1$的特征相似度矩阵，由于q与k都是由同一个x生成，因此视为一组正样本。
  • l_neg：将一个batch中的q和整个queue相乘，对于q而言，整个queue都视为负样本，结果为$N\times C$的相似度矩阵，K表示queue的长度。
  • labels：是一个$N\times 1$的全0向量，表示所有的分类结果都在第一位，即指向l_pos那一列
  • loss：求一个交叉熵，做分类，分类至第1列，就是使logits的第一个数值最大，即正例的相似度最高。
  • 动量更新：${\theta }_k=m\times {\theta }_k+(1-m)\times {\theta }_q$

• Loss函数（InfoNCEloss）

  • $l(i,j)=-log\frac{exp(s_{i,j}/\tau )}{{\sum }_{k=1}^{2N}{1}_{k\ne i}exp(s_{i,j}/\tau )}$
  • 实际上就是一个在局部样本中进行多分类的交叉熵loss。与传统分类网络不同之处在于
    • 表示概率的向量元素，并非回归，而是通过余弦相似度计算所得。
    • softmax的分类数量并非是整个数据集的种类量。与当前采样数量有关(局部)。

• 总结

  • 特点：动量更新的编码器能够较为平滑的生成样本负例，产生一种近似同一个模型的特征编码，将其保存在一个内存队列queue中，一定程度上能够解决对大batch_size的需求，在小batch_size上也能有很好的效果，是我等的福音。
  • shuffle BN：为了解决BN层信息泄露问题，作者提出了Shuffling BN的概念，而这个层必须在多GPU卡上才能进行。首先作者将不同GPU的sub mini batch进行汇聚，然后对汇聚后的batch进行打乱，并且记下打乱后的索引（便于恢复shuffle结果），在提取了Key编码特征后，通过之前记下的索引，恢复shuffle结果。注意到这个过程只对Key编码器进行，而不考虑Query编码器。（了解即可，后来用的不多）。

三、MoCo-v2

• 模型改进

  • 在backbone后面增加投影层（project head）：MLP
  • 更强的图像增广操作

• 总结
  特点：将SimClr的优势结合在MoCo-v1上

四、MoCo-v3

• 伪代码
  

  • 自预测，通过predict head结果（query）=>project head的结果(key)

• 模型改进

  • 采用vit作为自监督模型的backbone
  • 结合了BYOL的自预测头的方法
  • 保留了动量模型，但是却不再使用MoCo-v1,v2中的queue，而是直接采用大batch_size进行端到端的对比学习。

• 训练现象及解决技巧
  作者花了大量的篇幅探讨了，在自监督对比学习中，以vit作为backbone，存在训练存在严重波动的问题，尤其是对于大batch_size。通过可视化各层梯度的情况，发现梯度变化的峰值发生在vit模型第一层，即patch project层，采用固定patch project层的参数的方法解决该问题。
  

• 总结

  • 作者在摘要第一句就表示，本文没有提出新的方法。
  • 主要工作：引入vit，并结合BYOL、SimSiam中的模块进行结构改进，并做了大量的实验。
  • 探讨了vit带来的训练波动问题。

五、全文总结

• 通过这个系列的学习，能够对自监督训练的框架有一下几个问题有更清晰的认识。
  • 采用大batch_size训练的深层原因。
  • MoCo是如何在研究背景的两种架构中进行改进的。
  • InfoNCEloss与多分类交叉熵的联系。
  • 为何在交叉熵中的温度系数$\tau $是必须的，读者可以自行计算一下loss的上下限情况。
  • 自监督对比学习框架的发展史：单层分类=>MLP(project head)=>project head+predict head。
• 亚马逊大佬表示，MoCo-v2在对比学习领域依旧是比较实用的存在，结构较新，而且使用queue来降低我们对大batch_size的依赖，比较推荐使用MoCo-v2做自己的项目。

参考视频

• https://www.bilibili.com/video/BV1C3411s7t9?spm_id_from=333.337.search-card.all.click

参考文献

[1] He K , Fan H , Wu Y , et al. Momentum Contrast for Unsupervised Visual Representation Learning[C]// 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). IEEE, 2020.
[2] Chen, X., Fan, H., Girshick, R., and He, K., “Improved Baselines with Momentum Contrastive Learning”, 2020, arXiv:2003.04297.
[3] Chen, X., Xie, S., and He, K., “An Empirical Study of Training Self-Supervised Vision Transformers”, 2021, arXiv:2104.02057.