深度了解自监督学习,就看这篇解读 !微软首创:运用在 image 领域的BERT

2021-07-05 18:38:53 浏览数 (1)

作者丨科技猛兽 来源丨极市平台 编辑丨极市平台

极市导读

本文介绍的这篇工作是把 BERT 模型成功用在 image 领域的首创,也是一种自监督训练的形式,所以取名为视觉Transformer的BERT预训练模型。这个工作用一种巧妙的办法把 BERT 的训练思想成功用在了 image 任务中。 >>加入极市CV技术交流群,走在计算机视觉的最前沿

深度了解自监督学习,就看这篇解读 !Hinton团队力作:SimCLR系列

本文目录

1 BERT 方法回顾 2 BERT 可以直接用在视觉任务上吗? 3 BEiT 原理分析 3.1 将图片表示为 image patches 3.2 将图片表示为 visual tokens 3.2.1 变分自编码器 VAE 3.2.2 BEIT 里的 VAE:tokenizer 和 decoder 3.2.3 BEIT 的 Backbone:Image Transformer 3.2.4 类似 BERT 的自监督训练方式:Masked Image Modeling 3.2.5 BEIT 的目标函数:VAE 视角 3.2.6 BEIT 的架构细节和训练细节超参数 3.2.7 BEIT 在下游任务 Fine-tuning 3.2.8 实验

科技猛兽:Self-Supervised Learning系列解读 (目录)

https://zhuanlan.zhihu.com/p/381354026

Self-Supervised Learning,又称为自监督学习,我们知道一般机器学习分为有监督学习,无监督学习和强化学习。而 Self-Supervised Learning 是无监督学习里面的一种,主要是希望能够学习到一种通用的特征表达用于下游任务 (Downstream Tasks)。其主要的方式就是通过自己监督自己。作为代表作的 kaiming 的 MoCo 引发一波热议, Yann Lecun也在 AAAI 上讲 Self-Supervised Learning 是未来的大势所趋。所以在这个系列中,我会系统地解读 Self-Supervised Learning 的经典工作。

今天介绍的这篇工作是把 BERT 模型成功用在 image 领域的首创,也是一种自监督训练的形式,所以取名为视觉Transformer的BERT预训练模型。这个工作用一种巧妙的办法把 BERT 的训练思想成功用在了 image 任务中,涉及的知识点包括 BERT (第1节),VAE (第3.2.1节) 等等,为了方便阅读本文也会对它们进行简单讲解。

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

Unsupervised Pre-train, Supervised Fine-tune.

下面首先借助 BERT 模型理解一下这句话的意思。

1 BERT 方法回顾

Self-Supervised Learning 超详细解读 (一):大规模预训练模型BERT (https://zhuanlan.zhihu.com/p/378360224)里面我们介绍了 BERT 的自监督预训练的方法,BERT 可以做的事情也就是Transformer 的 Encoder 可以做的事情,就是输入一排向量,输出另外一排向量,输入和输出的维度是一致的。那么不仅仅是一句话可以看做是一个sequence,一段语音也可以看做是一个sequence,甚至一个image也可以看做是一个sequence。所以BERT其实不仅可以用在NLP上,还可以用在CV里面。所以BERT其实输入的是一段文字,如下图1所示。

图1:BERT的架构就是Transformer 的 Encoder

接下来要做的事情是把这段输入文字里面的一部分随机盖住。随机盖住有 2 种,一种是直接用一个Mask 把要盖住的token (对中文来说就是一个字)给Mask掉,具体是换成一个特殊的字符。另一种做法是把这个token替换成一个随机的token。

图2:把这段输入文字里面的一部分随机盖住

接下来把这个盖住的token对应位置输出的向量做一个Linear Transformation,再做softmax输出一个分布,这个分布是每一个字的概率,如下图3所示。

那接下来要怎么训练BERT呢?因为这时候BERT并不知道被 Mask 住的字是 "湾" ,但是我们知道啊,所以损失就是让这个输出和被盖住的 "湾" 越接近越好,如下图4所示。

图3:把这个盖住的token对应位置输出的向量做一个Linear Transformation

图4:让这个输出和被Mask 住的 token 越接近越好

其实BERT在训练的时候可以不止是选取一个token,我们可以选取一排的token都盖住,这就是 SpanBERT 的做法,至于要盖住多长的token呢?SpanBERT定了一个概率的分布,如图5所示。有0.22的概率只盖住一个token等等。

图5:SpanBERT定了一个概率的分布

除此之外,SpanBERT还提出了一种叫做Span Boundary Objective (SBO) 的训练方法,如下图6所示,意思是说:

图6:Span Boundary Objective (SBO)

盖住一串token以后,用这段被盖住的token的左右2个Embedding去预测被盖住的token是什么。SBO把盖住的部分的左右两边的Embedding吃进来,同时还输入一个数字,比如说3,就代表我们要还原被盖住的这些token里面的第3个token。

就是通过上面的图1-图6的方法,让 BERT 看很多的句子,随机盖住一些 tokens,让模型预测盖住的tokens是什么,不断计算预测的 token 与真实的 token 之间的差异,利用它作为 loss 进行反向传播更新参数,来达到 Self-Supervised Learning 的效果。

Self-Supervised Learning 训练好 BERT 以后,如何在下游任务上使用呢?

我们就以情感分析为例,要求输入一个句子,输出对应的情感类别。

BERT是怎么解Sentiment Analysis的问题呢?给它一个句子,在这个句子前面放上 class token,这步和 ViT 是一模一样的。同样地,我们只取输出的Sequence里面的class token对应的那个vector,并将它做Linear Transformation Softmax,得到类别class,就代表这个句子的预测的情感,如下图7所示。

值得注意的是,对于这种下游任务你需要有labelled data,也就是说 BERT 其实没办法凭空解Sentiment Analysis的问题,也是需要一部分有监督数据的。我们此时的情感分析模型包括:

  1. BERT部分
  2. Linear Transformation部分

只是BERT部分的初始化来自 Self-Supervised Learning,而 Linear Transformation 部分采样的是随机初始化。这两部分的参数都用Gradient Descent来更新。

图7:使用BERT做情感分析

下图8其实是个对比,就是BERT部分不用预训练模型的初始化 (scratch) 和用了预训练模型的初始化 (fine-tune) 的不同结果,不同颜色的线代表GLUE中的不同任务。不用预训练模型的初始化会导致收敛很慢而且loss较高,说明预训练模型的初始化的作用。

图8:预训练模型的初始化结果

2 BERT 可以直接用在视觉任务上吗?

上面的 BERT 都是在 NLP 任务上使用,因为 NLP 任务可以把每个词汇通过 Word2Vec 自动转化成一个固定大小的 token,我们随机盖住一些 token,让模型根据这个不完整的句子来预测被盖住的 token 是什么。那么一个自然而然的问题是:对于图片来讲,能否使用类似的操作呢?

第1个困难的地方是:视觉任务没有一个大的词汇表。 在 NLP 任务中,比如图3所示,假设我们盖住词汇 "湾",那么就想让模型根据这个不完整的句子来预测被盖住的 token 是 "湾",此时我们有个词汇表,比如这个词汇表一共有8个词,"湾" 是第3个,则 "湾" 这个 token 的真值就是

text{GT}=[0,0,1,0,0,0,0,0]

,只需要让模型的输出和这个

text{GT}

越接近越好。

但是 CV 任务没有这个词汇表啊,假设我盖住一个 patch,让模型根据这个不完整的 image 来预测被盖住的 patch 是什么。那么对应的这个

text{GT}

是什么呢?

BEIT 通过一种巧妙的方式解决了这个问题。

假设这个问题可以得到解决,我们就能够用 masked image modeling 的办法 (和BERT类似,盖住图片的一部分之后预测这部分) 训练一个针对图片的预训练模型,这个预训练模型就也可以像 BERT 一样用在其他各种 CV 的下游任务中啦。

3 BEIT 原理分析

论文名称:BEIT: BERT Pre-Training of Image Transformers

论文地址:

https://arxiv.org/pdf/2106.08254.pdf

本文提出的这个方法叫做 BEIT,很明显作者是想在 CV 领域做到和 NLP 领域的 BERT 一样的功能。在第1篇文章中提到,训练好的 BERT 模型相当于是一个 Transformer 的 Encoder,它能够把一个输入的 sentence 进行编码,得到一堆 tokens。比如输入 "台湾大学",通过 BERT 以后会得到4个 tokens。并且这4个 tokens 也结合了sentence 的上下文。

那 BEIT 能不能做到类似的事情呢?,即能够把一个输入的 image 进行编码,得到一堆 vectors,并且这些个 vectors 也结合了 image 的上下文。

答案是肯定的。BEIT 的做法如下:

在 BEIT 眼里,图片有 2 种表示的形式:

image → image patches | visual tokens

在预训练的过程中,它们分别被作为模型的输入和输出,如下图9所示。

图9:图片有 2 种表示的形式:image patches or visual tokens

BEIT的结构可以看做2部分,分别是:

  • BEIT Encoder
  • dVAE

BEIT Encoder 类似于 Transformer Encoder,是对输入的 image patches 进行编码的过程,dVAE 类似于 VAE,也是对输入的 image patches 进行编码的过程,它们的=具体会在下面分别详细介绍。

3.1 将图片表示为 image patches

将图片表示为 image patches 这个操作和 Vision Transformer 对图片的处理手段是一致的。首先把

color{crimson}{xin mathbb{R}^{H times W times C}}

的图像分成

N=HW/P^2

个展平的2D块

color{orange}{x^{p} in mathbb{R}^{N times (P^2 C)}}

式中,

C

是 channel 数,

(H,W)

是输入的分辨率,

(P,P)

是块大小。每个 image patch 会被展平成向量并通过线性变换操作 (flattened into vectors and are linearly projected)。这样一来,image 变成了一系列的展平的2D块的序列,这个序列中一共有

N=HW/P^2

个展平的2D块,每个块的维度是

(P^2cdot C)

实作时

P=16,H=W=224

,和 ViT 一致。

问:image patch 是个扮演什么角色?

答: image patch 只是原始图片通过 Linear Transformation 的结果,所以只能保留图片的原始信息 (Preserve raw pixels)

3.2 将图片表示为 visual tokens

这一步是啥意思呢?BEIT的一个通过 dVAE 里面一个叫做 image tokenizer 的东西,把一张图片

color{crimson}{xin mathbb{R}^{H times W times C}}

变成离散的 tokens

color{teal}{z = [z_1, dots, z_N] in mathcal{V}^{h times w}}

。字典

mathcal{V} = { 1, dots, |mathcal{V}| }

包含了所有离散 tokens 的索引 (indices)。

要彻底理解如何将图片表示为 visual tokens,那就得先从 VAE 开始讲起了,熟悉 VAE 的同学可以直接跳过3.2.1。

  • 3.2.1 变分自编码器 VAE

VAE 跟 GAN 的目标基本是一致的——希望构建一个从隐变量

Z

生成目标数据

X

的模型,但是实现上有所不同。更准确地讲,它们是假设了

Z

服从某些常见的分布(比如正态分布或均匀分布),然后希望训练一个模型

X=g(Z)

,如下图10所示,这个模型能够将原来的概率分布映射到训练集的概率分布,也就是说,它们的目的都是进行分布之间的变换。

图10:生成模型的难题就是判断生成分布与真实分布的相似度,因为我们只知道两者的采样结果,不知道它们的分布表达式

图10里面的

Z

服从标准的正态分布,那么我就可以从中采样得到若干个

Z_{1},Z_{2},...,Z_{n}

,然后对它做变换得到

tilde{X_{1}}=g(Z_{1}),tilde{X_{2}}=g(Z_{2}),...,tilde{X_{n}}=g(Z_{n})

。注意这些

tilde{X_{i}}

都是通过

g

重构出来的数据集,那如何衡量

tilde{X_{i}}

的分布与目标的数据集分布是不是一样的呢?注意在这里我们只有一堆重构出来的数据

tilde{X_{i}}

,但并不知道

tilde{X_{i}}

的分布是啥,所以没法用KL散度来衡量

tilde{X_{i}}

的分布与目标的数据集分布的关系,因为KL散度是根据两个概率分布的表达式来算它们的相似度的。我们只有一批从构造的分布采样而来的数据

left{ tilde X_{1},tilde X_{2},...,tilde X_{n} right}

,还有一批从真实的分布采样而来的数据

left{ X_{1},X_{2},...,X_{n} right}

(也就是我们希望生成的训练集)。我们只有样本本身,没有分布表达式,当然也就没有方法算KL散度。

上面的假设是直接从正态分布中采样,实际情况是

Z

由真实的分布采样而来的数据

left{ X_{1},X_{2},...,X_{n} right}

计算得到,并希望它接近标准正态分布。之后的步骤不变,假设

p(X|Z)

描述了一个由

Z

来生成

X

的模型,而我们假设

Z

服从标准正态分布,也就是

p(Z)sim N(0,1)

。那么就可以按照图10的做法,从先从标准正态分布中采样一个

Z

,然后根据

Z

来算一个

X

。接下来就是结合自编码器来实现重构,保证有效信息没有丢失,再加上一系列的推导,最后把模型实现。框架的示意图如下图11所示。表达式为:

p(X)=sum_{Z}^{}{p(X|Z)p(Z)}tag{1}

图11:VAE的传统理解

但如果像这个图的话,我们其实完全不清楚:究竟经过重新采样出来的

Z_{k}

,是不是还对应着原来的

X_{k}

。换句话说,采样得到的几个值和原来的几个样本的对应关系没有了。比如在图11里面,

Z_1

X_1

是对应的吗,显然不是,因为

Z_1

是从正态分布里面随机采样得到的。

所以我们如果直接最小化

D(tilde X_{k},X_{k})^{2}

(这里

D

代表某种距离函数)是很不科学的,而事实上你看代码也会发现根本不是这样实现的。

所以 VAE 到底是如何实现的呢?

在整个VAE模型中,我们并没有去使用

p(Z)

(隐变量空间的分布)是正态分布的假设,我们用的是假设

p(Z|X)

(后验分布)是正态分布,如下图12所示。

图12:事实上,vae是为每个样本构造专属的正态分布,然后采样来重构

具体来说,给定一个真实样本

X_{k}

,我们假设存在一个专属于

X_{k}

的分布

p(Z|X_{k})

,并进一步假设这个分布是(独立的、多元的)正态分布。为什么要强调“专属”呢?因为我们后面要训练一个生成器

X=g(Z)

,希望能够把从分布

p(Z|X_{k})

采样出来的一个

Z_{k}

还原为

X_{k}

。如果假设

p(Z)

是正态分布,然后从

p(Z)

中采样一个

Z

,那么我们怎么知道这个

Z

对应于哪个真实的

X

呢?现在

p(Z|X_{k})

专属于

X_{k}

,我们有理由说从这个分布采样出来的

Z

应该要还原到

X_{k}

中去。

这时候每一个

X_{k}

都配上了一个专属的正态分布,才方便后面的生成器做还原。但这样有多少个

X

就有多少个正态分布了。我们知道正态分布有两组参数:均值

mu

和方差

sigma^{2}

(多元的话,它们都是向量),那我怎么找出专属于

X_{k}

的正态分布

p(Z|X_{k})

的均值和方差呢?就是通过一个神经网络来拟合出来。我们构建两个神经网络

μ_{k}=f_{1}(X_{k})

logsigma_{k}^{2}=f_{2}(X_{k})

来算它们了。我们选择拟合

logsigma_{k}^{2}

而不是直接拟合

sigma_{k}^{2}

,是因为

sigma_{k}^{2}

总是非负的,需要加激活函数处理,而拟合

logsigma_{k}^{2}

不需要加激活函数,因为它可正可负。到这里,我能知道专属于

X_{k}

的均值和方差了,也就知道它的正态分布长什么样了,然后从这个专属分布中采样一个

Z_{k}

出来,然后经过一个生成器得到

tilde X_{k}=g(Z_{k})

,现在我们可以放心地最小化

D(tilde X_{k},X_{k})^{2}

,因为

Z_{k}

是从专属

X_{k}

的分布中采样出来的,这个生成器应该要把开始的

X_{k}

还原回来。

图13:均值方差通过一个神经网络来拟合出来

为什么要让均值方差拟合网络的输出分布

p(Z|X)

接近标准正态分布

N(0,I)

如果没有这个约束的话,最终会得到什么结果呢?目标函数是最小化

D(tilde X_{k},X_{k})^{2}

,这个重构过程受到噪声的影响,因为

Z_{k}

是通过重新采样过的,不是直接由encoder算出来的。显然噪声会增加重构的难度,不过好在这个噪声强度(也就是方差)通过一个神经网络算出来的,所以最终模型为了重构得更好,肯定会想尽办法让方差为0。而方差为0的话,也就没有随机性了,生成的数据的范围就只在

left{ X_{1},X_{2},...,X_{n} right}

**里面了,模型会慢慢退化成普通的AutoEncoder,**所以不管怎么采样其实都只是得到确定的结果(也就是均值),只拟合一个当然比拟合多个要容易,而均值是通过另外一个神经网络算出来的。

说白了,模型会慢慢退化成普通的AutoEncoder,噪声不再起作用。

所以,VAE 这里还让所有的

p(Z|X)

,即均值方差拟合网络的输出分布接近标准正态分布

N(0,I)

,这样就防止了噪声为零,同时保证了模型具有生成能力。怎么理解“保证了生成能力”呢?如果所有的

p(Z|X)

都很接近标准正态分布

N(0,1)

,那么根据定义:

p(Z)=sum_{X}^{}{p(Z|X)p(X)}=sum_{X}^{}{N(0,1)p(X)}=N(0,1)tag{2}

这样我们就能达到我们的先验假设:

p(Z)

是标准正态分布,也就克服了退化成普通的AutoEncoder的问题。

那怎么均值方差拟合网络的输出分布

p(Z|X)

接近标准正态分布

N(0,I)

如果没有外部知识的话,其实最直接的方法应该是在重构误差的基础上中加入额外的loss:

L_{mu}=||f_{1}(X_{k})||^{2}\L_{sigma^{2}}=||f_{2}(X_{k})||^{2}tag{3}

因为它们分别代表了均值

μ_{k}

和方差的对数

logsigma_{k}^{2}

,达到

N(0,1)

就是希望二者尽量接近于0了。不过,这又会面临着这两个损失的比例要怎么选取的问题,选取得不好,生成的图像会比较模糊。所以,原论文直接算了一般(各分量独立的) 正态分布与标准正态分布的KL散度

KL(N(μ,σ^{2})||N(0,1))

作为这个额外的loss:

一元正态分布的情形即可,根据定义我们可以写出

KL(N(μ,σ^{2})||N(0,1))
=int_{}^{}frac{1}{sqrt{2pisigma^{2}}}e^{-(x-mu)^{2}/2sigma^{2}}(logfrac{e^{-(x-mu)^{2}/2sigma^{2}}/sqrt{2pisigma^{2}}}{e^{-(x)^{2}/2}/sqrt{2pi}})dx
=int_{}^{}frac{1}{sqrt{2pisigma^{2}}}e^{-(x-mu)^{2}/2sigma^{2}}(logleft{ frac{1}{sqrt{sigma^{2}}}expleft{ frac{1}{2}[x^{2}-(x-mu)^{2}/sigma^{2}] right} right})dx
=frac{1}{2}int_{}^{}frac{1}{sqrt{2pisigma^{2}}}e^{-(x-mu)^{2}/2sigma^{2}}[-logsigma^{2} x^{2}-(x-mu)^{2}/sigma^{2}]dx
=frac{1}{2}({-logsigma^{2} mu^{2} sigma^{2}-1})

整个结果分为三项积分,第一项实际上就是

-logsigma^{2}

乘以概率密度的积分(也就是1),所以结果是

-logsigma^{2}

;第二项实际是正态分布的二阶矩,熟悉正态分布的朋友应该都清楚正态分布的二阶矩为

mu^{2} sigma^{2}

;而根据定义,第三项实际上就是**“-方差除以方差=-1”**。所以总结果就是上面的最后一行。

所以正态分布与标准正态分布的KL散度

KL(N(μ,σ^{2})||N(0,1))

这个额外的loss计算结果为:

L_{mu,sigma^{2}}=frac{1}{2}sum_{i=1}^{d}{mu_{i}^{2} sigma_{i}^{2}-logsigma_{i}^{2}-1}tag{4}

这里的

d

是隐变量

Z

的维度,而

mu_{i}

sigma_{i}

分别代表一般正态分布的均值向量和方差向量的第

i

个分量。直接用这个式子做补充loss,就不用考虑均值损失和方差损失的相对比例问题了。显然,这个loss也可以分两部分理解:

begin{align} L_{mu,sigma^{2}}=L_{mu} L_{sigma^{2}}\L_{mu}=frac{1}{2}sum_{i=1}^{d}{mu_{i}^{2}}=frac{1}{2}||f_{1}(X)||^{2}\L_{sigma^{2}}=frac{1}{2}sum_{i=1}^{d}{(sigma_{i}^{2}-logsigma_{i}^{2}-1)} end{align} tag{5}

所以 VAE 模型的总的损失函数为:

L=E_{xsim tilde p(x)}[E_{zsim tilde p(z|x)}[color{purple}{-logq(x|z)} color{crimson}{KL(p(z|x)||q(z))}]] tag{6}

式中,紫色部分

color{purple}{-logq(x|z)}

就是重构损失

D(tilde X_{k},X_{k})^{2}

,红色部分

color{crimson}{KL(p(z|x)||q(z))}

就是均值方差拟合网络的输出分布

p(Z|X)

接近标准正态分布

N(0,I)

的这部分。

  • 3.2.2 BEIT 里的 VAE:tokenizer 和 decoder

上面我们了解了 VAE 模型的训练过程,那么我们回到之前的问题上面,BEIT 是如何将图片表示为 visual tokens的呢?

具体而言,作者训练了一个 discrete variational autoencoder (dVAE)。训练的过程如下图14所示。读者可以仔细比较一下这个 dVAE 和上文介绍的 VAE 的异同,dVAE 虽然是离散的 VAE,但它和 VAE 的本质还是一样的,都是把一张图片通过一些操作得到隐变量,再把隐变量通过一个生成器重建原图。下表就是一个清晰的对比,我们可以发现:

  • VAE使用图13所示的均值方差拟合神经网络得到隐变量
  • dVAE使用Tokenizer得到隐变量
  • VAE使用图12所示的生成器重建原图
  • dVAE使用Decoder重建原图

model

得到隐变量的模块

重建原图的模块

VAE

均值方差拟合神经网络

生成器

dVAE

Tokenizer

Decoder

图14:训练 discrete variational autoencoder (dVAE) 的过程

所以dVAE中的Tokenizer就相当于是VAE里面的均值方差拟合神经网络,dVAE中的Decoder就相当于是VAE里面的生成器。

所以,dVAE 的训练方式其实可以和 VAE 的一模一样。

问:这里的 visual token 具体是什么形式的?

答: 作者把一张 224×224 的输入图片通过 Tokenizer 变成了 14×14 个 visual token,每个 visual token 是一个位于[1,8192]之间的数。就像有个 image 的词汇表一样,这个词汇表里面有 8192 个词,每个 16×16 的image patch会经过 Tokenizer 映射成

|mathcal{V}|

里面的一个词。因为 visual token 是离散的数,所以优化时没法求导,所以作者采用了 gumbel softmax 技巧,想详细了解 gumbel softmax trick 的同学可以参考下面的链接:

科技猛兽:PyTorch 32.Gumbel-Softmax Trick

https://zhuanlan.zhihu.com/p/166632315

  • 3.2.3 BEIT 的 Backbone:Image Transformer

BEIT 的总体结构如下图15所示,BEIT 的 Encoder 结构就是 Transformer 的 Encoder,模型架构是一样的。图片在被分成

N=HW/P^2

个展平的2D块

color{orange}{x^{p} in mathbb{R}^{N times (P^2 C)}}

之后,通过线性变换得到

{E} x^{p}_{i}

,其中

{E} in mathbb{R}^{(P^2 C) times D}

。在concat上一个 special token [S]。这里作者还给输入加上了 1D 的位置编码

{E} _{pos} in mathbb{R}^{N times D}

,所以总的输入张量可以表示为:

{H}_0 = [ {e}_{texttt{[S]}} , {E} x^{p}_{i} , dots , {E} x^{p}_{N} ] {E}_{pos} tag{7}

图15:BEIT 的总体结构

输入 BEIT 的 Encoder (就是 Transformer 的 Encoder) 之后,张量依次通过

L

个 Encoder Block:

{H}^l = mathrm{Transformer}({H}^{l-1} ) tag{8}

式中

l=1,2,...,L

。最后一层输出

{H}^{L} = [ {h}^{L}_{texttt{[S]}}, {h}^{L}_{1}, dots, {h}^{L}_{N} ]

作为 image patches 的 encoded representations,

{h}^{L}_{i}

代表第

i

个 image patch的编码表示。

  • 3.2.4 类似 BERT 的自监督训练方式:Masked Image Modeling

至此,我们介绍了 BEIT 的两部分结构:

  • BEIT Encoder
  • dVAE

下面就是 BEIT 的训练方法了。既然BEIT 是图像界的 BERT 模型,所以也遵循着和 BERT 相似的自监督训练方法。BERT 的自监督训练方法忘了的同学请再看一遍图1-图6。

BERT 看很多的句子,随机盖住一些 tokens,让 BERT 模型预测盖住的tokens是什么,不断计算预测的 token 与真实的 token 之间的差异,利用它作为 loss 进行反向传播更新参数,来达到 Self-Supervised Learning 的效果。

BEIT 使用了类似 BERT 的自监督训练方式:Masked Image Modeling,如图15所示,即:

BEIT 看很多的图片,随机盖住一些 image patches,让 BEIT 模型预测盖住的patches是什么,不断计算预测的 patches 与真实的 patches 之间的差异,利用它作为 loss 进行反向传播更新参数,来达到 Self-Supervised Learning 的效果。

具体做法是:

  1. 给定输入图片
{x}

  1. 使用3.2.1节的做法把它变成
N

image patches

{{x}^{p}_{i}}_{i=1}^{N}

  1. 使用3.2.2节的做法把它变成
N

visual tokens

{z_{i}}_{i=1}^{N}

  1. 随机盖住40% 的 image patches,盖住的位置可以表示为
mathcal{M} in {1,dots,N}^{0.4 N}

  1. 把盖住的这40%的image patches 替换成可学习的编码
{e}_{texttt{[m]}} in mathbb{R}^{D}

  1. 现在这个输入的 image patches 就可以表示成:
x^{mathcal{M}} = { {x}^p_i : i notin mathcal{M} }_{i=1}^N bigcup { {e}_{texttt{[M]}} : i in mathcal{M} }_{i=1}^N\
  1. 把这个
x^{mathcal{M}}

通过

L

层的 BEIT Encoder,得到

{ {h}^{L}_{i} }_{i=1}^{N}

,表示输入image patches的编码表示。

  1. 盖住的位置的输出
left{boldsymbol{h}_{i}^{L}: i in mathcal{M}right}_{i=1}^{N}

去通过一个分类器,去预测盖住的这个

mathrm{patch}

的相应的 visual token, 就像图3所示在 BERT 里面把盖住的部分通过一个分类器, 去预测盖住的token:

p_{mathrm{MIM}}left(z^{prime} mid x^{mathcal{M}}right)=operatorname{softmax}_{z^{prime}}left(boldsymbol{W}_{c} boldsymbol{h}_{i}^{L} boldsymbol{b}_{c}right)

。式中

x^{mathcal{M}}

是盖住之后的所有 image patches,

boldsymbol{W}_{c} in mathbb{R}^{|mathcal{V}| times D}, boldsymbol{b}_{c} in mathbb{R}^{|mathcal{V}|}

, 这里

|mathcal{V}|=8192, quad D

是模型的 Embedding dimension.

  1. BERT 的训练目标是最小化计算预测的 token 与真实的 token 之间的差异,所以BEIT的目标也是最小化计算预测的 token 与真实的 token 之间的差异。那其实还不完全一致,在 BEIT 里面,假设我盖住第
i

个patch,毫无疑问它对应的 visual token 应该是

z_i

,这时候我希望 Encoder 输出的第

i

个位置的东西通过分类器之后是

z_i

的概率最大,即:

begin{align} max sum_{x in mathcal{D}} E_{mathcal{M}} left[ sum_{i in mathcal{M}} logcolor{teal}{ p_{text{MIM}}( z_i | x^{mathcal{M}} )} right] label{eq:objective} end{align} tag{9}

式中,

D

是全部的无标签训练数据,

mathcal{M}

是随机盖住的位置,

x^{mathcal{M}}

是盖住以后的corrupted image。

9 式是什么意思呢?

i in mathcal{M}

就是对盖住的每个 patches,BEIT 的 Encoder 在这个位置的输出

{h}^{L}_{i}

通过线性分类器

text{softmax}_{z'} ({W}_{c} {h}^{L}_{i} {b}_{c})

之后得到预测的 visual token真实 patches 对应的 visual token 越接近越好,如下图16所示。

问:真实 patches 对应的 visual token 是怎么得到的呢?

答: 如3.2.2节介绍。训练一个 dVAE,其中的 Tokenizer 的作用就是把 image patches 编码成 visual tokens,通过 Tokenizer 来实现。

图16:BEIT的训练方法:对盖住的每个 patches,BEIT 的 Encoder 在这个位置的输出通过线性分类器之后得到预测的 visual token 与真实 patches 对应的 visual token 越接近越好

下面的问题是如何随机盖住40% 的 image patches?

BEIT 并不是完全随机地盖住40%,而是采取了 blockwise masking 的方法,如下图17所示。

就是每次循环先通过 Algorithm 1计算出

s,r,a,t

,然后盖住

iin[t,t a),jin[l,l b)

的部分,直到盖住的部分超过了 40% 为止。

图17:Blockwise masking 的方法

  • 3.2.5 BEIT 的目标函数:VAE 视角

下面的问题是:BEIT 具体是取去优化什么目标函数呢?

回顾上式 6,VAE 模型的总的损失函数为:

L=E_{xsim tilde p(x)}[E_{zsim tilde p(z|x)}[color{purple}{-logq(x|z)} color{crimson}{KL(p(z|x)||q(z))}]] tag{6}

式中,紫色部分

color{purple}{-logq(x|z)}

就是重构损失

color{purple}{D(tilde X_{k},X_{k})^{2}}

,红色部分

color{crimson}{KL(p(z|x)||q(z))}

就是均值方差拟合网络的输出分布

p(Z|X)

接近标准正态分布

N(0,I)

的这部分。

  • dVAE 模型通过 Tokenizer 把 input image 变成一些 visual tokens,这个过程可以用
color{darkgreen}{q_{phi}(z|x)}

来表示。

  • 通过 Decoder 把 visual tokens 重建成 reconstructed image,这个过程可以用
color{darkgreen}{p_{psi}(x|z)}

来表示。

  • BEIT的 Encoder 也把 masked image 变成 visual tokens,这个过程可以用
color{darkgreen}{p_{theta}(z|tilde x)}

来表示。

那么

color{darkgreen}{p(x | tilde{x})}

代表什么含义呢?

因为

tilde{x}

代表 masked image,

x

代表 original image,所以

color{darkgreen}{p(x | tilde{x})}

代表给定一张masked image,能够重建回原图的概率,且对于

color{darkgreen}{p(x | tilde{x})}

我们有evidence lower bound (ELBO):

begin{align} sum_{(x_i, tilde{x}_i) in mathcal{D}} log{p(x_i | tilde{x}_i)} geq sum_{(x_i, tilde{x}_i) in mathcal{D}} big( underbrace{E_{z_isim q_{phi}( mathbf{z}|x_i)}[log p_{psi}(x_i|z_i)]}_{text{Visual Token Reconstruction}} - KL[q_{phi}(mathbf{z}|x_i) , p_{theta}(mathbf{z}| tilde{x}_i) ] big) label{eq:elbo} end{align} tag{10}

在 10 式中也标出了重构损失 Visual Token Reconstruction。传统 VAE 的重构损失 (Reconstruction loss)

color{purple}{D(tilde X_{k},X_{k})^{2}}

在 dVAE 里面可以写成:

color{purple}{E_{zsim q_{phi}( mathbf{z}|{x})}[log p_{psi}({x}|{z})]} tag{11}

所以整个优化过程分为2步:

第1步是去优化这个重构损失 (Reconstruction loss),就是更新 Tokenizer 和 Decoder 的参数,同时保持 BEIT的 Encoder 参数不变。

-E_{z_isim q_{phi}( mathbf{z}|x_i)}[log p_{psi}(x_i|z_i)] tag{obg.1}

第2步是去优化 BEIT 的 Encoder 参数,让预测的 visual token真实 patches 对应的 visual token 越接近越好,同时保持 Tokenizer 和 Decoder 的参数不变。

-log p_{theta}(hat{z}_i | tilde{x}_i) tag{obj.2}

所以总的优化目标可以写成:

上式11就是 BEIT 的总目标函数,使用 Gradient Ascent 更新参数。

所以,BEIT 遵循 BERT 的训练方法,让 BEIT 看很多的图片,随机盖住一些 image patches,让 BEIT 模型预测盖住的 patches 是什么,不断计算预测的 patches 与真实的 patches 之间的差异,利用 12 式进行反向传播更新参数,来达到 Self-Supervised Learning 的效果。

不同的是,BERT 的 Encoder 输入是 token,输出还是 token,让盖住的 token 与输出的预测 token 越接近越好;而 BEIT 的 Encoder 输入是 image patches,输出是 visual tokens,让盖住的位置输出的 visual tokens 与真实的 visual tokens 越接近越好。真实的 visual tokens 是通过一个额外训练的 dVAE 得到的。

  • 3.2.6 BEIT 的架构细节和训练细节超参数

BEIT Encoder 的具体架构细节: 12层 Transformer,Embedding dimension=768,heads=12,FFN expansion ratio=4,Patch Size=16,visual token总数,即词汇表大小

|mathcal{V}| =8192

。Mask 75个 patches,一个196个,大约占了40%。

BEIT Encoder 的具体训练细节: 在 ImageNet-1K上预训练。

参数

Batch Size

2000

Epochs

800

优化器和参数

Adam, 0.9, 0.999

Learning rate

1.5e-3

Warmup epochs

10

cosine learning rate decay

weight decay=0.5

数据增强

random resized cropping, horizontal flipping, color jittering。

  • 3.2.7 BEIT 在下游任务 Fine-tuning

使用 Self-Supervised Learning 预训练完的 BEIT,作者展示了在2种下游任务上微调的结果,分类和分割。这里只以分类为例展示做法。以分类为例如下图18所示,我们拿着训练好的 BEIT Encoder,给它添加一个 分类层 (Linear Transformation),池化层 (Avg),和激活函数,我们只微调分类层 (Linear Transformation),池化层 (Avg),和激活函数的参数,Encoder的参数保持不变 BEIT 在下游任务 Fine-tuning

图18:BEIT 在下游分类任务 Fine-tuning

  • 3.2.8 实验

分类实验

BEIT 实验的具体做法遵循3.2.7节的BEIT 在下游任务 Fine-tuning的做法,展示的都是预训练模型在具体小数据集上面 Fine-tune之后得到的结果。分类实验在CIFAR-10和ImageNet这两个数据集上进行,超参数设置如下图19所示:

图19:CIFAR-10和ImageNet超参数

下图20是实验在CIFAR-10和ImageNet这两个数据集上的性能以及与其他模型的对比。所有的模型大小都是 "base" 级别。与随机初始化训练的模型相比,作者发现预训练的BEIT模型在两种数据集上的性能都有显著提高。值得注意的是,在较小的CIFAR-100数据集上,从头训练的ViT仅达到48.5%的准确率。相比之下,通过Pre-train的帮助,BEIT达到了90.1%。结果表明,BEIT可以大大降低有标签数据 (labeled data) 的需求。BEIT还提高了ImageNet上的性能。

图20:BEIT在CIFAR-10和ImageNet这两个数据集上的性能以及与其他模型的对比

此外,作者将BEIT与21年几个最先进的 Transformer 自监督方法进行比较,如 DINO 和 MoCo v3 (这2个模型也会在这个系列中解读)。我们提出的方法在ImageNet微调上优于以往的模型。BEIT在ImageNet上的表现优于DINO,在CIFAR-100上优于MoCo v3。此外,作者评估了我们提出的方法与 Intermediate Fine-tuning。换句话说,我们首先以自监督的方式对BEIT 进行预训练,然后用标记数据在 ImageNet 上对预训练的模型进行 Fine-tune。结果表明,在ImageNet上进行 Intermediate Fine-tuning 后获得额外的增益。

问:图20中的 Supervised Pre-Training on ImageNet 和 Supervised Pre-Training, and Intermediate Fine-tuning on ImageNet有什么区别?

答: 二者都是使用全部的 ImageNet-1K 数据集。前者是只训练分类器的参数,而 BEIT 预训练模型参数不变。后者是既训练分类器的参数,又微调 BEIT 预训练模型参数。

作者也在 384×384 高分辨率数据集上面作 Fine-tune 了 10个epochs,同时patch的大小保持不变,也就是用了序列长度增加了。结果如下图21所示,在ImageNet上,更高的分辨率可以提高1个点的。更重要的是,当使用相同的输入分辨率时,用 ImageNet-1K 进行预训练的BEIT-384 甚至比使用 ImageNet-22K 进行监督预训练的 ViT-384 表现更好。

图21:Top-1 accuracy on ImageNet-1K

作者进一步扩大了 BEIT 的规模 (扩大到与 ViT-L 相同)。如上图21所示,在ImageNet上,从头开始训练时,ViT-384-L 比 ViT-384差。结果验证了 Vision Transformer 模型的 data hungry 的问题。解决方法就是用更大的数据集 ImageNet-22K,用了以后 ViT-384-L 最终比ViT-384 涨了1.2个点。相比之下,BEIT-L比 BEIT 好2个点,BEIT-384-L 比 BEIT-384 好1.7个点,说明大数据集对BEIT的帮助更大。

对比实验:

消融实验分别是在ImageNet (分类) 和 ADE20K (分割) 任务上进行的,自监督方式训练 epochs是300。

第1个探索Blockwise masking的作用。Blockwise masking 指的是图17的方法,发现它在两种任务中都是有利的,特别是在语义分割上。

第2个探索 recover masked pixels的作用,recover masked pixels指的是盖住一个 image patch,BEIT 的 Encoder 模型不输出visual token,而是直接进行 pixel level的回归任务,就是直接输出这个 patch,发现这样也是可以的,只是精度稍微变差了。这说明预测 visual tokens 而不是直接进行 pixel level的回归任务才是 BEIT 的关键。

第3个探索 1,2 的结合方案,去掉Blockwise masking,以及直接进行 pixel level的回归任务,这个性能是最差的。

第4个探索不进行自监督预训练,即直接盖住100%的image patches,性能也会下降。

图22:BEIT 对比实验

下图23是BEIT模型不同reference points的attention map,可视化的方法是拿出BEIT的最后一个layer,假定一个参考点,随机选定它所在的patch,比如是第57个patch,然后把attention map的第57行拿出来,代表这个第57号patch attend to所有patch的程度,再reshape成正方形就得到了下图23。

可以发现仅仅是预训练完以后,BEIT 就能够使用 self-attention 来区分不同的语义区域。这个性质表明了为什么 BEIT 能够帮助下游任务的原因。通过BEIT获得的这些知识有可能提高微调模型的泛化能力,特别是在小数据集上。

图23:不同reference points的attention map

总结:

BEIT 遵循 BERT 的训练方法,让 BEIT 看很多的图片,随机盖住一些 image patches,让 BEIT 模型预测盖住的patches是什么,不断计算预测的 patches 与真实的 patches 之间的差异,利用它作为 loss 进行反向传播更新参数,来达到 Self-Supervised Learning 的效果。

不同的是,BERT 的 Encoder 输入是 token,输出还是 token,让盖住的 token 与输出的预测 token 越接近越好;而 BEIT 的 Encoder 输入是 image patches,输出是 visual tokens,让盖住的位置输出的 visual tokens 与真实的 visual tokens 越接近越好。真实的 visual tokens 是通过一个额外训练的 dVAE 得到的。

参考:

变分自编码器(一):原来是这么一回事 - 科学空间|Scientific Spaces

https://spaces.ac.cn/archives/5253

本文亮点总结

1.BERT 可以做的事情也就是Transformer 的 Encoder 可以做的事情,就是输入一排向量,输出另外一排向量,输入和输出的维度是一致的。那么不仅仅是一句话可以看做是一个sequence,一段语音也可以看做是一个sequence,甚至一个image也可以看做是一个sequence。

2.不同的是,BERT 的 Encoder 输入是 token,输出还是 token,让盖住的 token 与输出的预测 token 越接近越好;而 BEIT 的 Encoder 输入是 image patches,输出是 visual tokens,让盖住的位置输出的 visual tokens 与真实的 visual tokens 越接近越好。

0 人点赞