MoCo V3：视觉自监督迎来Transformer
何凯明从 CVPR 2020 上发表的 MoCo V1（Momentum Contrast for Unsupervised Visual Representation Learning），到前几天挂在arxiv上面的 MoCo V3（An Empirical Study of Training Self-Supervised Visual Transformers），MoCo一共走过了三个版本。
今天介绍 MoCo 系列第三版，MoCo v1 和 v2 是针对 CNN 设计的，而 MoCo v3 是针对 Transformer 结构设计的，反映了 MoCo 系列对视觉模型的普适性。
[TOC]
MoCo V3 原理分析
MoCo v3 的算法原理不应该是这篇论文的重点，这篇论文的重点应该是将目前无监督学习最常用的对比学习应用在 ViT 上。MoCo v3 相比 v2 去掉了 memory queue，转而像SimCLR 那样采用large batch来取得稍好一点的结果，从结构上 encoder f_q 借鉴 BYOL（Bootstrap Your Own Latent: A New Approach to Self-Supervised Learning） 那样增加了一个 prediction head（两层FC），在ResNet上效果有稍许提升。最后通过大量的实验，去证明如何去克服自监督中引入 ViT 训练不稳定的问题。
MoCo V3 算法原理
作者给出的结论是：影响自监督 ViT 模型训练的关键是：instability，即训练的不稳定性。这种训练的不稳定性所造成的结果并不是训练过程无法收敛 (convergence)，而是性能的轻微下降 (下降1%-3%的精度)。
MoCo v3 损失函数和 v1 和 v2 相同都是使用 InfoNCE，表达式如下：
不一样的是 MoCo V3 在网络结构的组成 Framework 有所差异，具体如图所示。因为引入了 ViT 视觉Transformer结构，所以对数据的输入不再是一张完成的图片，而是 image Patch。另外 Transformer对于长序列具有 Attention 的 Q、K、V 结构能够存储和记忆大量的信息，因此取消了 Memory Queue，直接利用大 Batch Size 来进行学习训练。
MoCo v3 的训练方法和 MoCo v1/2 的训练方法的差异是：
- 取消 Memory Queue，用大 Batch Size：MoCo V3 的 Framework 里面没有 Memory Queue，这就意味着 MoCo v3 所观察的负样本都来自一个 Batch 的图片。换句话讲，只有当 Batch size 足够大时，模型才能看到足够的负样本，所以 MoCo v3 中 Batch size = 4096 这样一个巨大的 Batch size。
- 学习 BYOL 添加 Prediction head：Encoder f_q 除了 Backbone 和预测头 Projection head 以外，还添加了个 Prediction head，是遵循了 BYOL 这篇论文的方法。即 Encoder f_q（ViT(BP) + Projection head + Prediction head），Encoder f_k（ViT(Momentum) + Projection head）。
- Contrastive loss的更新方式修改：对于同一张图片的2个增强版本 x1,x2 ，分别通过 Encoder f_q 和 Momentum Encoder f_k 得到 q1,q2 和 k1,k2。让 q1,k2 通过 Contrastive loss (式 1) 进行优化 Encoder f_q的参数，让 q2,k1 通过 Contrastive loss (式 1) 进行优化 Encoder f_q的参数。Momentum Encoder f_k 则跟 MoCo V1 版本相同通过动量更新。即：
MoCo V3 算法分析
- 数据增强：
在ImageNet中 有一堆无标签的数据，拿出一个 MiniBatch，代码表示为 x，也就是 N 张图片，分别进行两种不同的数据增强，得到 x_1 和 x_2，此时 x_1 是 N 张图片，x_2 也是 N 张图片。
for x in loader: # load a minibatch x with N samples
x1, x2 = aug(x), aug(x) # augmentation- 分别通过 Encoder 和 Momentum Encoder：
x_1 分别通过 Encoder 和 Momentum Encoder 得到特征 q_1 和 k_1，维度是 [N,C]，这里特征空间由一个长度为 C=128 的向量表示。
x_2 分别通过 Encoder 和 Momentum Encoder 得到特征 q_2 和 k_2，维度是 [N,C] ，这里特征空间由一个长度为 C=128 的向量表示。
q1, q2 = f_q(x1), f_q(x2) # queries: [N, C] each
k1, k2 = f_k(x1), f_k(x2) # keys: [N, C] each- Contrastive loss 的定义：
对两个维度是 [N,C] 的矩阵（比如是q_1和k_2）做矩阵相乘，得到维度是 [N,N] 的矩阵，其对角线元素代表的就是 positive sample 的相似度，就是让对角线元素越大越好，所以目标是整个 [N,N] 的矩阵越接近单位阵越好。
def ctr(q, k):
logits = mm(q, k.t()) # [N, N] pairs
labels = range(N) # positives are in diagonal
loss = CrossEntropyLoss(logits/tau, labels)
return 2 * tau * loss- Contrastive loss 优化：
通过一个 Contrastive loss 优化 q_1 和 k_2，通过另一个 Contrastive loss 优化 q_2 和 k_1，并反向传播更新 f_q 的参数：
loss = ctr(q1, k2) + ctr(q2, k1) # symmetrized
loss.backward()
update(f_q) # optimizer update: f_q- Momentum Encoder的参数使用动量更新：
f_k = m*f_k + (1-m)*f_q # momentum update: f_kMoCo V3 提升ViT训练稳定性
重头戏主要在MoCo v3在ViT上的实验，下面主要是对实验部分进行介绍。由于作者给出的结论是：影响自监督ViT模型训练的关键是：instability，即训练的不稳定性。 而这种训练的不稳定性所造成的结果并不是训练过程无法收敛 (convergence)，而是性能的轻微下降 (下降1%-3%的精度)。下面具体来看看每个实验部分的内容。
不稳定性测试
- Batch size 过大使得训练不稳定
如下图，Encoder 架构换成 ViT-B/16 ，Learning rate=1e-4，在 ImageNet 数据集上训练 100 epochs 的结果。作者使用了4种不同的 Batch size：1024, 2048, 4096, 6144 的结果。可以看到当 bs=4096 时，曲线出现了 dip 现象 (稍稍落下又急速升起)。这种不稳定现象导致了精度出现下降。当 bs=6144 时，曲线的 dip 现象更大了，可能是因为跳出了当前的 local minimum。这种不稳定现象导致了精度出现了更多的下降。
正常在CNN架构上，Batch size越大，那么自监督学习中负样本数量越大，能够学习更多的负样本特征，但是使用 ViT结构取消 Memory Queue，用大 Batch Size后却出现了 dip 现象，继续实验。
- Learning rate 过大使得训练不稳定
如下图，Encoder 架构换成 ViT-B/16 ，Batch size=4096，在 ImageNet 数据集上训练 100 epochs 的结果。作者使用了4种不同的 Learning rate：0.5e-4, 1.0e-4, 1.5e-4 的结果。可以看到当Learning rate 较大时，曲线出现了 dip 现象 (稍稍落下又急速升起)。这种不稳定现象导致了精度出现下降。
- LARS optimizer 的不稳定性
如下图，使用了 LARS 优化器，分别使用了4种不同的 Learning rate：3e-4, 5e-4, 6e-4, 8e-4 的结果。结果发现当给定合适的学习率时，LARS的性能可以超过AdamW，但是当学习率稍微变大时，性能就会显著下降。而且曲线自始至终都是平滑的，没有 dip 现象。所以最终为了使得训练对学习率更鲁棒，作者还是采用 AdamW 作为优化器。因为若采用 LARS，则每换一个网络架构就要重新搜索最合适的 Learning rate。
提升稳定性的方法
既然上面重点分析了 Batch Size、Learning rate 和 LARS optimizer对不稳定性的分析，论文实验中就给出了提升稳定性的方法：random patch projection。
导致训练出现不稳定的这些 dip 跟梯度暴涨 (spike) 有关，第1层会先出现梯度暴涨的现象，结果几十次迭代后，会传到到最后1层。作者解决的办法是冻结第1层的参数，也就是patch embedding那层，随机初始化后，不再更新这一层的参数。
实验结果通过下面的图可以看到，使用 MoCo v3 or SimCLR, BYOL 方法，Encoder 架构换成 ViT-B/16 ，Batch size=4096，在 ImageNet 数据集上训练 100 epochs 的结果，不同的是冻结了patch embedding那层的参数，使用了随机参数初始化。
不论是 MoCo v3 还是 SimCLR, BYOL 方法，冻结 patch embedding 那层的参数都能够提升自监督 ViT 的训练稳定性。除此之外， gradient-clip 也能够帮助提升训练稳定性，其极限情况就是冻结参数，真的可以避免出现 dip 现象。
对比实验
有意思的是这篇论文的写作方式，首先通过实验发现问题，然后通过实验去解决问题，然后再做一些对比实验，深度学习发文章真的是实验对比非常重要。
作者对 1) position embedding 的具体形式，2) class token 的必要性，3) Prediction head 的必要性和 4) momentum 超参数的影响分别做了不同的对照实验。其中 position embedding 和 class token 的影响分析比较重要，另外两个对精度影响在千分之一之间。
- 位置编码的具体形式
在 Transformer 结构里面position embedding很重要，在无监督训练过程去除位置编码，效果下降了1个多点，说明 ViT 的学习能力很强，在没有位置信息的情况下就可以学习的很好；从另外一个角度来看，也说明 ViT 并没有充分利用好位置信息。
- class token 的必要性
使用 class token 的性能是76.5，而简单地取消 class token，并换成 Global Average Pooling 会下降到69.7，这时候最后一层后面有个LN层。如果把它也去掉，性能会提升到76.3。说明 class token 并不是必要的，LN的选择也很重要。个人猜测 负样本中的 Layer Norm 数据归一化很重要，否则会引起负样本数据不均衡。
最后作者跟采用了 Big ResNet 的方法进行对比，以 VIT-L 为backbone的 MoCo v3 完胜。
小结
MoCo v3 的改进主要1）取消了 Memory Queue 的机制，2）添加了 Prediction head，3）更新Contrastive loss 优化 Encoder 参数方式。在 自监督学习引入 ViT 的过程中发现了训练不稳定性的问题，通过 random patch embedding 的方式暂时解决了目前问题，但是更大的 Batch Size仍然会引起 dip 问题还需要进一步研究啦。
引用
[1] Hadsell, Raia, Sumit Chopra, and Yann LeCun. "Dimensionality reduction by learning an invariant mapping." 2006 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR'06). Vol. 2. IEEE, 2006.
[2] Chen, Ting, et al. "A simple framework for contrastive learning of visual representations." International conference on machine learning. PMLR, 2020.
[3] He, Kaiming, et al. "Momentum contrast for unsupervised visual representation learning." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2020.
[4] Chen, Xinlei, et al. "Improved baselines with momentum contrastive learning." arXiv preprint arXiv:2003.04297 (2020).
[5] Grill, Jean-Bastien, et al. "Bootstrap your own latent-a new approach to self-supervised learning." Advances in neural information processing systems 33 (2020): 21271-21284.
[6] 何凯明新作MoCo V3！！！探讨一下它的前世与今生 - 知乎
[7] 如何评价Kaiming He团队的MoCo v3？ - 知乎
[8] https://zhuanlan.zhihu.com/p/38