【论文精读】BEiT

摘要

• 提出一种掩码图像建模任务，以自监督方式预训练视觉transformer。
• 对BEIT进行了预训练，并对下游任务进行微调实验，如图像分类和语义分割。
• 发现自监督BEIT的自注意力机制可以学会区分语义区域和对象边界。

框架

       给定输入图像$x$，BEiT将其编码为上下文向量表示。如上图，BEiT通过自监督学习掩码图像建模 (MIM) 任务进行预训练，恢复编码向量的masked image patches。对于下游任务 (如图像分类和语义分割)，在经预训练的BEiT上添加任务层，并对特定数据集上的参数进行有监督微调。

图像表示

       本方法中，图像有两种表示视图，即images patches和visual tokens。两者分别在预训练任务中作为输入和输出表示。

Images patch

       输入图像$x$被splited为一个序列的patches，适应标准Transformer的输入格式。具体，先将图像$x\in R^{H\times W\times C}$reshape为$N=HW/P^2$个patches为$x^p\in R^{N\times (P^{2}C)}$，其中$C$为通道数，$(H,W)$为输入图像分辨率，$(P,P)$为每个patch的分辨率。然后images patches{xip}i=1N{ \{x^p_i\} }^N_{i=1}{xip​}i=1N​被flattened为向量并经过线性投影层得到Patch Embedding，作为BEiT输入特征。

       实验中$P=16$，每幅 224×224 图像被splited成$14\times 14$的images patches grid$(N=14\times 14)$，每个grid表示的patch尺寸为$16\times 16$。

Visual token

       类似于自然语言处理，图像被表示为由image tokenizer获得的离散token序列。具体，采取离散变分自编码器(dVAE)学习到的图像tokenizer，将图像$x\in R^{H\times W\times C}$tokenized为$z=[z_1,...,z_N]\in {\nu }^{h\times w}$，其中词表 (vocabulary) ν={1,...,∣ν∣}ν={\{1,...,|ν|\}}ν={1,...,∣ν∣}包含离散的token索引。

       如dVAE，在visual token的学习过程中，有tokenizer和decoder两个模块。tokenizer $q_{\varphi }(z|x)$将图像像素$x$映射为离散tokens $z$。decoder $p_{\psi }(x|z)$学习基于visual token $z$重建输入图像$x$。重建的优化目标写成 $E_{z\sim q_{\varphi }(z|x)}[log{p}_{\psi }(x|z)]$。由于潜在的visual tokens是离散的，模型不可微。因此，采用Gumbel-softmax训练模型。在dVAE训练过程中，$q_{\varphi }$会加uniform先验。

       实验中每幅图像被tokenized为$14\times 14$的visual tokens grid，和images patches数相同。词表大小设置为 $|\nu |=8192$。

主干网络：Transformer

       使用标准Transformer作为骨干网络。Transformer的输入是images patches序列{xip}i=1N{ \{x^p_i\} }^N_{i=1}{xip​}i=1N​，images patcher经线性投影获得Patch Embedding$E{x}_i^p$，其中$E\in R^{(P^{2}C)\times D}$；然后，在输入序列前添加一个特殊标记$[S]$，在将标准的可学习1D position embedding $E_{pos}\in R^{N\times D}$添加到Patch Embedding中，得到Transformer的输入向量$H_0=[e_{[S]},E{x}_i^p,...,E{x}_N^p]+E_{pos}$。

       编码器包含$L$层的Transformer blocks$H^l=Transformer(H^{l-1}),l=1,...,L$。最后一层的输出向量$H^L=[h_{[S]}^L,h_1^L,...,h_N^L]$即为经编码的images patches(共N+1 个向量)，其中$h_i^L$是第$i$个图像块的images patcher向量。

预训练 BEiT：掩码图像建模

       结合上述，进行整体建模。给定一个输入图像$x$，将其split为 N个images patches({xip}i=1N)({\{x^p_i\}}^N_{i=1})({xip​}i=1N​)，并 tokenized为N个visual tokens({zi}i=1N)({\{ z_i\}}^N_{i=1})({zi​}i=1N​)。然后随机mask 40%（掩码率为 0.4）的images patches，其中masked的位置表示为M∈{1,...,N}0.4NM∈{ \{ 1,...,N \}}^{0.4N}M∈{1,...,N}0.4N，这些位置使用可学习的嵌入$e_{[M]}\in R^D$来替换（用于训练）。

       被masked后的图像xM={xip:i∉M}i=1N⋃{e[M]:i∈M}i=1Nx^M={\{ x^p_i:i∉M \}}^N_{i=1}⋃{\{ e_{[M]}:i∈M \}}^N_{i=1}xM={xip​:i∈/M}i=1N​⋃{e[M]​:i∈M}i=1N​将馈入$L$层Transformer中，得到的最后隐层向量{hiL}i=1N{\{h^L_i\}}^N_{i=1}{hiL​}i=1N​为输入patch的编码表示。对于每个mask位置的{hiL:i∈M}i=1N{\{h^L_i:i∈M\}}^N_{i=1}{hiL​:i∈M}i=1N​，使用softmax分类器预测相应的visual tokens $p_{MIM}(z^{\prime }|x^M)=softma{x}_{z\prime }(W_c{h}_i^L+b_c)$，其中$x^M$是被masked的images patches，$W_c\in R^{|\nu |\times D}$且$b_c\in R^{|\nu |}$。

       预训练的目标为最大化masked patches在图像中正确的visual tokens $z_i$的对数似然(log-likelihood)。
$$\max \sum _{x\in D}{\mathbb{E}}_M[\sum _{i\in M}\log p_{MIM}(z_i|x^M)]$$

       其中，D为训练语料库，M为随机masked的位置，$x^M$为根据M来masked的已损坏图像。

       本算法掩码采用blockwise masking，而非随机选择patches进行mask。如下图所示，每次先计算出$s,r,a,t$，在mask掉$i\in [t,t+a),j\in [l,l+b)$的部分。重复上述2个步骤，直到获得足够的masked patches。

       MIM主要受到MLM的启发，MLM是NLP中最成功的预训练目标之一。blockwise masking也被广泛应用于 BERT-like模型。然而，直接使用像素级自编码（即恢复 masked patches 的像素）进行视觉预训练，会促使模型聚焦于短程依赖和高频细节。BEiT通过预测离散visual tokens来克服上述问题，该tokens会将细节总结为高级抽象。

VAE分析

       BEiT预训练从理论上可视为变分自动编码器训练。设$x$为原始图像，$\hat{x}$为masked图像，$z$为visual tokens。考虑VAE的对数似然$p(x|\hat{x})$的下界 (ELBO)，带入解释即从其mask图像中恢复原始图像：
$$\sum _{(x_i,{\hat{x}}_i)\in D}\log p(x_i|{\hat{x}}_i)\ge \sum _{(x_i,{\hat{x}}_i)\in D}(\underset{VisualTokenReconstruction}{\underbrace{{\mathbb{E}}_{z_i\sim q_{\varphi }(z|x_i)}[\log p_{\psi }(x_i|z_i)]}}-D_{KL}[q_{\varphi }(z|x_i),p_{\theta }(z|{\hat{x}}_i)])$$

       其中，$q_{\varphi }(z|x)$表示图像tokenizer，即基于输入图像采用dVAE的Tokenizer获取visual tokens；$p_{\psi }(x|z)$表示输入visual tokens解码原始图像，即通过dVAE的Decoder把visual tokens重建成reconstructed image；$p_{\theta }(z|\hat{x})$表示基于masked图像恢复visual tokens，此即BEIT的预训练任务。

       总结为两阶段过程学习模型：第一阶段，使用uniform作为先验最小化图像重建（Visual Token Reconstruction）损失。即保持$p_{\theta }$固定的情况下，学习$q_{\varphi }$和$p_{\psi }$。第二阶段，保持$q_{\varphi }$和$p_{\psi }$固定的情况下，学习先验$p_{\theta }$。若$q_{\varphi }(z|x)$已为最优，则${\hat{z}}_i=argma{x}_z{q}_{\varphi }(z|x)$，则上条公式可重写为：
$$\sum _{(x_i,{\hat{x}}_i)\in D}\underset{Stage1:VisualTokenReconstruction}{\underbrace{({\mathbb{E}}_{z_i\sim q_{\varphi }(z|x_i)}[\log p_{\psi }(x_i|z_i)]}}+\underset{Stage2:MaskedImageModeling}{\underbrace{\log p_{\theta }({\hat{z}}_i|{\hat{x}}_i)}})$$
       Stage 2即为BEiT的预训练目标。

预训练设置

       BEIT的网络架构与ViT-Base相同。使用768隐层大小的12层Transformer，注意力头数为12。FFN的中间尺寸是3072，默认patch size尺寸为$16\times 16$，采样已训练好的图像tokenizer。视觉token词表大小为8192。

       在ImageNet-1K数据集上自监督预训练BEIT，其中包含大约1.2M张图像。数据增强策略包括随机调整大小的裁剪，水平翻转，颜色抖动。输入图像尺寸为$224\times 224$，输入被分割为$14\times 14$的images patches和相同数量的visual token。随机掩码最多75个patch(总图像补丁的40%)。

       预训练运行大约500k step(即800 epoch)，batchsize为2k。采用$\beta 1=0.9$, $\beta 2=0.999$的Adam优化。学习率设置为1.5e-3，预热10 epochs余弦学习率衰减，权重衰减是0.05。采用$rate=0.1$的stochastic depth（训练时每个batch随机dropout网络的一些layers，测试时用完整网络)，禁用dropout。

       适当的初始化对稳定Transformer训练很重要。故在一个小范围内随机初始化所有参数，例如[−0.02, 0.02]。对于第$l$层Transformer，用缩放率为$1/\sqrt{2l}$重新缩放自注意力模块和FFN的输出矩阵。

下游视觉任务微调 BEiT

       预训练BEiT后，可以Transformer末尾追加一个任务层，用于微调下游任务。以图像分类和语义分割为例。

• 图像分类：对于图像分类，直接使用一个简单的线性分类器作为任务层。具体地，类别概率计算为 softmax(avg({hiL}i=1NWc))softmax(avg(\{{h^L_i}\}^N_{i=1}W_c))softmax(avg({hiL​}i=1N​Wc​))，其中$h_i^L$是第$i$个images patcher的最终编码向量，$W_c\in R^{D\times C}$是一个参数矩阵，$C$是类别数。通过有监督训练，对有标签数据进行最大似然估计。
• 语义分割：对于语义分割，遵循SETRPUP中使用的任务层。具体地，使用预训练BEiT作为主干编码器，合并几个反卷积层 (deconv) 作为解码器来产生分割。
• 中间微调：经过自监督的预训练，BEiT可以在中间数据集上进一步训练，然后在目标下游任务上微调模型 (自监督预训练 → 中间数据集微调 → 下游任务微调)，可以直接采用BEIT的方法。

实验

图像分类实验

       以下为在ILSVRC-2012 ImageNet数据集上用1k类和1.3M张图像评估BEIT。

       BEIT实验的做法遵循下游任务Fine-tuning的做法，为预训练模型在具体小数据集上面Fine-tune之后得到的结果。分类实验在CIFAR-10和ImageNet这两个数据集上进行，训练配置如上图。

       上图为CIFAR-100上的测试top-1精度。观察到，在较小的CIFAR-100数据集上，从头训练的ViT仅达到48.5%精度，BEIT通过预训练实现了90.1%，实验结果表明BEIT可以大大降低标注工作量。BEIT也优于MoCo v3、DINO。另外，在ImageNet-1K上的中间微调可改善在CIFAR-100上的结果。

       上图显示了在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的帮助更大。

       当输入图像尺寸为384×384时，在数据集上Fine-tune10个epochs，patch的大小保持不变的情况下。观察到，更高的分辨率可以提高1个精度点。当使用相同的输入分辨率时，用ImageNet-1K进行预训练的BEIT-384比使用 ImageNet-22K进行监督预训练的 ViT-384表现更好。

       与最先进的监督预训练在ImageNet微调上进行了比较。除了使用ImageNet-1K还在ImageNet-22K上对BEIT进行了预训练150 epochs，在自监督预训练之后，在ImageNet-22K进行了90 epoch的中间微调。此外，使用一个内部数据集，其中约有7000万张带标签的图像，作为ImageNet-22K的直接替代。BEIT +为增加了LayerScale和相对位置偏差模块。

       上图结果看到。 BEIT-L在ImageNet-22K上进行了微调，取得了与在谷歌JFT-3B上训练的ViT-L相当的性能。此外，BEIT -L在对内部70M数据集进行中间微调后，在ImageNet上获得了89.5% top-1的精度。结果表明，BEIT预训练大大减少了所需的标记工作，并推进了大尺寸视觉transformer的新技术水平。

       上图显示了在ADE20K语义分割基准上的微调结果。在Swin中，使用与BEiT相同的任务层(即UperNet)并在640 × 640分辨率下评估模型。BEiT-L模型在ADE20K上获得最先进的性能。

       上图二比较了从零开始训练和先训练再微调两种范式的收敛曲线。微调BEIT不仅实现了更好的性能，而且收敛速度比从头训练DeiT快得多。

消融实验

       上图为在ImageNet (分类) 和 ADE20K (分割) 任务上进行的消融实验，行1观察到Blockwise masking在两种任务中都是有利的，特别是在语义分割上。行2为将预训练改为masked patches 的像素回归问题，观察到精度下降了一个点，行3在此之上减少 blockwise masking 的使用进一步降低了精度，发现blockwise masking 对像素级自动编码更有帮助，可缓解短距离依赖。行4预训练恢复100%的visual tokens会影响下游任务的性能，行5在此之上减少 masking 的使用会进一步降低精度。行6比较了不同训练 steps下的 BEiT，显示对模型进行更长时间的预训练可以进一步提高下游任务的性能。

       上图显示使用LayerScale和相对位置偏差都提高了ImageNet分类和ADE20K语义分割的性能。 根据经验注意到，当将模型扩展到数十亿参数时，vanilla Transformer是最稳定的，因此我们不将LayerScale用于超大型模型。

       上图为对比了在ImageNet-1K上训练的DALL-E的tokenizer和dVAE的tokenizer的性能，实验显示本文提出的tokenizer获得了更好的重构损失和ImageNet微调性能。

可视化

       上图绘制了图像中不同参考点的自注意力分布图，可视化通过网络最后一层的query、key的积计算出的注意力分数产生。观察到BEIT中自注意力机制可以分离对象，即使预训练完全不依赖于任何人工标注。

reference

Bao, H. , Dong, L. , & Wei, F. . (2021). Beit: bert pre-training of image transformers.