MoCo系列

Dec 29, 2024
2 views
Self-Supervised

总结下 Self-Supervised Learning 的方法,用 4 个英文单词概括一下就是:

Unsupervised Pre-train, Supervised Fine-tune.

在预训练阶段我们使用无标签的数据集 (unlabeled data),因为有标签的数据集很贵,打标签得要多少人工劳力去标注,那成本是相当高的,所以这玩意太贵。相反,无标签的数据集网上随便到处爬,它便宜。在训练模型参数的时候,我们不追求把这个参数用带标签数据从初始化的一张白纸给一步训练到位,原因就是数据集太贵。于是 Self-Supervised Learning 就想先把参数从 一张白纸 训练到 初步成型,再从 初步成型 训练到 完全成型。注意这是2个阶段。这个训练到初步成型的东西,我们把它叫做 Visual Representation。预训练模型的时候,就是模型参数从 一张白纸 到 初步成型 的这个过程,还是用无标签数据集。等我把模型参数训练个八九不离十,这时候再根据你 下游任务 (Downstream Tasks) 的不同去用带标签的数据集把参数训练到 完全成型,那这时用的数据集量就不用太多了,因为参数经过了第1阶段就已经训练得差不多了。

第1个阶段不涉及任何下游任务,就是拿着一堆无标签的数据去预训练,没有特定的任务,这个话用官方语言表达叫做:in a task-agnostic way。第2个阶段涉及下游任务,就是拿着一堆带标签的数据去在下游任务上 Fine-tune,这个话用官方语言表达叫做:in a task-specific way

以上这些话就是 Self-Supervised Learning 的核心思想,如下图1所示。

image

1. MoCo v1

论文名称:Momentum Contrast for Unsupervised Visual Representation Learning

论文地址:

https://arxiv.org/pdf/1911.05722.pdfarxiv.org/pdf/1911.05722.pdf

开源地址:

facebookresearch/mocogithub.com/facebookresearch/moco

image

MoCo 系列也遵循这个思想,预训练的 MoCo 模型也会得到 Visual Representation,它们可以通过 Fine-tune 以适应各种各样的下游任务,比如检测和分割等等。MoCo在 7 个检测/语义分割任务(PASCAL VOC, COCO, 其他的数据集)上可以超过他的有监督训练版本。有时会超出很多。这表明在有监督与无监督表示学习上的差距在许多视觉任务中已经变得非常近了。

自监督学习的关键可以概括为两点:Pretext Task,Loss Function,在下面分别介绍。

1.1 自监督学习的 Pretext Task

Pretext Task 是无监督学习领域的一个常见的术语,其中 "Pretext" 翻译过来是"幌子,托词,借口"的意思。所以 Pretext Task 专指这样一种任务:这种任务并非我们所真正关心的,但是通过完成它们,我们能够学习到一种很好的表示,这种表示对下游任务很重要。

The term "pretext" implies that the task being solved is not of genuine interest, but is solved only for the true purpose of learning a good data representation.

我这里举几个例子:

(1) BERT 的 Pretext Task:在训练 BERT 的时候,我们曾经在预训练时让它作填空的任务,

如下图2所示,把这段输入文字里面的一部分随机盖住。就是直接用一个Mask把要盖住的token (对中文来说就是一个字)给Mask掉,具体是换成一个特殊的字符。接下来把这个盖住的token对应位置输出的向量做一个Linear Transformation,再做softmax输出一个分布,这个分布是每一个字的概率。因为这时候BERT并不知道被 Mask 住的字是 "湾" ,但是我们知道啊,所以损失就是让这个输出和被盖住的 "湾" 越接近越好。

image

通过图2这种方式训练 BERT,得到的预训练模型在下游任务只要稍微做一点 Fine-tune,效果就会比以往有很大的提升。

所以这里的 Pretext Task 就是填空的任务,这个任务和下游任务毫不相干,甚至看上去很笨,但是 BERT 就是通过这样的 Pretext Task学习到了很好的 Language Representation,很好地适应了下游任务。

(2) SimCLR 的 Pretext Task:在训练 SimCLR 的时候,我们曾经在预训练时试图教模型区分相似和不相似的事物,如下图3所示,假设现在有1张任意的图片 \(x\) ,叫做Original Image,先对它做数据增强,得到2张增强以后的图片 \(x_i, x_j\) 。接下来把增强后的图片\(x_i, x_j\)输入到Encoder里面,注意这2个Encoder是共享参数的,得到representation \(h_i, h_j\) ,再把 \(h_i, h_j\) 继续通过 Projection head 得到 representation \(z_i, z_j\) ,这里的2个 Projection head 依旧是共享参数的,且其具体的结构表达式是:

\[ z_i=g(h_i)=W^{(2)}\sigma (W^{W(1)}h_i) \]

接下来的目标就是最大化同一张图片得到的 \(z_i^1,z_j^1\),最小化不同张图片得到的 \(z_i^1,z_j^1,z_i^2,z_j^2,z_i^3,z_j^3,z_i^4,z_j^4,z_i^5,z_j^5,...\) 。

image

通过图3这种方式训练 SimCLR,得到的预训练模型在下游任务只要稍微做一点 Fine-tune,效果就会比以往有很大的提升。

所以这里的 Pretext Task 就是试图教模型区分相似和不相似的事物,这个任务和下游任务毫不相干,甚至看上去很笨,但是 SimCLR 就是通过这样的 Pretext Task学习到了很好的 Image Representation,很好地适应了下游任务。

还有一些常见的 Pretext Task 诸如denoising auto-encoders,context autoencoders,cross-channel auto-encoders等等,这里就不一一介绍了。

1.2 自监督学习的 Contrastive loss

Contrastive loss 来自于下面这篇 Yann LeCun 组的工作,如何理解这个对比损失呢?

http://yann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdfyann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdf

常见的损失函数是 cross entropy loss,它适合于数据的 label 是 one-hot 向量的形式。此时网络结构的最后一层是softmax,输出得到各个类的预测值。比如现在有3个类:dog, cat, horse,它们的 label 分别对应着 (1,0,0), (0,1,0), (0,0,1),cross entropy loss会让 dog 图片的输出尽量接近 (1,0,0),让 cat 图片的输出尽量接近 (0,1,0),让 horse 图片的输出尽量接近 (0,0,1)。但是这也存在一个问题,就是假设再来3个类,分别是sky,car和bus。那么按道理dog与horse的距离 应该比dog与sky的距离近,因为dog与horse都属于动物;car与bus的距离 应该也比car与cat的距离近,因为car与bus都属于车类。但是 cross entropy loss确是一视同仁地把 dog与horse的距离 和 dog与sky的距离 看作是一样的。

Contrastive loss 的初衷是想让:1 相近的样本之间的距离越小越好。2 相远的样本之间的距离越大越好。

如果神经网络的损失函数只满足条件1,那么网络会让任何的输入都输出相同的值,不论输入是 dog, cat, horse 还是 sky,car, bus,输出都是一样的,这确实满足了相近的样本之间的距离越小越好,但是却使得网络丧失了分类能力。

如果神经网络的损失函数只满足条件 1 和 2,是不是就完善了呢?

实际上如果想让相远的样本之间的距离越大越好,就需要一个边界,否则如果 dog 是 (1,0,0),那么 假设第1轮训练网络输出 cat 是 (0,1,0),第2轮训练网络输出 cat 是 (0,5,0)...这样下去dog 与 cat 之间的距离越来越大,网络却没法收敛。

Contrastive loss 改进的思路就是让相远的样本之间的距离越大越好,但是这个距离要有边界,即要求:1 相近的样本之间的距离越小越好。2 相远的样本之间的距离越大越好,这个距离最大是 � 如下图4弹簧图所示:黑色实心球代表与蓝色球相近的样本,白色空心球代表与蓝色球相远的样本,蓝色箭头的长度代表力的大小,方向代表力的方向。

(a):Contrastive loss 使得相近的样本接近。

(b):横轴代表样本之间的距离,纵轴代表loss值。Contrastive loss 使得相近的样本距离越小,loss值越低。

(c):Contrastive loss 使得相远的样本疏远。

(d):横轴代表样本之间的距离,纵轴代表loss值。Contrastive loss 使得相远的样本距离越大,loss值越低,但距离存在上界,就是红线与x轴的交点 � ,代表距离的最大值。

(e):一个样本受到其他各个独立样本的作用,各个方向的力处于平衡状态,代表模型参数在 Contrastive loss 的作用下训练到收敛。

image

Contrastive loss 用公式表示为:

\[ L(W, Y, \vec{X_1},\vec{X_2})=(1-Y)\frac{1}{2}(D_W)^2+(Y)\frac{1}{2}\{max(0, m-D_W)\}^2\\ D_W(\vec{X_1},\vec{X_2})=||G_W(\vec{X_1}-\vec{X_2})||_2 \]

式中, \(\vec{X_1}, \vec{X_2}\) 是2个样本,\(D_W\) 是 \(\vec{X_1}\) 与 \(\vec{X_2}\) 在潜变量空间的欧几里德距离, \(Y=0\) 代表是相近的样本,此时要求 \(D_W\) 尽量接近0。 \(Y=1\) 代表是相远的样本,此时要求 \(D_W<m\) 时与 \(m\) 越接近越好,但是距离一旦超过 \(m\) 即失效,不会再接着更新参数使得距离越来越大。

有意思的是那些属于不同类,但两两距离天生大于m的样本对。 LeCun 的对比损失完全忽视这些样本对,大大减少了计算量。此外,Contrastive loss 提供了一个不同类别样本点之间距离的下界 \(m\) 。

1.3 MoCo v1之前的做法

了解了 Pretext Task 和 Contrastive loss,接下来就要正式介绍MoCo v1的原理了。

image

如上图5所示,输入 \(x^q\) 通过编码器 Encoder 得到 query \(q\) ,还有几个输入 \(x^{k_0}, x^{k_1}, x^{k_2}\) 通过编码器 Encoder 得到 key \(k_0, k_1, k_2\) 。假设只有一个 key \(k_+\) 和 \(q\)  是匹配的。根据上面的 Contrastive loss 的性质,只有当 \(q\) 和相匹配的 \(k_+\) 相近,且与其他不匹配的 \(k_0, k_1,k_2\) 相远时, Contrastive loss 的值才会最小。这个工作使用点积作为相似度的度量,且使用 InfoNCE 作为Contrastive loss,则有:

\[ \mathcal{L}_q=-log\frac{exp(q\cdot k_+/\tau)}{\sum_{i=0}^Kexp(q\cdot k_i/\tau)} \]

式中, \(\tau\)是超参数,这个式子其实非常像把 q 分类成第 \(k_+\) 类的cross entropy loss。

这里的 \(x^q, x^k\) 可以是image,可以是image patches等, \(q=f_q(x^q), k=f_k(x^k)\), \(f_q, f_k\) 是Encoder,也可以是很多种架构。

[1 原始的端到端自监督学习方法]:对于给定的一个样本 \(x^q\) , 选择一个正样本 \(x^{k_+}\) (这里正样本的对于图像上的理解就是 \(x^q\) 的 data augmentation版本)。然后选择一批负样本 (对于图像来说,就是除了 \(x^q\) 之外的图像),然后使用 loss function \(\mathcal{L}\) 来将 \(x^q\) 与正样本之间的距离拉近,负样本之间的距离推开。样本 \(x^q\) 输入给 Encoder \(f_q\) ,正样本和负样本都输入给 Encoder \(f_k\) 。

这样其实就可以做自监督了,就可以进行端到端的训练了。实际上这就是最原始的图像领域的自监督学习方法,如下图6所示,方法如上面那一段所描述的那样,通过loss function \(\mathcal{L}_q\) 来更新2个Encoder \(f_q, f_k\) 的参数。

原始的自监督学习方法里面的这一批负样本就相当于是有个字典 (Dictionary),字典的key就是负样本,字典的value就是负样本通过 Encoder \(f_k\) 之后的东西。那么现在问题来了:

问:这一批负样本,即字典的大小是多大呢?

答:负样本的规模就是 batch size,即字典的大小就是 batch size。

举个例子,假设 batch size = 256,那么对于给定的一个样本 \(x^q\) , 选择一个正样本 \(x^{k_+}\) (这里正样本的对于图像上的理解就是 \(x^q\) 的 data augmentation版本)。然后选择256个负样本 (对于图像来说,就是除了 \(x^q\) 之外的图像),然后使用 loss function \(\mathcal{L}_q\) 来将 \(x^q\) 与正样本之间的距离拉近,负样本之间的距离推开。

毫无疑问是 batch size 越大效果越好的,这一点在 SimCLR 中也得到了证明。但是,由于算力的影响 batch size 不能设置过大,因此很难应用大量的负样本。因此效率较低。

image

图6:原始的端到端自监督学习方法

image

针对很难应用大量的负样本的问题,有没有其他的解决方案呢?下面给出了一种方案,如下图7所示。

[2 采用一个较大的memory bank存储较大的字典]:对于给定的一个样本 \(x^q\) , 选择一个正样本 \(x^{k_+}\) (这里正样本的对于图像上的理解就是 \(x^q\) 的 data augmentation版本)。采用一个较大的 memory bank 存储较大的字典,这个 memory bank 具体存储的是所有样本的representation。(涵盖所有的样本,比如样本一共有60000个,那么memory bank大小就是60000,字典大小也是60000),采样其中的一部分负样本 \(k^{sample}\) ,然后使用 loss function \(\mathcal{L}_q\) 来将 \(x^q\) 与正样本之间的距离拉近,负样本之间的距离推开。这次只更新 Encoder \(f_q\) 的参数,和几个采样的key值 \(k^{sample}\) 。因为这时候没有了 Encoder \(f_k\) 的反向传播,所以支持memory bank容量很大。

但是,你这一个step更新的是 Encoder \(f_q\) 的参数,和几个采样的key值 \({\color{purple}k^{sample_1}}\),下个step更新的是 Encoder \(f_q\) 的参数,和几个采样的key值 \({\color{red}k^{sample_2}}\),问题是 \({\color{purple}k^{sample_1}}\neq {\color{red}k^{sample_2}}\) ,也就是:Encoder \(f_q\) 的参数每个step都更新,但是某一个 \(k_i\) 可能很多个step才被采样到更新一次,而且一个epoch只会更新一次。这就出现了一个问题,即:每个step编码器都会进行更新,这样最新的 query 采样得到的 key 可能是好多个step之前的编码器编码得到的 key,因此丧失了一致性。

image

从这一点来看,[1 原始的端到端自监督学习方法] 的一致性最好,但是受限于batchsize的影响。而 [2 采用一个较大的memory bank存储较大的字典] 的字典可以设置很大,但是一致性却较差,这看起来似乎是一个不可调和的矛盾。

1.4 MoCo v1 的做法

[3 MoCo方法]:

kaiming大神利用 momentum (移动平均更新模型权重) 与queue (字典) 轻松的解决这个问题。为了便于读者理解,这里结合kaiming大神提供的伪代码一起讲解 (下面加粗的黑体字母是代码中的变量)。

image

image

首先我们假设 Batch size 的大小是 \(N\) ,然后现在有个队列 Queue,这个队列的大小是 \(K(K>N)\) ,注意这里 \(K\) 一般是 \(N\) 的数倍,比如 \(K=3N,K=5N,... ,\)但是 \(K\) 总是比 \(N\) 要大的 (代码里面 \(N\)=65536 ,即队列大小实际是65536)。

下面如上图8所示,有俩网络,一个是 Encoder \(f_q\) ,另一个是Momentum Encoder \(f_{MK}\) 。这两个模型的网络结构是一样的,初始参数也是一样的 (但是训练开始后两者参数将不再一样了)。 \(f_q\) 与 \(f_{MK}\) 是将输入信息映射到特征空间的网络,特征空间由一个长度为 \(C\) 的向量表示,它们在代码里面分别表示为:\(f_q , f_k\)  C

代码里的 k 可以看作模板,q 看作查询元素,每一个输入未知图像的特征由 f_q 提取, 现在给一系列由 \(f_k\) 提取的模板特征 (比如狗的特征、猫的特征) ,就能使用$ f_q$ 与 \(f_k\) 的度量值来确定 \(f_q\) 是属于什么。

1) 数据增强:

现在我们有一堆无标签的数据,拿出一个 Batch,代码表示为 x,也就是 \(N\) 张图片,分别进行两种不同的数据增强,得到 x_q 和 x_k,则 x_q 是 \(N\) 张图片,x_k 也是 \(N\) 张图片。

**for** x **in** loader: *# 输入一个图像序列x,包含N张图,没有标签
        *x_q **=** aug(x) *# 用于查询的图 (数据增强得到)
        *x_k **=** aug(x) *# 模板图 (数据增强得到),自监督就体现在这里,只有图x和x的数据增强才被归为一类*

2) 分别通过 Encoder 和 Momentum Encoder:

x_q 通过 Encoder 得到特征 q,维度是 \(N,C\) ,这里特征空间由一个长度为 \(C=128\) 的向量表示。

x_q 通过 Momentum Encoder 得到特征 k,维度是 \(N,C\) 。

    q **=** f_q**.**forward(x_q) *# 提取查询特征,输出NxC
        *k **=** f_k**.**forward(x_k) *# 提取模板特征,输出NxC*

3) Momentum Encoder的参数不更新:

    *# 不使用梯度更新f_k的参数,这是因为文章假设用于提取模板的表示应该是稳定的,不应立即更新
        *k **=** k**.**detach()

4) 计算 \(N\) 张图片的自己与自己的增强图的特征的匹配度:

    *# 这里bmm是分批矩阵乘法
        *l_pos **=** bmm(q**.**view(N,1,C), k**.**view(N,C,1)) *# 输出Nx1,也就是自己与自己的增强图的特征的匹配度*

这里得到的 l_pos 的维度是 (N, 1, 1)N 个数代表\(N\) 张图片的自己与自己的增强图的特征的匹配度。

5) 计算 \(N\) 张图片与队列中的 \(K\) 张图的特征的匹配度:

    l_neg **=** mm(q**.**view(N,C), queue**.**view(C,K)) *# 输出Nxk,自己与上一批次所有图的匹配度(全不匹配)*

这里得到的 l_neg 的维度是 (N, K),代表\(N\) 张图片与队列 Queue 中的\(K\)张图的特征的匹配度。

6) 把 4, 5 两步得到的结果concat起来:

    logits **=** cat([l_pos, l_neg], dim**=**1) *# 输出Nx(1+k)*

这里得到的 logits 的维度是 (N, K+1),把它看成是一个矩阵的话呢,有 N 行,代表一个 Batch里面的 N 张图片。每一行的第1个元素是某张图片自己与自己的匹配度,每一行的后面K 个元素是某张图片与其他 K 个图片的匹配度,如下图9所示,图9展示的是某一行的信息,这里的 K=2

image