链接：https://arxiv.org/pdf/2208.06366
论文：BEIT V2: Masked Image Modeling with Vector-Quantized Visual Tokenizers

Introduction

1. Motivation：Masked image modeling (MIM) 任务在自监督表征学习上取得了不错的成绩，但是现有方法大多是基于low-level image pixels，需要探索模型的high-level semantics。
2. 创新点：引入了矢量量化知识蒸馏（Vector-Quantized Knowledge Distillation ，VQ-KD）算法来离散语义空间，同时引入一个patch聚合策略（a patch aggregation strategy）鼓励模型关联所有patch到[CLS]
3. 实验效果：
   

Details

整体架构与BEiT相似，依旧是包含一个visual tokenizer对图像进行离散表示（visual tokens），训练目标则是重建相应位置的masked visual tokens。

1. Image Representation：依旧使用ViT，将输入切为patches，flattened and linearly projected to input embeddings，最终输入transformer。
2. Training Visual Tokenizer
   
   a. 提出vector-quantized knowledge distillation (VQ-KD)方法训练visual tokenizer， 架构如上图所示，包括tokenizer和decoder两部分。
   b. tokenizer将输入图像映射为一系列visual tokens，进行离散化，与patch数量对应。其包含一个Vision Transformer encoder, and a quantizer。想通过encoder将图像编码为向量，然后quantizer查找最近邻的表示。寻找最近邻embedding公式如下（quantizer的目的是将向量映射到固定词表，便于在后续mask任务中预测）：
   
   先对向量进行l2正则，使用余弦相似度计算。
   c. 向量量化后，过l2正则输入decoder，decoder模型为多层Transformer，其目标是重建Teacher模型的语义特征（Teacher模型可以为DINO或者CLIP）。再最大化decoder的输出和teacher模型的输出的余弦相似度。
   d. 由于量化过程不可微，直接使用梯度拷贝（从decoder的输入到encoder的输出），直观上，量化器为encoder输出查找最近的编码，所以该codebook embeddings的提督对encoder的优化方向有效。
   
   第一项为decoder输出与teacher输出的余弦相似度损失，sg表示stop-gradient，前向传递过程中为恒等式，同时在反向传播期间具有零梯度。第二项和第三项分别代表前向和反向。
3. Improving codebook utilization.
   向量量化训练期间很容易遇到codebook的坍缩，只有一小部分的codes可以使用，
   a. 量化过程将码本嵌入空间的维数减少到32-d, 在被送入解码器之前被映射回高维空间
   b. 指数移动平均可以使VQ-KD的训练稳定。
4. Pretraining BEIT V2
   a. 给定输入图像x，大约40%的patches会被block-wisely masked，masked position会被标记，masked patch会被一个shared learnable embedding取代。
   b. prepend a learnable [CLS] token，在预测时对应位置加一个全连接层。
   
   c. Pretraining global representation:为了使CLS更好地进行全局图像特征表示，消除patch-level的预训练对image-level的表示差异，选取最后L层的CLS表示，以及encoder第l层的patch表示，拼接作为一个浅层（2-layer）Transformer的输入，进行掩码预测。MIM Head的参数共享，两个loss相加（原始MIM loss和过完浅层transformer的loss）。直观上，这样做的好处在于由于使用训练不充分的中间表示，导致CLS更好地涵盖全局信息，使MIM loss更低。该新增的两层transformer仅辅助训练，inference会被丢弃。

实验

1. 两阶段训练参数
   a. Visual tokenizer training
   ViT-B/16，decoder为三层Transformer，和encoder头数与维度都相同，Teacher使用CLIP-B/16 train VQ-KD on ImageNet-1k with 224×224 resolution。code size K is set as 8192，code dimension D as 32。
   b. Masked image modeling
   ImageNet-1K，set l = 9 for ViT-B/16, l = 21 for ViT-L/16，40% mask
   pre-train的vit encoder和modeling阶段不是同一个
2. 对比实验
   
   
3. 消融实验
   a. VQ-KD的消融
   
   b. Patch aggregation消融
   
   c. VQ-KD targets
   
4. 可视化实验