Stare at What You See: Masked Image Modeling without Reconstruction 阅读笔记

论文地址：2211.08887.pdf (arxiv.org)
该论文收录于CVPR2023

亮点

在Masked Image Modeling(MIM)中提出MaskAlign使得不用在masked token上进行任何重建就能实现SOTA效果，同时不用重建意味着减少了大量的计算，使得效率得到了很大的提高。
提出动态对齐模块Dynamic Align (DA) module，在增加很少的计算量的条件下，解决student model和teacher model之间输入不一致的问题。

简介

Masked Image Modeling(MIM)中最近的一些工作，引入了语义丰富的Teacher Model（如CLIP）。通过利用Teacher Model提取的特征，作为Masked Image特征重建的target。在masked regions中的重建，暗示着迫使图像编码器学习图像中的语义相关性。重建的过程会带来大量的计算，增大了预训练的开销。而且强大的teacher model提取出的特征包含了丰富的语义相关性。所以作者提出了叫做MaskAlign的MIM范式，就是不用重建，只是将由student model提取的特征与teacher model提取的特征对齐。

因此，MaskAlign使得student model在特征对齐中学到teacher model对特征的良好表示，还能在掩码建模（masked modeling）的过程中学到完整图像的信息——因为在完整图像和掩码视图（masked view）之间存在特征的一致性，使得student model需要在比teacher model少得多的信息中推断整体的语义。

然而，输入到student model中的信息比到teacher model中的信息少的多，这就会导致每个层的特征无法对齐。作者提出了动态对齐模块（Dynamic Align (DA) module）来增强学生特征（student's model）。DA模块动态的聚合不同级别学生的特征，并于teacher model保持一致。这种方法还能转移到不对称的师生模型结构中。

详细方法

模型输入

MaskAlign中包含一个随机初始化的student model和一个经过预训练的语义丰富的frozen teacher model去提取特征作为监督。作者采用的是标准ViT作为student model，ViT （如 CLIP-ViT, DINO）作为teacher model。

teacher model的输入：一幅图像 $I \in \mathbb{R}^{C \times H \times W}$ 被分成 $ N = $ 个patch：$\mathcal{I} = \{x_i^p\}_{i=1}^N$ ，$\mathcal{I} \in \mathbb{R}^{N\times(P^2C)}$ ，其中patch size为 $(P, P)$ 。然后把patchs $\mathcal{I}$ 带上位置信息，并线性投影成为tokens，输入到teacher model中。
student model的输入：过程与teacher model的类似除了要对图像进行mask处理。要丢弃掉 $r\%$ 的patch,只喂进去可见的patch：$\mathcal{V} = \{x_i^p\}_{i=1}^{N(1-r\%)}$ 到student model中。

掩蔽策略

掩蔽策略（Masking Strategy）也就是从完整图像patch $\mathcal{I}$ 中形成mask view $\mathcal{V}$ 的策略，主要分为两种：

随机掩蔽：将部分patch以均匀分布的形式在 $\mathcal{I}$ 进行采样形成 $\mathcal{V}$。
注意掩蔽（Attentive masking）：在teacher model的引导下，形成将那些大概率是重要的区域遮盖住的图像的$\mathcal{V}$ 。

这两中策略的对比，在后面实验部分有提到。

动态对齐

因为student model和teacher model的输入不一致（因为student model只输入了 $N(1-r\%)$ 个token而teacher model输入了 $N$ 个token），所以student model和teacher model的每一层输出都会出现对不齐的情况。

为解决上述问题，作者提出了动态对齐模块（Dynamic Alignment （DA）module）。Transformer由一系列的block组成，在student model的每个块后面加上了一个适配器 $A_i$ 将每个块的输出 $x_i$ 投影到teacher model的特征空间。这个适配器是一个轻型模型如2层MLP。

作者为了动态地聚合不同级的学生特征，并与多级教师特征对齐，使用了一个动态对齐矩阵—— $W$ , $W$ 是一个 $S \times T$ 的矩阵，带有项 $w_{ij}$ 。其中 $S,\ T$ 分别是student model中的block的数量和teacher model中的block的数量。DA模块的计算如下所示： \[ \hat{y} = \{\sum_{i=0}^Sw_{ij}A_i(x_i)\}_{j=0}^T \] 这样就得到了student model多级特征的一组线性组合 $\hat{y}$ 。预训练阶段，梯度可以反向传播到动态对齐矩阵 $W$ 。在下游任务中，会弃用整个DA模块。

作者还对比了没有动态对齐矩阵的情况叫做逐层对齐（layer-wise alignment）。在实验中显示这有没有动态对齐矩阵，对计算效率几乎没有影响。

损失函数

为了抑制教师特征的数量级，在生成对齐target $\tilde{y}$ 的时候，先对第 $i$ 层的特征 $y_i$ 归一化处理，得到 $\tilde{y}$ 的过程如下： \[ \tilde{y} = \{ \mathrm{Normalize}(y_i)\}_{j=0}^T \]

最后在学生特征和教师特征之间用了平滑L1 loss： \[ \mathcal{L}_{Align}(\hat{y}, \tilde{y}) = \begin{cases} \frac{1}{2}(\hat{y} - \tilde{y})^2, \quad |\hat{y} - \tilde{y}| \le 1\\ (|\hat{y} - \tilde{y}| - \frac{1}{2}), \quad otherwise \end{cases} \]