自监督学习Visual Transformers(ViT)的训练经验(Moco v3) -- 论文解析

    之前讲过一篇自监督学习：自监督学习（Self Supervised Learning），里面有提到几种把图像转成通用的embedding的方式，有CPC, SimCLR, 还有Moco。今天来详细说一下Moco v3，主要讲下面这篇文章：An Empirical Study of Training Self-Supervised Visual Transformers。

    这篇文章主要讲的是自监督ViT的训练技巧。之后一定会讲这篇文章的源代码哦，心动就关注吧。

    CV和NLP在预训练时有两点不同：

• NLP选择的通常是masked auto-encoder，而CV最近比较倾向孪生网络（Siamese networks）
• NLP用的主干结构是self-attentional Transformers，而CV的主干通常是deep residual networks（ResNets）

    这篇文章，主干结构用的是self-attentional Transformers，网络框架是基于Siamese networks（Moco和一些其他框架）。网络属于Vision Transformers (ViT)这个范畴。ViT属于一种新的网络结构，因此关于这种网络的训练技巧比较少，这篇文章主要就从一些很基础的方面，比如batch size，learning rate，optimizer这些方面来讨论，这种网络要怎么训练。实验反馈的结果是，Transformers搭配Contrastive Learning会比搭配auto-encoding效果好。下图中iGPT是Transformers搭配auto-encoding的模型，可以发现Moco v3优势是很大的。另外也可以发现，ViT越大，准确率越高，这表示ViT还持续有潜力。另外，ViT-Large在某些案例里可以直接打败监督学习。

ViT也可以打败ResNets和Contrastive Learning的搭配，下面这张是MoCo各个版本搭配ResNets训练的结果：

    作者发现，训练过程中的不稳定情况会略微降低这个网络最后的精确度，大概1%-3%。而通过一些鸡贼的操作，可以减轻这个问题带来的影响，从而提高精确度。我们就来看看，究竟是个什么问题，他们又是怎么解决的。

    我们看一下Moco v3架构的伪代码：

# f_q: encoder: backbone + proj mlp + pred mlp
# f_k: momentum encoder: backbone + proj mlp
# m: momentum coefficient
# tau: temperature
for x in loader: # load a minibatch x with N samples
  x1, x2 = aug(x), aug(x) # 将同一个样本随机进行不同的数据增强处理
  q1, q2 = f_q(x1), f_q(x2) # 用encoder分别编码
  k1, k2 = f_k(x1), f_k(x2) # 再 momentum encoder分别编码
  loss = ctr(q1, k2) + ctr(q2, k1) # 交叉计算loss再求
  loss.backward()
  update(f_q) # 更新encoder的参数
  f_k = m*f_k + (1-m)*f_q # 根据更新的encoder和之前的momentum encoder得到新的momentum encoder
# contrastive loss
def ctr(q, k):
  logits = mm(q, k.t()) # [N, N] 对角线上的出自同一幅图
  labels = range(N) # positives are in diagonal
  loss = CrossEntropyLoss(logits/tau, labels)
  return 2 * tau * loss

    损失函数在伪代码里是ctr，也可以看上图的公式，q和k+是同一幅图得到的embedding，q和k-则是不同图得到的embedding，q是encoder得到的，k是由momentum encoder得到的，容易理解的是，这样的损失函数，就是要让qk+的值尽量大，qk-的值尽量小，而momentum的设计目的我们之后再说。关于Cross Entropy Loss可以看Cross Entropy Loss。

    这里说一下encoder（f_q）和momentum encoder（f_k）的网络结构，encoder包含的是一个主干网络结构，即ViT，或者ResNet，另外还包含一个projection head和一个prediction head，而momentum encoder则只有projection head和主干网络，momentum encoder在获取encoder参数来调整自己参数的时候，不会用到prediction head。

接下来我们来看一下各种训练参数怎么配置比较合适：

1. Batch size

    通常来说，ViT是很大的模型，配置比较大的batch size往往会取得更好效果，下面是配置了不同batch size的一些实验结果：

为了展示ViT训练存在的不稳定性问题，他们给了上面的实验结果，这里涉及到一个概念，kNN accuracy，这个具体怎么计算，我们解析代码的时候再说，这里我们就把它当成评量自监督模型的方式就好，越高，表示模型越好。实验者将batch size逐渐放大，从1024到6144，发现当batch size变大的时候，kNN accuracy出现了很大的震荡，这种震荡是局部突然下降（作者推测是训练突然重启了，即跳出了原本的局部最优点，重新梯度下降，训练虽然没有偏离，但准确性决定于新的局部点），batch size越大，这种现象越严重（为了更好的捕捉这种震荡，训练的时候他们每十个batch就测一次kNN accuracy）。

2. Learning rate

    首先设定了一个base learning rate: lr，然后再根据batch size调整，调整方式是：lr*BatchSize/256。

    和之前调整batch size一样，learning rate也出现了局部突然下降又快速回升的现象，learning rate越大，这种现象越明显。另外也可以看出来，不是越稳定就能获得越好的结果。

3. Optimizer

    文章里提到了三种Optimizer，一种AdamW（常用来做ViT训练），一种LARS（在自监督学习经常用来匹配batch size很大的训练），最后一种LAMB（AdamW-counterpart of LARS）。实验在不同lr情况下对LAMB进行了测试，发现当lr=5e-4的时候，取得的效果会好于用AdamW的，但是看图我们会发现LAMB对lr的取值非常敏感。虽然曲线平滑，没有用AdamW时候突然下降的情况，但是训练中途会下降很长一段时间，作者们猜测是用LAMB会有梯度不可靠的情况。因此，这篇文章最终选择的Optimizer还是AdamW。

4. 如何提高训练过程中的稳定性

    接下来，我们看看怎么解决精确度突然下降的问题，首先当然要找到产生这个现象的原因：

他们去观察了网络每一层在训练过程中的梯度，发现是梯度的突增导致了精确度的陡降，还发现，总是第一层最先开始突增，大概几十个batch之后，最后一层的梯度也出现了同样的现象。这样不难推测，不稳定性极有可能是最前面的patch projection layer引发的，然后实验者就尝试在训练过程中冻结patch projection layer的参数，即这一层在最开始参数随机配置后就不再做梯度下降，调皮点说，就是既然你不好好学习，那我们就暂时不让你学习啦。开心不？

    答案是开心的。

    不难发现，冻住projection layer效果是明显的，准确率陡降的现象消失了，而且最终获得的准确率也提高了，甚至当lr增大的时候，这种提高还尤为明显。

    另外，他们发现这种不稳定性不止出现在MoCo里，SimCLR和BYOL里也有类似的问题。

同样的，将patch projection layer冻住，不让它学习，可以大大提高SimCLR和BYOL的稳定度，且可以最终提高模型的准确率。此外，还有SwAV，可以用这个方法使得SwAV可以使用更大的lr来训练模型，进而获得更高的精度。所以这个方法几乎可以用于所有的这一类自监督学习模型的训练。

    另外，他们还尝试用BatchNorm（BN），WeightNorm（WN），甚至对patch projection进行梯度裁剪（gradient clip），即设定一个阈值，当梯度超过这个阈值的时候，就让梯度等于这个阈值。实验发现，BN和WN是无用的，梯度裁剪是有用的，但是要设定好阈值，总的来说，还是干脆不让学习最有效。

    这个方法虽然有效地缓和了训练过程中的不稳定性的问题，但是本质上，冻结patch projection来克制不稳定性是一种回避问题的方法，当lr太大时，依然还是会出现不稳定的情况。第一层应当不是造成不稳定的本质原因，只是它不是Transformer的部分，好控制，并不算找到了解决不稳定性的根本办法。对于这个问题，应该要有更好的解决办法才对（小本本记下来，潜在课题）。

    到这里，文章的前四小节本啤应该是讲清楚了，之后的细节会穿插在代码里面说。

    觉得有用，右下角帮啤酒点个在看吧