15 BEiT 论文精读：BERT Pre-Training of Image Transformers

在上一篇文章中我们精读了 MAE。MAE 的核心思想非常直接把图像中的大部分 patch 遮住只让编码器看到少量可见 patch然后通过解码器重建被遮住区域的像素。从直觉上看MAE 更像是一个“图像补全任务”模型需要根据可见部分推测缺失区域长什么样。而 BEiT 的思路则更接近 NLP 中的 BERT。BERT 的预训练任务是 Masked Language Modeling也就是遮住句子中的一部分词让模型根据上下文预测这些被遮住的词。BEiT 想做的事情就是能不能把图像也变成类似“词”的离散 token然后像 BERT 一样做 masked token prediction这也是 BEiT 这篇论文最核心的贡献它把图像自监督预训练从“像素重建”推进到了“视觉 token 预测”。BEiT 的全称是 BERT Pre-Training of Image Transformers论文由 Hangbo Bao、Li Dong、Songhao Piao 和 Furu Wei 提出发表于 ICLR 2022。论文提出了一种面向 Vision Transformer 的 Masked Image Modeling 方法输入是被遮挡的图像 patch预测目标不是原始像素而是由图像 tokenizer 生成的离散视觉 token。1、为什么需要 BEiT在 ViT 中一张图像会被切分成固定大小的 patch。例如输入图像大小为 224×224patch size 为 16×16那么整张图像会被切成 14×14196 个 patch。每个 patch 被拉平成向量再经过线性映射变成 Transformer 可以处理的 token 序列。这让图像看起来很像 NLP 中的文本序列。但是这里有一个关键问题文本中的 token 本身是离散的比如 word、subword 或 BPE token而图像 patch 本质上是连续的像素块。BERT 之所以能够做 masked language modeling是因为它有一个明确的词表。比如某个词被 mask 之后模型只需要在词表中预测它是哪一个 token。可是图像没有天然的“词表”。如果我们直接让模型预测被遮住 patch 的像素值那么任务就变成了像素级回归。论文认为这种像素级恢复容易让模型过度关注局部纹理、高频细节和短距离依赖而不一定能学到更高层的语义表示。BEiT 因此选择了另一条路线先把图像转换成离散的 visual token再让 ViT 去预测这些 token。简单来说BEiT 试图回答的问题是如果文本可以通过“遮词预测”学到语言表示那么图像能不能通过“遮 patch 后预测视觉词”学到视觉表示2、BEiT 的核心思想BEiT 的整体流程可以概括为三步第一步将原始图像切分成 patch作为 ViT 编码器的输入。第二步使用一个提前训练好的图像 tokenizer把原始图像转换成离散的 visual token。第三步随机遮挡一部分图像 patch让 ViT 根据未遮挡区域预测被遮挡位置对应的 visual token。因此BEiT 中同一张图像实际上有两种表示表示形式作用Image patches作为 Transformer 编码器的输入Visual tokens作为 masked image modeling 的预测目标BEiT 并不是直接拿被 mask 的 patch 去预测像素而是预测该位置对应的离散视觉 token。论文中使用的 tokenizer 来自离散 VAE也就是 dVAE。该 tokenizer 会把图像映射成一个 14×14 的离散 token 网格并且视觉词表大小为 8192。也就是说对于一张 224×224 的图像图像被切成 14×14 个 patchtokenizer 也生成 14×14 个 visual token每个图像 patch 都对应一个离散视觉 token预训练时模型只需要预测被 mask 位置的 token 类别。从这个角度看BEiT 真的很像图像版 BERT。3 BEiT 和 BERT 的对应关系为了更直观理解 BEiT我们可以把它和 BERT 做一个对应BERTBEiT文本句子图像单词 / subword token图像 patch词表中的离散 tokendVAE 生成的 visual tokenMasked Language ModelingMasked Image Modeling预测被 mask 的词预测被 mask patch 对应的 visual tokenTransformer EncoderVision Transformer EncoderBERT 中模型输入的是带有 [MASK] 的文本序列输出是被 mask 位置对应的词表概率。BEiT 中模型输入的是带有 mask embedding 的图像 patch 序列输出是被 mask 位置对应的 visual token 概率。所以 BEiT 的关键并不只是“遮住 patch”而是引入了一个视觉词表让图像预训练任务变成了分类问题而不是像素回归问题。4、BEiT 的模型结构BEiT 的主干网络基本沿用标准 ViT。给定一张输入图像首先将其划分为固定大小的 patch。假设图像大小为 H×Wpatch 大小为 P×P那么 patch 数量为N H × W / P²对于 224×224 图像和 16×16 patchN196。论文实验中BEiT-Base 使用 12 层 Transformer隐藏维度为 768注意力头数为 12FFN 中间维度为 3072整体配置与 ViT-Base 保持一致。在输入端BEiT 会做以下操作1. 将图像 patch 拉平成向量2. 通过线性投影得到 patch embedding3. 加入位置编码4. 对被 mask 的 patch用一个可学习的 mask embedding 替代5. 将整个序列输入 Transformer Encoder。经过多层 Transformer 编码后BEiT 只在被 mask 的位置上接一个分类头用于预测对应的 visual token。这个设计和 BERT 很接近模型并不需要重建所有位置只需要在被遮挡的位置上完成 token 预测。5、Masked Image ModelingBEiT 的预训练任务BEiT 的预训练任务叫做 Masked Image ModelingMIM。具体来说给定一张图像BEiT 会随机遮挡大约 40% 的 image patch。论文中使用的是 blockwise masking也就是按照块状区域进行遮挡而不是完全独立地随机遮挡每个 patch。这样做更符合图像的空间连续性也能避免任务过于简单。设原始图像经过 tokenizer 得到 visual tokenz [z1, z2, ..., zN]其中每个 zi 都是视觉词表中的一个离散编号。设被 mask 的位置集合为 M那么 BEiT 的目标就是根据被遮挡后的图像 xM预测这些位置上的原始 visual tokenmaximize Σ log p(zi | xM), i ∈ M用更直白的话说模型看到一张被遮住部分区域的图像然后要猜出这些被遮住区域原本对应哪个视觉 token。这与 MAE 有明显不同。MAE 的目标是重建像素masked patch → pixel valuesBEiT 的目标是预测视觉 tokenmasked patch → discrete visual token因此MAE 更偏向低层视觉重建而 BEiT 更强调离散语义预测。6、为什么 BEiT 不直接预测像素这是理解 BEiT 的关键。对于图像来说像素包含了大量低层信息比如颜色、边缘、纹理、局部噪声和高频细节。如果模型的目标是恢复像素它可能会花很多能力去拟合这些低层信息。但对于下游任务比如分类、检测、分割我们真正希望模型学到的是更抽象的语义表示比如哪些区域属于同一个物体哪些局部结构组成了目标图像中不同区域之间有什么语义关系背景和前景应该如何区分。BEiT 使用 visual token 作为预测目标本质上是在像素和语义之间加入了一个离散瓶颈。这个 tokenizer 会把连续图像压缩成离散 token模型不再需要精确恢复每个像素而是学习预测更抽象的视觉类别。论文的消融实验也验证了这一点当去掉 visual token改成直接恢复 masked pixels 时BEiT 的性能明显下降。论文表 4 中完整 BEiT 在 300 epoch 预训练后 ImageNet top-1 为 82.86ADE20K mIoU 为 44.65而改为恢复像素后ImageNet top-1 降到 81.04ADE20K mIoU 降到 41.38。这说明BEiT 的关键并不是单纯 mask而是 mask 之后预测什么。7、BEiT 的两阶段训练逻辑从训练流程上看BEiT 可以理解为两个阶段。1. 第一阶段训练或使用图像 tokenizerBEiT 需要一个图像 tokenizer把原始图像转换成离散 visual token。论文中使用的是基于 dVAE 的 tokenizer。这个 tokenizer 的作用类似 NLP 中的 BPE tokenizer只不过它处理的不是文本而是图像。它会将图像压缩成离散 token 序列并通过 decoder 尝试从这些 token 重建图像。在 BEiT 论文中作者直接使用了已有的公开 tokenizer而不是从头训练一个新的 tokenizer。该 tokenizer 将每张图像转换为 14×14 的 visual token 网格词表大小为 8192。2. 第二阶段训练 BEiT 预测 visual token有了 tokenizer 之后就可以训练 ViT 编码器。此时 tokenizer 的作用是提供监督信号。训练时原始图像被 tokenizer 转换成 visual token图像 patch 中一部分被 mask被 mask 后的 patch 序列输入 ViTViT 预测被 mask 位置对应的 visual token使用分类损失训练模型。这就是 BEiT 的 BERT-style 视觉预训练。注意这里虽然使用了 tokenizer 产生的 token 作为目标但整个过程仍然不需要人工标注。因此 BEiT 仍然属于自监督学习方法。8、BEiT 的实验结果BEiT 的实验主要验证两件事第一BEiT 预训练能否提升 ViT 在图像分类上的表现。第二BEiT 学到的表示能否迁移到密集预测任务比如语义分割。在 ImageNet-1K 分类任务上BEiT-B 在 224×224 分辨率下达到 83.2% top-1 accuracy提高到 384×384 后BEiT-B 达到 84.6%。BEiT-L 在 224×224 下达到 85.2%在 384×384 下达到 86.3%。论文指出BEiT-L 在只使用 ImageNet-1K 自监督预训练的情况下达到了非常有竞争力的结果。在 ADE20K 语义分割任务上BEiT 也表现出较好的迁移能力。论文中 BEiT 的 mIoU 为 45.6加入 intermediate fine-tuning 后提升到 47.7。OpenReview这些实验说明BEiT 学到的表示不只是对分类有用也能迁移到需要空间理解能力的任务中。BEiT 论文中还有一个很有意思的分析自监督预训练之后模型的 self-attention map 能够在一定程度上自动分离物体区域。也就是说即使 BEiT 没有使用人工标注它的注意力头也可能关注到前景目标、物体边界或语义区域。这说明 masked image modeling 不只是让模型学会“补图”而是在学习图像中更高层的结构关系。这点也为后续很多 ViT 自监督方法提供了启发Transformer 的全局注意力机制本身具有建模长距离关系的能力而合适的自监督任务可以激发这种能力。9 BEiT 与 MAE 的区别因为我们上一篇刚讲过 MAE所以这里有必要把 BEiT 和 MAE 放在一起比较。对比维度BEiTMAE预训练任务预测 visual token重建原始像素预测目标离散 token 分类连续像素回归是否需要 tokenizer需要 dVAE tokenizer不需要 tokenizermask 比例约 40%通常更高解码器设计主要使用预测头使用轻量 decoder 重建像素目标特点更偏语义 token更偏像素重建复杂度需要额外 tokenizer框架更简单BEiT 更像 BERT因为它预测的是“离散词表”中的 token。MAE 更像 autoencoder因为它直接重建被遮挡区域的像素。这也导致二者的思想侧重点不同BEiT 关注如何为图像构造类似文本词表的离散监督信号MAE 关注如何通过高比例遮挡和非对称编码器-解码器让 ViT 高效学习图像表示如果说 MAE 的优势是简单、直接、高效那么 BEiT 的优势就在于它更接近 NLP 中 BERT 的预训练范式。BEiT 的贡献总结BEiT 的贡献可以总结为三点。1. 提出了图像版 BERT 的预训练框架BEiT 将 BERT 中的 masked language modeling 迁移到视觉领域提出 masked image modeling。它证明了 ViT 不一定只能依靠监督训练或对比学习也可以通过遮挡预测进行自监督预训练。2. 引入离散 visual token 作为预测目标这是 BEiT 最关键的设计。图像没有天然词表所以 BEiT 使用 dVAE tokenizer 构造视觉词表。模型预测的不是像素而是视觉 token从而将图像重建任务转化为 token 分类任务。3. 验证了 MIM 对分类和分割任务的有效性BEiT 在 ImageNet 分类和 ADE20K 语义分割上都取得了较好的结果说明 masked image modeling 学到的表示具有较好的迁移能力。BEiT 的局限性当然BEiT 也不是完美的。首先BEiT 依赖额外的图像 tokenizer。这个 tokenizer 的质量会直接影响预训练目标的质量。如果 tokenizer 生成的 visual token 不够语义化模型学到的表示也可能受限。其次BEiT 的训练流程相对更复杂。相比 MAE 直接重建像素BEiT 需要先有一个 tokenizer再进行 masked token prediction。这使得整个方法的实现链条更长。第三BEiT 的 visual token 虽然比像素更抽象但它是否真正对应高级语义仍然取决于 tokenizer 本身。后续 BEiT v2 就进一步改进了视觉 tokenizer希望让视觉 token 更具语义信息。因此BEiT 的核心思想非常重要但它也留下了一个关键问题图像的“词表”应该如何构造才能真正像语言 token 一样承载语义这个问题也影响了后续很多 masked image modeling 方法的发展。如何理解 BEiT 在 ViT 发展中的位置如果我们把 ViT 的发展过程串起来看BEiT 的位置非常重要。ViT 证明了图像可以切成 patch然后用 Transformer 处理。DeiT 证明了在更合理的数据增强、蒸馏和训练策略下ViT 可以在 ImageNet 上更高效地训练。BEiT 进一步证明了ViT 可以像 BERT 一样通过 masked token prediction 进行自监督预训练。MAE 则从另一个方向证明了即使不构造视觉词表只要设计好遮挡比例和编码器-解码器结构像素重建也可以成为非常强大的自监督任务。所以 BEiT 的历史意义在于它明确提出了“图像版 BERT”的路线让视觉 Transformer 的自监督预训练进入了 masked image modeling 阶段。