[论文解读] BEiT v2: Masked Image Modeling with Vector-Quantized Visual Tokenizers

Zhiliang Peng, Dong Li+3 more|arXiv (Cornell University)|Aug 12, 2022

Advanced Neural Network Applications被引用 114

一句话总结

BEiT v2 引入通过知识蒸馏训练的向量量化视觉标记器，为掩码图像建模提供语义目标，以及一种补丁聚合策略以增强全局表示；在 ImageNet 分类和 ADE20K 分割上达到最先进的结果。

ABSTRACT

Masked image modeling (MIM) has demonstrated impressive results in self-supervised representation learning by recovering corrupted image patches. However, most existing studies operate on low-level image pixels, which hinders the exploitation of high-level semantics for representation models. In this work, we propose to use a semantic-rich visual tokenizer as the reconstruction target for masked prediction, providing a systematic way to promote MIM from pixel-level to semantic-level. Specifically, we propose vector-quantized knowledge distillation to train the tokenizer, which discretizes a continuous semantic space to compact codes. We then pretrain vision Transformers by predicting the original visual tokens for the masked image patches. Furthermore, we introduce a patch aggregation strategy which associates discrete image patches to enhance global semantic representation. Experiments on image classification and semantic segmentation show that BEiT v2 outperforms all compared MIM methods. On ImageNet-1K (224 size), the base-size BEiT v2 achieves 85.5% top-1 accuracy for fine-tuning and 80.1% top-1 accuracy for linear probing. The large-size BEiT v2 obtains 87.3% top-1 accuracy for ImageNet-1K (224 size) fine-tuning, and 56.7% mIoU on ADE20K for semantic segmentation. The code and pretrained models are available at https://aka.ms/beitv2.

研究动机与目标

• 通过将语义视觉标记器用于从像素级到语义级的掩码图像建模推广。
• 训练一个 VQ-KD 标记器，用可学习的码本对语义空间进行离散化。
• 用补丁聚合策略提升 BEiT 预训练中的全局表示能力。
• 在 ImageNet 分类（微调和线性探针）和 ADE20K 语义分割上经验性验证 BEiT v2。
• 展示在模型规模和预训练时程上的鲁棒性与可迁移性。

提出的方法

• 在 VQ-KD 下训练一个视觉标记器，其中 ViT 编码器将图像映射到补丁特征，码本选择最近的码，解码器使用余弦相似度作为监督来学习还原教师特征。
• 将 OpenAI CLIP-B/16 作为教师提供语义引导用于 VQ-KD，并通过停止梯度来通过量化器对编码器进行反向传播。
• 通过 BEiT v2 的预训练，按照 40% 的掩码比例对被掩盖的补丁预测离散视觉标记，遵循补丁级标记预测。
• 引入补丁聚合机制，鼓励 [CLS] 标记聚合补丁信息，使局部补丁表示与全局图像表示对齐。
• 预训练目标函数将标准 MIM 损失（重建视觉标记）与聚焦 CLS 的损失相结合，以提升全局表示质量。
• 可选地在预训练期间包含一个浅层 Transformer 解码器，以加强基于 CLS 的全局表示，同时不增加推理成本。

实验结果

研究问题

• RQ1 语义、向量量化的视觉标记器是否能在掩码图像建模中优于像素级目标？
• RQ2 补丁聚合策略是否在 MIM 预训练中促进更好的全局图像表示？
• RQ3 BEiT v2 在不同模型规模和预训练时程下的 ImageNet 微调、线性探针以及 ADE20K 语义分割表现如何？
• RQ4 VQ-KD 目标选择（DINO/CLIP）对下游性能有何影响？
• RQ5 BEiT v2 对不同下游评估设置和鲁棒性基准的鲁棒性如何？

主要发现

• 在 ViT-B/16 的 300 轮训练下，BEiT v2 在 ImageNet-1K 的微调 top-1 为 85.0，ADE20k mIoU 为 52.7。
• 在 ViT-B/16 的 1600 轮训练下，BEiT v2 的 top-1（ImageNet）为 85.5，ADE20k mIoU 为 53.1；若中间在 ImageNet-21k 上微调，精度分别提高到 86.5 和 53.5。
• 在 ViT-L/16 的 300 轮训练下，BEiT v2 达到 86.6% 的 top-1（ImageNet）和 55.0% 的 ADE20k mIoU；1600 轮训练得到 87.3% 的 top-1 与 56.7% 的 ADE20k mIoU。
• 在 ImageNet-1K 的线性探针中，BEiT v2 达到 80.1% 的 top-1，优于在报告的比较中 BEiT、CAE、MVP 和 MoCo v3 的结果。
• BEiT v2 在若干 ImageNet 变体（对抗、改写、草图）上显示出较强的鲁棒性，相较于 MAE，在多个设定中取得显著提升。
• 消融研究表明，更深的 VQ-KD 解码器会降低码本使用率和下游性能，而更小的码本则提升使用率和准确性。