华为诺亚方舟实验室提出LocalMIM，局部重构自监督预训练3倍加速MAE

不同于已有全部MIM模型仅在编码器顶层施加重构损失，我们首次将重构任务引入多个选择的局部层，并提出多尺度重构：较低层重构细尺度信息，较高层重构粗尺度信息。就ImageNet-1K上Top-1微调准确率而言，我们的LocalMIM分别以3.1倍、5.6倍、3.6倍和6.4倍的加速达到MAE，MaskFeat，SimMIM192和GreenMIM的最佳表现。我们的方法是架构无关的，未来可以应用到更多新的主干网络中。

论文见：

https://arxiv.org/pdf/2303.05251v1.pdf

MindSpore代码：

https://gitee.com/mindspore/hub/blob/master/mshub_res/assets/noah-cvlab/gpu/1.8/localmim_v1.0_imagenet2012.md

引言

近年来，从MoCo，SimCLR等对比式方法到MAE, BEiT等掩码图像建模（MIM）方法，视觉自监督表征学习的研究如火如荼。随着vision transformer的发展，MIM方法凭借其在下游任务中更为优越的微调性能而成为关注的焦点。现有的MIM方法通常训练成本很高，而实践中我们期望其能够从海量无标记数据（如网络上随机爬取的图像）中学习通用的知识，所以高昂的预训练成本限制了其工业落地。

MIM方法的计算量在于编码器和解码器，由于解码器可以很小，已有的加速预训练的方法都通过降低编码器的计算量来加速编码过程，可分为：

• 编码器只处理可见patch，如MAE和GreenMIM；

• 降低输入图像分辨率来减少总的patch，如LoMaR，UM-MAE和FastMIM。

不同于已有的思路，我们转换视角，从表征学习过程本身来深入思考现有方法的不足。有理论[1]表明当我们对输入图像做片段划分并线性映射之后，由于映射过程的随机性，所获得的patch表征已经基本丢失patch之间的语义关联。vision transformer中的注意力机制通过后续patch之间的交互来学习出这种语义关联并获得比像素空间更优的表征空间。值得注意的是，处理patch之间关联的计算复杂度是O(N^2)的，所以模型想要较好地学到这种语义关联并非易事。现有全部MIM方法只在顶层引入重构任务，使得较低层无法获得直接的指导，从而只能通过缓慢的学习过程来学习patch表征及语义关联，拖累了整体的表征学习过程。尤其对于一些金字塔型主干网络，其较低层往往有着远多于顶层的patch（如Swin-224中3136（最底层）vs 49（顶层））。另一方面，实际上较低层在表征学习中扮演关键角色：

1）良好学习的较低层可以将知识传递给较高层以促进其学习；2）在下游任务微调时，较高层通常快速适应到新任务中，而较低层变化较慢，需要在预训练时就得到充分学习（无论是NLP还是CV下游任务中学习率layer decay通常效果最佳）。为了更直观地展现模型不同层对patch间语义关联的学习程度，我们考察了不同层的query patch与key patch之间的标准化互信息（Normalized Mutual Information，NMI），如图3所示。

较高的NMI值意味着注意力强烈地依赖于query patch，如图3所示已有的很多经典模型（BeiT，SimMIM和MaskFeat）在较低层的注意力并不像顶层一样强烈地依赖于query patch。

方法

LocalMIM框架图

通过上面的分析，我们发现在掩码图像建模中较低层的学习很关键，然而目前所有MIM方法都只显式地指导顶层的学习。为此，考虑到重构任务需要patch之间的语义推理才能完成，我们将重构任务引入多个局部层以显式地进行有意义的指导。实际上，特征蒸馏也可以显式指导多个局部层，但是特征蒸馏需要提前训练好的或动量更新的teacher网络，会显著增加计算量，而且简单的特征匹配比不上推理任务对学习语义关联带来的增益。进一步地，我们发现直接地将顶层的重构任务引入到多个局部层增益不明显，原因可能是多个不同的局部层需要学习不同粒度的信息。为此，我们考虑从原始输入中提取不同尺度的监督信号来指导多个局部层的学习。具体的，对于原始输入，为了获得监督信号，已有方法通常首先将x划分为不重叠的区域，该划分与构造编码器输入的划分对齐。然后使用恰当的特征描述算子（如像素标准化，HOG或预训练的codebook）提取每个区域的特征作为监督信号。我们定义此时监督信号的尺度为，通常可以认为在粗糙划分（p较大）下每个区域捕捉原始输入相对high-level的语义信息，比如目标的部分或整体形状，而精细划分（p较小）下每个区域捕捉相对low-level的语义信息，比如边、角或纹理。

另一方面，视觉主干网络设计的很多工作已经表明从底层到高层特征图的尺度由粗到细对各种视觉任务（检测，分割和分类）都有增益，通常较低层学习细粒度的信息而较高层学习粗粒度的信息。为此，我们不妨从原始输入中构造多尺度的监督信号（选择不同的p）来分别用于多个局部层的重构，令较低层重构细尺度的监督信号而较高层重构粗尺度的监督信号。值得注意的是，我们的方法与“加速编码”方法兼容，特别的我们采用MAE和GreenMIM中的做法，即编码器只输入可见patch。算法的整体流程如图2(a)所示。图2(b)中演示了某一特定尺度下的解码过程，解码器由三部分组成：推理部分（transformer blocks）+缩放部分（Deconvolution/Pool）+预测部分（MLP）。推理部分负责基于可见patch的表征推理被遮挡patch的信息；缩放部分是处理特征尺度与监督信号尺度不一致的情况，比如ViT这种柱状结构每层特征尺度不变而监督信号尺度是变化的，当不匹配时需要使用反卷积或池化操作进行上/下采样；预测部分负责整合放缩后的预测来作为最终输出。对于金字塔结构的模型，通常已经划分为多个stage，我们将重构任务用于每个stage的末端，对于柱状架构，我们参照金字塔架构的经验，选择部分层进行重构。

总的来说，局部层多尺度重构（LocalMIM[2]）既能显式地指导较低层以加速整体表征学习，又能促进对输入图像的多尺度理解。而且，该方法是架构无关的，理论上与各种主干网络适配，可用于未来更先进的主干网络中。

实验结果

我们目前在柱状架构ViT以及金字塔架构Swin上验证LocalMIM的有效性，考察其在分类、检测以及分割中的表现，如表1，2和3所示。出于其简单性，我们只考察像素标准化和HOG两种特征描述算子。

表1.在ImageNet-1K上的Top-1微调准确率

表 2. ADE20K上语义分割结果

表3. COCO上目标检测和实例分割结果

从表1中可以看到，LocalMIM比已有模型显著更加高效。具体的，就ImageNet-1K上top-1微调准确率而言，LocalMIM分别以3.1倍和5.6倍的加速达到MAE和MaskFeat的最佳表现，以3.6倍和6.4倍的加速达到SimMIM192和GreenMIM的最佳表现。相较于其他模型，LocalMIM也以显著更少的预训练时长实现了可比较的表现。从表2和3中可以看到，在下游的检测及分割任务上，LocalMIM以显著更少的计算成本获得更优的表现。更多的消融实验及实施细节可见论文。

除此之外，我们还做了一个有趣的实验，在训练过程中对选择的层进行梯度截断，即每阶段的参数只能接收来自该阶段重构任务的反传梯度，接收不到来自更高层的梯度，其结果如表4所示。

表4.LocalMIM进行梯度截断的训练可以达到与全局反向传播相同的表现

我们发现，即便没有全局的反向传播梯度只使用局部的监督梯度也能很好的指导主干网络每层的表征学习，这一方面展现了我们引入的局部监督任务的优越性，另一方面也展现了神经网络解耦训练的可能性。神经网络解耦训练能够训练任意深的网络而无须担心显存不够，而且也可以有效缓解梯度消失和爆炸问题。目前视觉领域常用的主干网络通常较浅，但未来也许会需要预训练非常深层的主干网络，可以作为进一步的探索。

参考文献

[1] How to Understand Masked Autoencoders. Cao and Xu, 2022.

[2] Masked Image Modeling with Local Multi-Scale Reconstruction. Wang et.al. CVPR 2023.