何凯明从 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_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] 何凯明新作MoCo V3!!!探讨一下它的前世与今生 - 知乎
[7] 如何评价Kaiming He团队的MoCo v3? - 知乎文章来源:https://www.toymoban.com/news/detail-446627.html
[8] https://zhuanlan.zhihu.com/p/38文章来源地址https://www.toymoban.com/news/detail-446627.html
到了这里,关于MoCo V3:视觉自监督迎来Transformer的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!