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，表达式如下：

$\mathcal{L}_{q}=-\log \frac{\exp \left(q \cdot k_{+} / \tau\right)}{\sum_{i=0}^{K} \exp \left(q \cdot k_{i} / \tau\right)} \tag{1}$

不一样的是 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 版本相同通过动量更新。即：

$loss = Contrastive(q_1, k_2) + Contrastive(q_2, k_1) \tag{2}$

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_k

MoCo 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] https://zhuanlan.zhihu.com/p/364446773

[7] https://www.zhihu.com/question/453203448/answer/1843542119

[8] https://zhuanlan.zhihu.com/p/38

posted @ 2022-07-18 22:17 ZOMI酱酱阅读(368) 评论(0) 编辑收藏举报

会员力量，点亮园子希望

刷新页面返回顶部