[论文解读] BEIT：适用于Image Transformer的新型预训练方法 （上：论文翻译）[原始代码已开源]

BEIT： BERT Pre-Training of Image Transformer

​ 我们提出了一种具有代表性的自监督模型——BeiT，它代表了Image Transformer的双向编码器表示方法。随着BERT在自然语言处理邻域的发展，我们提出了一种遮盖式图像的建模任务来预训练Vision Transformer。具体来说就是，每张图像有两种表示，一种是图像的patch（如 $16\times 16$的像素patch），另一种是视觉token（如不连续的位置表示）。我们首先将原始图像“标记化”为视觉token。然后我们随机遮盖一些图像的patch，并把它送入到Transformer的backbone中。预训练的目的是根据损坏的图像patch来恢复原始的视觉token。ERIT完成预训练后，我们直接在预训练的编码器上附加任务层来微调模模型在下游任务上的参数。实验结果表明，我们的方法模型在图像分类和语义分割上比先前的预训练方法更有竞争力。例如，基础版本的BEIT在ImageNet上实现了83.2%的准确率，显著优于相同配置下从头开始的DeiT（81.8%）。此外，large版本的BEIT仅使用ImageNet-1K是取得了86.3%的准确率，甚至优于在ImageNet-22k上预训练的ViT-L模型（85.2%）。代码和预训练模型可以从下面的链接中获取：https://aka.ms/beit.

1.Introduction

​ Transformer在计算机视觉已经取得了可喜的表现。但是，实证表明vision Transformer需要远比CNN更多的数据来训练。为了解决需要大量数据来训练的这个问题，自监督式的预训练是一个很有前途的方式，它可以利用大规模的图像数据。多种在Visoin Transformer上具有代表性的训练方法已经被提出，如对比学习（Chen et al.,2021;Xie et al.,2021）,自我蒸馏（Caron et al.,2021）。

​ 同时，BERT在自然语言处理（NLP）上已经取得了不错的进展。它遮盖语言建模任务的第一步是以某个比例随机遮盖文本中的token，然后根据损坏文本在Transformer编码后的结果来恢复被遮盖的token。收到BERT的影响，我们产生了将去噪自动编码的方法运用在vision Transforemr的预训练中，在视觉领域中还没有好好的研究过这个方法。想要把BERT中的预训练方法直接运用到图像数据中这对我们来说也是非常有挑战性的。首先，在vision Transformer中并没有预先的词汇作为输入单元，即图像patch。所以我们不能简单的使用softmax分类器来预测所有可能的被遮蔽斑块的候选人。与此相反，语言词汇（如单词和BPE），是定义明确且便于自动编码的。一个直接了当的做法是将此任务当作是一个回归任务来处理，来预测被遮挡的斑块的原始像素。然而，这种像素级的恢复任务往往会将建模能力浪费在预先训练短距离的以来和高频细节上。我们的目标是克服上述的问题使其能够在vision Transformer实现预训练。

​ 在这篇文章中，我们将介绍一种自监督的vision 表征模型——BEIT（基于图像Transformer的双向编码器表示）。收到BERT的启发，我们提出了一种名为遮盖图像建模（MIM）的预训练任务。如下图所示，MIM从图像中获取两种特征（图像patch，Vision token）。我们将图像分割成网格状的patch，并把它作为Transformer backbone的输入表征。此外，我们将原始的图像“标记化”为Vision token，这是由离散VAE（Ramesh et al.,2021）的潜在代码得到的。在与预训练过程中，我们按照某个比例随机的遮盖一些图像patch，并把残缺的图像输入到Transformer中。模型将去学着恢复原始图像的vision token来替代原始图像中被遮盖的部分。

​ 我们实验了自监督学，然后将预训练的BEIT模型在两个下游任务（图像分类，语义分割）上微调。实验结果表明，BEIT的结果优于从头开始训练和先前强大的预训练模型。此外，BEIT是对监督式预训练的一种补充。通过使用ImageNet标签进行中间微调，可以进一步提高BEIT的性能。消融研究表明，我们提出的技术对图像数据的BERT式预训练的有效性至关重要。除了性能之外，收敛速度的提升和微调的稳定性都可以减少在终端任务上的花费。另外，我们证明，自我监督的BEIT可以通过预训练学习合理的语义区域，释放出图像中包含的丰富的监督信号。

​ 我们的贡献总结如下：

​ 1.我们提出了一种遮盖图像的建模任务来自监督式的训练vision Transformer。我们还从变异自动编码器的角度提供了一个理论解释。

​ 2.我们预训练了BEIT并且在下游任务（图像分类、语义分割）上进行了广泛的微调实验。

​ 3.我们发现自监督式BEIT中的self-Attention机制可以学习区分语义区域和物体边界，尽管没有使用任何人类注释。

2. Methods

​ 给定一个输入图像$x$​​​，BEIT将其编码为上下文的矢量表征。就如上图Figure.1展示的一样，BEIT以自我监督的学习方式，通过遮蔽图像建模（MIM）任务进行预训练。MIM的主要目的是从编码的矢量中恢复被遮盖的图像patch。对于下游任务（如：图像分类、语义分割），我们在预训练的BEIT后添加任务层，并且在特殊的数据集上进行参数的微调。

2.1 Image Representations

​ 在我们的方法中，图像有两种表征方式，分别是图像patch，vision token。这两种方式分别作为预训练过程中的输入和输出表征。

2.1.1 Image Patch

​ 一张2D图像将会被分成patch的序列（[$h,w,c$]->[$n,patc{h}_L$​]），因此一个标准的Transformer可以直接接收这样的图像输入。正式的，我们将图像$x\in R^{H\times W\times C}$​拆分成$N=\frac{HW}{p^2}$​个patch，其中C是图像的通道数，（H,W）是图像的分辨率，（P，P）是每个patch的分辨率。图像的patch { x i p } i = 1 N \{x^p_i \}^N_{i=1} {xip​}i=1N​​​将会被展开成矢量，并且会被线性投影，这和BERT（Delib et al.,2019）中单词的处理方式很相似。​​图像patch保留了原始像素，在BEIT中将会作为输入的特征。

​ 在我们的实验中，我们以$14\times 14$​的patch网格来划分$224\times 224$​​的图像，每个patch的大小是$ 16\times 16$。​

2.1.2 Visual Token

​ 与自然语言相似，我们将图像表示为由 "图像标记器 "获得的离散标记序列，而不是原始像素。具体来说就是，我们将标记图像$x\in R^{H\times W\times C}$为$Z=[z_1,...,z_N]\in V^{w\times h}$,其中词表 v = { 1 , . . . , ∣ V ∣ } v=\{1,...,|V|\} v={1,...,∣V∣}包含离散的标记索引。遵循前人（Ramesh et al.,2021）的方法，我们通过离散变异自动编码器学习图像标记器（dVAE）。在Visual token的学习过程中有两个模块，分别是tokenizer和encoder。标记器$q_{\phi }(z|x)$根据视觉编码簿（visual codebook）图像像素x映射成离散的标记z。编码器$p_{\phi }(x|z)$根据visual token z来学习去重建输入图像x。重建的对象可以被写成$E_{z-q_{\phi }(z|x)}[log{p}_{\phi }(x|z)]$。由于潜在的visual token是离散的，模型的训练是无差异的。采用Gumbel-Softmax（Jang et al.,2017;Maddison et al.,2017）松弛法来训练模型参数。在dVAE训练过程中，一个统一的被事先放在$q_{\phi }$​上。更多关于image patch的信息可以从 Ramesh等人的论文中查看。

​ 我们将每张图像标记为一个$14\times 14$​的视觉标记token。请注意，一个图像的视觉标记数量和图像斑块数量是一样的。词表的尺寸被设置为$|V|=8192$​​。在我们的工作中，我们直接使用Ramesh等人论文中的描述的公开可以用的图像标记。

2.2 Backbone Network ：Image Transformer

​ 遵循ViT的设计原则，我们使用标准的Transformer作为我们的骨干网络。因此，在网络结构方面，结果可以直接与以前的工作进行比较。

​ Transformer的输入是图像patch的列表 { x i p } i = 1 N \{x^p_i\}^N_{i=1} {xip​}i=1N​​。然后对这些patch进行线性投影，以获得patch的矢量$E_{x_i^p}$（$E\in R^{(P^{2}C)\times D}$​）。此外，我们预加了一个特殊的token $[S]$到输入序列中。我们也将标准的可学习的1D的位置矢量$E_{pos}\in R^{N\times D}$保存下来并添加到patch列表中。最终，输入的矢量$H_0=[e_{[s]},E_{x_i^p},...,E_{x_N^p}]+E_{pos}$​将会被送入Transformer进行训练。编码器部分包含L个Transformer blocks$H^l=Transformer(H^{l-1}),(l=1,...,L)$。最后一层$H^L=[h_{[s]}^l,h_1^L,...,h_N^L]$​输出的矢量将会做为图像patch的编码表征，其中$h_i^L$是第i个图像patch矢量。

2.3 Pre-Training BEIT ：Masked Image Modeling

​ 我们提出了一种遮盖图像式建模任务来预训练BEIT。我们按照一定的比例随机遮盖图像的patch，然后预测与被遮蔽的patch相对应的Visual token。

​ Figure.1展示了我们的方法。如2.1小节提到的，给定一个输入图像$x$，我们将他分为$N$个图像patch( { x i p } i = 1 N \{x^p_i\}^N_{i=1} {xip​}i=1N​)，同时将它标记为$N$个visual token { z i } i = 1 N \{z_i\}^N_{i=1} {zi​}i=1N​。我们随机遮盖大约40%的图像patch，这白的位置标记为 M ∈ { 1 , . . . , N } 0.4 N M \in \{1,...,N\}^{0.4N} M∈{1,...,N}0.4N。接下来我们用可学习的向量$e_{[M]}\in R^D$来代替被遮盖的图像patch。这样残缺的图像patch（ x M = { x i p : i ∉ M } i = 1 N ⋃ { e [ M ] : i ∈ M } i = 1 N x^M = \{x^p_i:i\notin M \}^N_{i=1} \bigcup \{e_{[M]}:i \in M\}^N_{i=1} xM={xip​:i∈/​M}i=1N​⋃{e[M]​:i∈M}i=1N​）将会被输入到有L层的Transformer(见2.2小节)。最终的隐藏矢量 { H i L } i = 1 N \{H^L_i\}^N_{i=1} {HiL​}i=1N​被视为输入patch的编码表征。对于每个编码的位置 { h i L : i ∈ M } i = 1 N \{h^L_i:i \in M\}^N_{i=1} {hiL​:i∈M}i=1N​,我们使用softmax分类器来预测相应的visual token $p_{MIN}(z^{{}^{\prime }}|x^M)=softma{x}_{z^{{}^{\prime }}}(W_c{H}_i^L+b_c)$，其中$x^M$是遮盖后的图像，$W_c\in R^{|V|\times D}$，$b_c\in R^{||V}$。预训练的目的是，在给定损坏图像的情况下，最大限度地提高正确的视觉标记$z_i$​​的对数可能性:
$$max\sum _{x\in D}{E}_M[\sum _{i\in M}log{p}_{MIN}(z_i|x^M)]$$

其中D是训练的语料库，M代表随机遮盖的位置，$x^M$表示被遮盖后的图像。

​ 在我们的工作中，我们并没有采用随机遮盖位置M的patch，而是采用顺时针给的遮盖方式。正如下图伪代码总结的一样，每次都有一个图像patch被遮盖。对于每个区块，我们设置最小的区块数量为16。然后我们随机选择一个长宽比来遮盖区块。我们重复上述步骤直到获取到足够的masked patch（实验中我们的比例为0.4,直到获取mask patch超过0.4倍的patch数量位置）。

​ MIM任务是遮盖式语言建模（Devlin et al.,2021）的启发，这是自然语言处理中最成功的预训练目标之一。此外，顺势针遮盖也被广泛的应用在BERT相似的模型中(Joshi et al.,2020;Bao et al.,2020;Raffel et al.,2020)。然而，直接使用像素级自动编码进行视觉预训练，促使模型关注短距离的依赖性和高频细节。BEIT通过预测离散的visual token来克服上述的问题，它将细节总结为高层次的抽象概念。第3.3节中的消融研究表明，我们提出的方法明显优于像素级自动编码。

2.4 From the Perspective of Variational Autoencoder

​ BEIT的预训练可以被看作是变异自动编码器（Kingma and Welling ,2014）的训练。$x$代表原始的影像，$\tilde{x}$表示遮盖后的图像，$z$代表visual token。考虑到对数可能性$p(x|\tilde{x})$​的证据下限（ELBO），即从其损坏的版本中恢复原始图像:
$$\sum _{(x_i,{\tilde{x}}_i)\in D}\log p(x_i|{\tilde{x}}_i)\ge \sum _{(x_i,{\tilde{x}}_i)\in D}\underset{VisualTokenReconstruction}{\underbrace{E_{z_i-q_{\phi }(z|x_i)}[\log p_{\phi }(x_i|z_i)]}}-D_{KL}[q_{\phi }(z|x_i),p_{\theta }(z|{\tilde{x}}_i)]$$

第一个公式中$q_{\phi }(z|x)$表示获得visual token的图像标记器；第二个公式中$p_{\phi }(x|z)$对输入的visual token进行解码，以获得原始图像；第三个公式中$p_{\theta }(x|\tilde{x})$表示根据被遮挡的图像恢复visual token，这是我们的MIM预训练任务。

​ 我们按照一个类似于二阶段的程序来训练模型（van den Oord et al.,2017;Razavi et al.,2019）。在第一个stage，我们得到的图像标记器是一个离散变量自动编码器(Ramesh et al.,2021)。具体来说就是，第一阶段使重建损失（$-E_{z_i-q_{\varphi }(z|x_i)}[\log p_{\phi }(x_i|z_i)]$​）最小化，就如公式(2)所诉。在第二个stage，我们在保持$q_{\varphi }$​和$p_{\phi }$​不变的情况下学习先验$p_{\theta }$​​。我们将$q_{\varphi }(z_{|}{x}_i)$​简化为一个单点分布，与最可能的visual token ${\hat{z}}_i=argma{x}_z{q}_{\varphi }(z|x_i)$。那么方程（2）可以改写为:
$$\sum _{(x_i,{\tilde{x}}_i)\in D}(\underset{stage1:VisualTOkenReconstruction}{\underbrace{E_{z_i-q_{\varphi }(z|x_i)}[\log p_{\phi }(x_i|z_i)]}}+\underset{stage2:MaskedImageModeling}{\underbrace{\log p_{theta}({\hat{z}}_i|{\tilde{x}}_i)}})$$
其中第二项是我们的BEIT预训练目标。

2.5 Pre-Training Setup

​ BEIT的网络结构于ViT-base版相同，以便进行公平的比较。我们使用了一个有责12层Transformer，768的隐藏尺寸和12个注意力头。前馈网络的中级尺寸是3072。我们使用默认的$16\times 16$​的​输入patch尺寸。我们所使用的图像标记器（image tokenizer）是直接借用Ramesh训练的。visual token的词表尺寸是7192。

​ 我们在ImagetNet-1K（包含1.2M张图片，Russakovsky et al.,2015）上进行预训练。我们的增强策略包含随机裁剪缩放（random resized cropping），横向翻转（horizontal flipping）,颜色抖动（color jittering,Wu et al.,2018）。请注意，我们在自监督学习过程中并没有使用标签。在实验中我们使用$224\times 224$的分辨率。所以输入的图像被分成$14\times 14$​的图像patch，以及相同数量的visual token。我们随机遮盖最多75个patch（随机遮盖40%的图像patch）。

​ 整个预训练的过程使用2k的batch size训练了500k steps（800 epochs）。优化器Adam的参数是${\beta }_1=0.9,{\beta }_2=0.999$​​。学习率的初始设置1.5e-3,初始时使用10个epochs 的warmup，随后使用余弦退火的策略。weight decay是0.05。​我们采用比例为0.1的随机深度（Huang et al.,2016），并禁用dropout。整个500k steps的训练过程在使用16张 Tesla V100（32G）后大约需要5天时间。

2.6 Fine-Tuning BEIT on DownStream Vision Tasks

​ 在预训练完BEIT后，我们在Transformer后添加了一个任务层，并且在下游任务上对参数进行微调，就和BERT一样。在我们的实验中我们用图像分了和语义分割作为示例。在其他视觉任务上利用BEIT的预训练-再微调范式是很直接的。

Image Classification 对于图像分类任务，我们直接使用一个简单的线性分类器作为任务层。具体来说，我们使用使用 Averge Polling层来汇总Transformer输出的表征，然后将全局的信息直接输入到softmax分类器中。类别概率的计算方法是softmax（ a v g ( { h i L } i L W c ) avg(\{h^L_i\}^L_i W_c) avg({hiL​}iL​Wc​)）,其中$h_i^L$是最后的编码矢量的第i个图像patch，$W_c\in R^{D\times C}$是参数矩阵，C是标标签的数量。我们通过更新BEIT和softmax分类器的参数来最大化标记数据的可能性。

Semantic Segmentation 对于语义分割，我们使用SETR-PUP(zheng et al.,2020)作为任务层。具体来说，我们使用预训练的BEIT作为backbone encoder,将几个反卷积层（deconvolution）作为解码器，以产生分割效果。模型也是和图像分类一样的端到端的微调。

Intermediate fine-tune 在自监督预训练后，我们可以在一个有数据量的中间数据集（如 ImageNet-1k）上进一步训练BEIT，然后再目标下游任务上进行微调。这种中间微调是NLP（Pruksachatkun et al.2020）中BERT微调的常见做法。我们直接参照了BEIT的方法。

欢迎各位多翻译中出现的错误进行之争，详细论文待后续解读，现有的翻译基本可以理解作者的大体意思
原始论文见：https://arxiv.org/pdf/2106.08254v1.pdf
原文代码见：https://github.com/microsoft/unilm/tree/master/beit
后续更多的论文翻译解读正在进行中，敬请期待…