BEIT: BERT Pre-Training of Image Transformers
代码：unilm/beit at master · microsoft/unilm (github.com)
期刊/会议：ICLR 2022
摘要
本文介绍了一种自监督视觉表示模型BEIT，即图像transformer的双向编码器表示。继自然语言处理领域开发的BERT之后，我们提出了一个掩码图像建模(mask image modeling)任务来预训练vision transformer。具体来说，每张图像在我们的预训练中有两个视图，即图像补丁(image patch)(如16×16像素)和视觉标记(visual token)(即离散标记)。我们首先将原始图像“标记化(tokenize)”为视觉标记。然后我们随机屏蔽(randomly mask)一些图像补丁(image patch)，并将它们输入transformer骨干网络。预训练的目标是根据损坏的图像补丁恢复原始视觉标记。在预训练BEIT之后，我们通过在预训练的编码器上附加任务层，直接对下游任务的模型参数进行微调。在图像分类和语义分割上的实验结果表明，该模型与之前的预训练方法相比，取得了较好的效果。
1、简介
Transformer在计算机视觉中取得了良好的性能。然而，实证研究表明，Vision transformer比卷积神经网络需要更多的训练数据。为了解决数据饥饿问题，自监督预训练是利用大规模图像数据的一种有前途的解决方案。针对vision transformer已经探索了几种方法，如对比学习和自蒸馏。
同时，BERT在自然语言处理方面取得了巨大的成功。它的掩码语言建模任务(masked language modeling tasking)首先随机掩码文本中一定比例的token，然后根据损坏文本的Transformer编码结果恢复mask token。在BERT的激励下，我们转向去噪自动编码的思想来预训练vision transformer，这在视觉领域还没有得到很好的研究。直接将BERT-style预训练应用于图像数据具有挑战性。首先，对于vision Transformer的输入单元，即图像patch，没有预先存在的词汇表。因此，我们不能简单地使用一个softmax分类器来预测所有可能的mask patch候选。相比之下，语言词汇，如单词和BPE，定义良好，易于自动编码预测。一个简单的替代方案是将该任务视为一个回归问题，它预测被mask patch的原始像素。然而，这种像素级的恢复任务往往会浪费建模能力在预训练的短程依赖(short-range dependencies)和高频细节(high-frequency detail)。我们的目标是克服上述问题，进行vision transformer的预训练。
在这项工作中，我们介绍了一种自监督视觉表示模型BEIT，它代表图像transformer的双向编码器表示。受BERT的启发，我们提出了一个预训练任务，即掩模图像建模(masked image modeling, MIM)。如图1所示，MIM对每个图像使用了两个视角，image patch和visual token。我们将图像分割成一个由image patch组成的网格，这些patch是Transformer骨干网络的输入表示形式。此外，我们将图像“tokenize”为离散的visual token，这是由离散VAE的潜在代码获得的。在预训练过程中，我们随机mask一定比例的图像patch，并将mask后的图像输入到Transformer。该模型学习恢复原始图像的visual token，而不是mask patch的原始像素。
我们执行自监督学习，然后在两个下游任务(即图像分类和语义分割)上对预训练的BEIT进行微调。实验结果表明，BEIT优于从零开始训练和以前的强自监督模型。此外，BEIT是对有监督的预训练的补充。通过使用ImageNet标签进行中间微调，BEIT的性能可以进一步提高。消融实验表明，我们提出的技术对于图像数据的BERT-style预训练的有效性至关重要。除了性能外，改进的收敛速度和微调的稳定性降低了下游任务的训练成本。此外，我们证明了自监督BEIT可以通过预训练学习合理的语义区域，释放图像中包含的丰富的监督信号。
我们的贡献主要如下：
- 我们提出了一个掩码图像建模任务，以自监督的方式预训练vision transformer。我们还从变分自编码器(variational autoencoder)的角度进行了理论解释。
- 我们预训练了BEIT，并对下游任务(如图像分类和语义分割)进行广泛的微调实验。
- 我们提出了自监督BEIT的自注意机制学习区分语义区域和对象边界，尽管没有使用任何人工标注。
2、方法
给定一个输入图像x xx, BEIT将其编码为情景化向量表示。如图1所示，BEIT由掩码图像建模(MIM)任务以自监督学习方式进行预训练。MIM的目标是基于编码向量恢复掩码图像patch。对于下游任务(如图像分类和语义分割)，我们在预先训练的BEIT上附加任务层，并对特定数据集上的参数进行微调。
2.1 图像表征
在我们的方法中，图像有两个表示视图，即image patch和visual token。这两种类型分别用作预训练期间的输入和输出表示。
2.1.1 image patch
在我们的实验中，我们将每张224 × 224的图像分割为14 × 14个图像patch网格，其中每个补丁为16 × 16。
2.2 骨干网络：图像transformer
继ViT之后，我们采用标准Transformer作为骨干网络。因此，结果可以直接与之前的工作在网络架构方面进行比较。
2.3 预训练BEIT：掩码图像建模
我们提出了一种掩码图像建模(MIM)任务。我们随机屏蔽(random mask)一定比例的图像patch，然后预测对应于被mask patch的visual token。
其中D 为训练语料库，M 为随机掩码位置，x M 为根据M进行掩码的损坏图像。
在我们的工作中，我们采用逐块屏蔽(blockwise masking)，而不是随机选择掩码位置M的patch。根据算法1的总结，每次都要mask一组图像patch。对于每个块，我们将patch的最小数量设置为16。然后我们随机为掩蔽块选择一个纵横比。我们重复上述两步，直到获得足够多的mask patch，即0.4N，其中N为图像patch总数，0.4为掩码比。
MIM任务在很大程度上受到掩码语言建模(masked language model, MLM)的启发，这是自然语言处理中最成功的预训练目标之一。此外，block-wise(或n-gram) masking也广泛应用于BERT-like模型。然而，直接使用像素级自动编码(即恢复被mask patch的像素)进行视觉预训练会使模型专注于短距离依赖关系和高频细节。BEIT通过预测离散的visual token来克服上述问题，它将细节总结为高级抽象。3.3节中的消融实验表明，我们提出的方法显著优于像素级自编码。
BEIT预训练可视为变分自编码器训练。设x 表示原始图像，
表示掩码图像，z 表示visual token。考虑到对数似然
的证据下限(evidence lower bound, ELBO)，即从其损坏的版本中恢复原始图像:
第二项是我们的BEIT预训练的目标。
2.5 预训练的设置
BEIT的网络架构遵循ViT-Base的网络架构，便于比较。我们使用一个隐藏大小为768的12层transformer和12个attention head。前馈网络的中间层为3072。我们使用默认的16 × 16输入patch大小。我们直接借用DaLL-E训练的图像tokenizer。visual token的词汇量为8192。
我们在ImageNet-1K的训练集上预训练BEIT，该训练集包含约1.2M张图像。我们的数据增强策略包括随机调整大小裁剪，水平翻转，颜色抖动(color jittering)。请注意，我们没有在自监督学习中使用标签。实验中使用的分辨率为224 × 224。因此输入被分割为14 × 14个图像patch，以及等量的visual token。我们random mask最多75个patch(即大约占总图像patch的40%)。
预训练运行大约500k步(即800个epoch)， 2k批处理大小。采用Adam ， β 1 = 0.9 ， β 2 = 0.999 进行优化。学习率设置为1.5e-3，warmup10个epoch，余弦学习率衰减。权重衰减为0.05。我们采用了随机深度(random depth)，速率为0.1，并禁用了dropout。使用16张Nvidia Telsa V100 32GB GPU卡，50k步训练大约需要5天。
我们发现适当的初始化对于稳定Transformer非常重要，特别是对于大规模的预训练。我们首先在一个小范围内随机初始化所有参数，例如[−0.02,0.02]。然后，对于第L个Transformer层，我们将自注意模块和前馈网络的输出矩阵(即每个子层内的最后一个线性投影)重新缩放
2.6 对下游视觉任务进行微调
在预先训练BEIT之后，我们将一个任务层附加到Transformer之上，并对下游任务(如BERT)的参数进行微调。我们以图像分类和语义分割为例。利用BEIT在其他视觉任务上利用预训练-然后微调范式是很简单的。
语义分割。对于语义分割，我们遵循SETRPUP中使用的任务层。具体来说，我们使用预训练的BEIT作为骨干编码器，并将几个反卷积层作为解码器来产生分割。该模型也是端到端的微调类似于图像分类。
中间微调。在自我监督的预训练之后，我们可以在数据丰富的中间数据集(即我们工作中的ImageNet-1K)上进一步训练BEIT，然后在目标下游任务上对模型进行微调。这种中间微调是NLP中BERT微调的常见做法。我们直接按照BEIT的方法。
3、实验
我们对图像分类和语义分割进行了完整的微调实验。此外，我们提出了用于预训练的各种消融实验，并分析了BEIT学习到的表征。我们也在附录D中报告了ImageNet上的线性探索。
3.1 图像分类
图像分类任务将输入的图像分类为不同的类别。我们在ILSVRC-2012 ImageNet数据集上使用1k个类和1.3万张图像评估BEIT。在我们的微调实验中，我们直接遵循了DeiT的大部分超参数，以便进行公平的比较。与从头开始训练相比，我们减少了微调时间，因为BEIT已经进行了预训练。因此，我们使用更大的学习率与分层衰减。详细的超参数概述在附录H中。
表1报告了图像分类的top-1精度。我们将BEIT与随机初始化、监督预训练和以前的自监督学习方法训练的vision transformer进行比较。除了iGPT参数为1.36B外，所有比较模型均为base大小。在ImageNet上进行预训练以进行比较，但ViT-JFT300M是在谷歌的内部300M图像上进行预训练。
与随机初始化训练的模型相比，我们发现预先训练的BEIT显著提高了在两个数据集上的性能。BEIT提高了在ImageNet上的性能，显示了在丰富资源设置下的有效性。
此外，我们将BEIT与以前最先进的Transformer自监督方法进行了比较，例如DINO和MoCo v3。我们提出的方法在ImageNet微调方面优于以前的模型。其中，iGPT-1.36B使用了更多的参数(即1.36B vs 86M)， ViT-JFT300M在更大的语料上进行预训练(即300M vs 1.3M)，其他ViT-Base则在ImageNet-1K上进行预训练。iGPT-1.36B和ViT-JFT300M是最具可比性的方法，同样遵循vision transformer的自动编码预训练。具体来说，iGPT使用聚集的图像token作为图像GPT或图像BERT的输入和输出。相比之下，我们使用图像patch作为输入来保存原始像素，并使用离散的visual token作为预测瓶颈。ViT-JFT300预测每个mask token的平均3位颜色，而不是由dVAE学习的visual token。我们还以多任务学习方式预训练BEIT和DINO的自我监督任务，详见附录E。
此外，我们评估了我们提出的方法与中间微调。换句话说，我们首先以自监督的方式预训练BEIT，然后在ImageNet上使用标记数据对预训练的模型进行微调。结果表明，BEIT是监督预训练的补充，在ImageNet上进行中间微调后获得额外的增益。
在384x384分辨率大小下进行微调。在分辨率为224 × 224的微调之后，我们又在384×384图像上对模型进行了10个epoch的微调。我们遵循DeiT的标准高分辨率设置，除了使用更少的epoch。请注意，对于224 × 224和384 × 384图像，我们保持patch大小相同。因此，transformer的输入序列长度随着分辨率的提高而变长。表1显示，较高的分辨率使BEIT结果在ImageNet上提高了1+点。更重要的是，当使用相同的输入分辨率时，在ImageNet-1K上预训练的BEIT384甚至优于使用ImageNet-22K的监督预训练ViT384。
缩放到更大的尺寸。我们进一步将BEIT扩大到大尺寸(与ViT-L相同)。如表1所示，在从头训练时，ViT384-L在ImageNet上比ViT384差。结果验证了vision transformer的数据饥饿问题。在ImageNet-22K上的监督式预训练部分缓解了这个问题，ViT384-L最终比ViT384高出1.2%。相比之下，BEIT-L优于BEIT 2.0%, BEIT384-L优于BEIT384 1.7%。换句话说，将BEIT从base扩展到large的好处大于使用ImageNet-22K进行监督式预训练。更重要的是，BEIT384与在ImageNet-22K上进行监督预训练的ViT384相比，随着规模从base(即0.6)扩大到large(即1.1)，BEIT的改进更大。结果表明，BEIT倾向于对超大型模型(如1B或10B)提供更多帮助，特别是当标注数据不足以对此类大型模型进行有监督的预训练时。
收敛曲线。图2比较了从零开始训练和先训练再微调范式的收敛曲线。我们发现，对BEIT进行微调不仅可以获得更好的性能，而且收敛速度也比从头训练DeiT快得多。此外，微调BEIT可以在很短的时间内达到合理的数字。
3.2 语义分割
语义分割的目的是为输入图像的每个像素预测相应的类别。我们使用25K图像和150个语义类别在ADE20K基准上评估BEIT。我们报告了所有语义类别的mIoU的度量。如2.6节所述，我们直接遵循任务层和SETR-PUP 中描述的大多数超参数。在ADE20K上，我们使用Adam作为优化器。学习率设置为1e-3，类似于图像分类的分层衰减。我们对160K步进行微调。批大小为16。详细的超参数描述在附录I中。
如表3所示，我们将BEIT与依赖于ImageNet标记数据的监督预训练进行了比较。我们发现，我们提出的方法比有监督的预训练获得了更好的性能，尽管BEIT不需要对预训练进行人工标注。此外，我们在ImageNet上对BEIT进行中间微调，即我们首先在ImageNet上对预训练的BEIT进行微调，然后在ADE20K上对模型进行微调。结果表明，中间微调进一步提高了语义分割的BEIT。
3.3 消融实验
我们进行消融实验，以分析BEIT中每个组件的贡献。模型基于图像分类(即ImageNet)和语义分割(即ADE20K)进行评估。我们将消融研究的默认预训练步骤设置为300个epoch，这是之前实验中使用的总步骤的37.5%。
表4报告了各种模型变量的结果。首先，我们通过随机采样掩码(random sample masked)位置来消除分块掩码(blockwise masking)。我们发现区块掩码(blockwise masking)在这两种任务中都是有益的，特别是在语义分割上。其次，我们通过预测被mask patch的原始像素来去除visual token的使用，即将预训练任务变成一个像素回归问题来恢复被mask patch。我们提出的掩码图像建模任务显著优于朴素像素级自动编码。与表1的结果相比，在两项任务上，消融的结果都不如从零开始训练vision Transformer。结果表明，visual token的预测是BEIT的关键因素。第三，我们将visual token和blockwise masking一起使用。我们发现blockwise masking对像素级的自动编码更有帮助，这减轻了短距离依赖的缺陷。第四，恢复所有visual token会损害下游任务的性能。第五，我们将BEIT与不同的训练步骤进行比较。对模型进行更长的预训练可以进一步提高下游任务的性能。
3.4 自注意图分析
我们证明了BEIT中的自注意机制可以分离物体，尽管我们的预训练根本不依赖于任何手动标注。探测图像(probing image)取自MS COCO语料库，以避免出现在预训练数据中。
如图2所示，我们绘制了图像中不同参考点的自注意图。可视化结果由最后一层的query-key点积计算的注意力分数产生。对于每个参考点，我们使用相应的patch作为查询，并显示它处理哪个patch。经过预训练后，BEIT学会了在没有任何特定任务监督的情况下，使用自注意头来区分语义区域。该属性部分地表明了BEIT能够帮助下游任务的原因。BEIT获得的这些知识可能会提高微调模型的泛化能力，特别是在小规模数据集上。
4、相关工作
自我监督视觉表征学习。多年来，已经引入了各种方法来以自监督的方式来预训练视觉模型。先前的研究设计了巧妙的clever pretext tasks，如预测patch顺序、着色和预测旋转角度。此外，有人提出对图像中的部分patch进行掩码，并对每个掩码位置的掩码patch进行真假分类。该方法类似于Jigsaw预训练的mask version。最近的研究方向遵循对比范式。该模型通常将不同的数据增强视为图像的不同视图，然后使正对的表示相似，同时将负对推开。为了在对比学习中获得足够信息量的负样本，方法通常依赖于较大的记忆库或大batch_size。BYOL和SimSiam 进一步消除了对负样本的要求，使用各种技术来避免表示崩溃。另一种方法使用聚类来组织图像示例。
自监督的vision transformer。由于数据饥饿的问题，预训练vision transformer最近受到了极大的关注。iGPT首先通过k-means聚类RGB像素创建一个9位调色板，然后使用聚类token来表示图像。接下来iGPT使用BERT和GPT的任务来预训练transformer。相比之下，我们提出的方法使用图像patch作为输入，而不会丢失像素级信息。此外，我们的visual token是通过离散的VAE而不是聚类获得的。ViT使用mask patch预测任务进行初步探索，预测mask patch的3位平均颜色。报告像素级自动编码性能更差，尽管它是BERT从NLP到CV的最直接的转换。我们提出的模型利用dVAE学习的visual token，而不是使用启发式设计的预训练任务，这不仅实现了更好的性能，而且在理论上也有更好的动机。除掩码自编码外，其他主流研究工作采用对比学习和自蒸馏。相比之下，BEIT在训练前吞吐量(附录E)和内存消耗方面可以实现数倍的改进。BEIT的优势使其吸引了visual transformer的规模。
5、总结
我们为vision transformer引入了一个自监督的预训练框架，在图像分类和语义分割等下游任务上实现了强大的微调结果。我们表明，所提出的方法对于使BERT-like预训练(即带有掩码输入的自动编码)对vision transformer工作良好至关重要。我们还展示了自动获得的语义区域知识的有趣属性，而不使用任何人类标注的数据。在未来，我们希望在数据大小和模型大小方面扩大BEIT预训练。此外，我们将以更统一的方式进行多模态预训练，使用相似的目标和文本和图像的共享架构。
