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

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 = \frac{HW}{P^2}$ 个patch： I = { x i p } i = 1 N \mathcal{I} = \{x_i^p\}_{i=1}^N I={xip​}i=1N​ ，$\mathcal{I}\in {\mathbb{R}}^{N\times (P^{2}C)}$ ，其中patch size为 $(P,P)$ 。然后把patchs $\mathcal{I}$ 带上位置信息，并线性投影成为tokens，输入到teacher model中。
• student model的输入：过程与teacher model的类似除了要对图像进行mask处理。要丢弃掉 $r\%$ 的patch,只喂进去可见的patch： V = { x i p } i = 1 N ( 1 − r % ) \mathcal{V} = \{x_i^p\}_{i=1}^{N(1-r\%)} V={xip​}i=1N(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模块的计算如下所示：
y ^ = { ∑ i = 0 S w i j A i ( x i ) } j = 0 T \hat{y} = \{\sum_{i=0}^Sw_{ij}A_i(x_i)\}_{j=0}^T y^​={i=0∑S​wij​Ai​(xi​)}j=0T​
这样就得到了student model**多级特征的一组线性组合 $\hat{y}$ **。预训练阶段，梯度可以反向传播到动态对齐矩阵 $W$ 。在下游任务中，会弃用整个DA模块。

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

损失函数

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

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

实验结果

• 实验参数见原文4.1

与SOTA对比

图像分类结果

目标检测和实例分割结果

消融实验

对齐策略

掩蔽策略

Teacher Model

适配器细节

预训练速度

自注意力权重可视化

我的个人博客地址：https://pixelpilot666.github.io ，欢迎访问。