BEIT:基于图像重建进行预训练!微软提出BEIT,Top-1准确率达86.3%!代码已开源!

2021-12-02 16:40:36 浏览数 (1)

写在前面

在本文中,作者介绍了一种自监督视觉表示模型BEITB idirectional E ncoder representation from I mage T ransformers)。继BERT在自然语言处理领域开发之后,作者提出了一种用于预训练视觉Transformer的masked image modeling 任务。

具体来说,在本文的预训练中,每个图像都有两个视图,即图像 patche (如16×16像素)和视觉token (即离散token)。作者首先将原始图像“tokenize”为视觉token。然后随机mask一些图像patch,并将它们输入主干Transformer中。预训练的目标是基于损坏的图像patch恢复原始视觉token

在预训练BEIT之后,作者通过在预训练编码器上附加任务层,直接微调下游任务上的模型参数。在图像分类和语义分割方面的实验结果表明,该模型与以前的预训练方法相比取得了较好的效果。

1. 论文和代码地址

BEIT: BERT Pre-Training of Image Transformers

代码语言:javascript复制
论文:https://arxiv.org/abs/2106.08254
代码:https://github.com/microsoft/unilm/tree/master/beit

2. Motivation

Transformer在计算机视觉方面取得了非常不错的性能。然而,实证研究表明,视觉Transformer比卷积神经网络需要更多的训练数据。为了解决data-hungry的问题,自监督预训练是利用大规模图像数据的一种很有前景的解决方案。

目前,BERT在自然语言处理方面取得了巨大成功。它的 masked language modeling任务首先随机mask了文本中部分token,然后根据损坏文本的Transformer编码结果恢复mask的token。

受BERT的启发,作者采用去噪自编码的思想对视觉Transformer进行预训练,这一点在CV界还没有得到很好的研究。直接对图像数据进行Bert形式的预训练是一个挑战。对于vision Transformer的输入单元,即图像patch,没有预先存在的词汇表。

因此,不能简单地使用softmax分类器来预测所有可能的mask patch候选。相反,语言词汇,如单词和BPE,定义明确,易于自编码预测。一个简单的替代方法是将任务视为回归问题,直接预测mask patch的原始像素。然而,这种像素级恢复任务往往会在预训练的局部相关性和高频细节上浪费建模能力 。本文的目标是克服上述问题,对视觉Transformer进行预训练。

在这项工作中,作者提出了一种自监督视觉表示模型BEIT。受BERT的启发,作者提出了一个预训练任务,即 masked image modeling(MIM) 。MIM对每个图像使用两个视图,即图像patch和视觉token。

作者将图像分割成一个网格的patches,这些patch是主干Transformer的输入表示。此外,作者将图像“tokenize”为离散视觉token,这是通过离散VAE的潜在代码获得的。在预训练期间,作者随机mask部分图像patch,并将损坏的输入到Transformer。该模型学习恢复原始图像的视觉token,而不是mask patch的原始像素。

3. 方法

给定一个输入图片x,BEIT将其编码为上下文化的向量表示。如上图所示,BEIT通过掩蔽图像建模(MIM)任务以自监督学习方式进行预训练。MIM的目标是基于编码矢量恢复被mask的图像块。对于下游任务(如图像分类和语义分割),作者在预训练的BEIT上附加任务层,并微调特定数据集上的参数。

3.1 Image Representations

在本文的方法中,图像有两个表示视图,即图像patch视觉token 。这两种类型分别作为预训练的输入和输出表示。

3.1.1 Image Patch

2D图像被分割成一系列patch,以便标准Transformer可以直接接受图像数据。形式上,作者reshape图片为个patch,

其中C是通道数,(H,W)是输入图像分辨率,(P,P)是每个patch的分辨率。图像patch

被展平成向量并线性投影,这类似于BERT中的单词嵌入。图像块保留原始像素,并在BEIT中用作输入特征。

在本文的实验中,作者将每幅224×224的图像分割成一个14×14的图像网格,其中每个图像网格的大小为16×16。

3.1.2 Visual Token

与自然语言类似,作者将图像表示为基于“image tokenizer”获得的离散token序列,而不是原始像素。具体来说,作者将图片

标记化为

其中词汇

包含离散的token索引。

作者通过离散变分自编码器(DVAE)学习image tokenizer。在视觉token学习过程中,有两个模块,即tokenizer解码器 。tokenizer

据视觉码本(即词汇表)将图像像素x映射为离散的tokens z。解码器

学习基于token z重建输入图像x。

重建目标可以写成

。由于潜在的视觉token是离散的,模型训练是不可微的。因此,作者采用Gumbel-softmax用来训练模型。此外,在dVAE训练过程中,在

上加入一个统一的先验。

作者将每个图像标记为14×14的视觉token grid。词汇表大小设置为

预训练目标是最大限度地提高损坏图像中正确视觉token的对数可能性:

max sum_{x in mathcal{D}} mathbb{E}_{mathcal{M}}left[sum_{i in mathcal{M}} log p_{mathrm{MIM}}left(z_{i} mid x^{mathcal{M}}right)right]

其中,为训练语料库,表示随机mask的位置,表示根据 得到的损坏图像。

作者在本文中采用了blockwise mask ,而不是随机选择mask的patch。如上面的算法所示,每次都会mask一块图像patch。对于每个block,作者将最小patch数设置为16。然后,随机选择mask patch的纵横比。重复上述两个步骤,直到获得足够的mask patch,即0.4N,其中N是图像patch的总数,0.4是 masking ratio。

MIM任务受masked language modeling的启发,blockwise mask也广泛应用于类BERT模型中。然而,直接使用像素级自编码进行视觉预训练会促使模型关注局部相关性和高频细节。BEIT通过预测离散的视觉token克服了上述问题,它将细节总结为了高级抽象。

3.4 Fine-Tuning BEIT on Downstream Vision Tasks

在预训练BEIT之后,作者在Transformer上附加一个任务层,并微调下游任务上的参数。本文以图像分类和语义分割为例。通过BEIT,在其他视觉任务上利用预训练然后微调范式也是很简单的。

Semantic segmentation

对于语义分割,作者遵循SETR-PUP中使用的任务层。具体来说,使用预训练的BEIT作为主干编码器,并结合几个deconvolution层作为解码器来产生分割。该模型也进行了端到端的微调,类似于图像分类。

Intermediate fine-tuning

在自监督的预训练之后,可以在数据丰富的中间数据集上进一步训练BEIT,然后在目标下游任务上微调模型。这种中间微调是NLP中BERT微调的常见做法,因此可以直接遵循BEIT的方法。

4.实验

4.1 Image Classification

上表展示了不同方法在CIFAR-100和ImageNet-1K数据集上的Top-1准确率。可以看出,相比于其他自监督学习方法,本文的方法能够实现更高的性能。

Fine-tuning to 384×384 resolution

上表展示了ImageNet数据集上,用384×384的分辨率进行微调的实验结果,可以看出,本文的方法依旧能够达到更高的实验性能。

Convergence curves

上图展示了DeiT和BEiT的收敛曲线,可以看出,BEiT比DeiT收敛更快,性能更高。

4.2 Semantic Segmentation

上表展示了本文方法在ADE20K数据集上进行语义分割任务的实验结果。

4.3 Ablation Studies

上表展示了本文提出的不同模块和训练策略,在ImageNet分类和ADE20K语义分割上进行消融实验的结果。

4.4 Analysis of Self-Attention Map

上图展示了BEiT不同的query点进行Self-Attention的Attention Map可视化结果。

5. 总结

在这项工作中,作者提出了一个用于视觉Transformer的自监督预训练框架,对下游任务(如图像分类和语义分割)微调之后,实现了非常好的实验性能。

作者表明,所提出的方法对于使类似于BERT的预训练(即 auto-encoding with masked input)能够很好地用于图像Transformer是至关重要的。

这篇文章跟Kaiming最新的那篇文章不同之处是,这篇文章把图片离散化成了一个一个的数字,所以可以和文本一样进行处理,另外,Kaiming的文章对patch进行了一定程度的移除,所以可以省计算量,而这篇文章却没有。

0 人点赞