[论文解读] Masked Image Modeling with Local Multi-Scale Reconstruction

Haoqing Wang, Yehui Tang+4 more|arXiv (Cornell University)|Mar 9, 2023

Advanced Neural Network Applications被引用 4

一句话总结

本文提出 LocalMIM，一种掩码图像建模方法，通过在视觉变换器的多个局部层上应用多尺度重建，加速自监督表示学习。通过让低层重建细粒度信号、高层重建粗粒度信号，并使用轻量级解码器，该方法显式地引导跨尺度的特征学习，在显著降低预训练计算量（例如，ViT-B 仅需 0.7 GPU 小时）的同时，保持下游任务的高精度，实现最先进性能。

ABSTRACT

Masked Image Modeling (MIM) achieves outstanding success in self-supervised representation learning. Unfortunately, MIM models typically have huge computational burden and slow learning process, which is an inevitable obstacle for their industrial applications. Although the lower layers play the key role in MIM, existing MIM models conduct reconstruction task only at the top layer of encoder. The lower layers are not explicitly guided and the interaction among their patches is only used for calculating new activations. Considering the reconstruction task requires non-trivial inter-patch interactions to reason target signals, we apply it to multiple local layers including lower and upper layers. Further, since the multiple layers expect to learn the information of different scales, we design local multi-scale reconstruction, where the lower and upper layers reconstruct fine-scale and coarse-scale supervision signals respectively. This design not only accelerates the representation learning process by explicitly guiding multiple layers, but also facilitates multi-scale semantical understanding to the input. Extensive experiments show that with significantly less pre-training burden, our model achieves comparable or better performance on classification, detection and segmentation tasks than existing MIM models.

研究动机与目标

• 为解决现有掩码图像建模（MIM）方法仅依赖顶层重建所导致的高计算成本和学习速度慢的问题。
• 通过显式地在编码器的低层和高层均引入多尺度重建，提升表示学习效率。
• 通过为低层分配细粒度监督、为高层分配粗粒度监督，实现多尺度语义理解。
• 在不引入额外组件（如码本或教师网络）的前提下，减轻预训练负担。
• 验证梯度隔离训练在深层网络中实现解耦、高效反向传播的有效性。

提出的方法

• 该方法在多个局部层上应用重建任务——具体而言，低层对来自小输入区域的细粒度监督进行重建，而高层则对来自大区域的粗粒度监督进行重建。
• 通过在不同尺度（例如，28²、14²、7²）上划分输入图像，并应用特征描述符（如像素值或 HOG 特征），提取多尺度监督信号。
• 使用包含 Transformer 块和转置卷积/池化层的轻量级解码器，将预测结果重采样至与监督信号相同的尺度，从而实现端到端训练。
• 采用非对称的编码器-解码器策略，仅将可见的 patch 输入编码器，从而在预训练过程中最小化计算成本。
• 应用梯度隔离训练，即在每个局部重建层后停止反向传播，实现网络不同部分的解耦训练。
• 该方法在柱状 ViT 和分层 Swin 变压器架构上进行了评估，并对层选择和监督尺度进行了消融研究。

实验结果

研究问题

• RQ1 在多个局部层而非仅顶层应用重建任务，是否能加速 MIM 中的自监督表示学习？
• RQ2 为低层和高层分别分配不同尺度的监督（细粒度 vs. 粗粒度），是否能提升多尺度语义理解能力与模型性能？
• RQ3 局部多尺度重建是否能在不损失精度的前提下，减少预训练时间和计算成本？
• RQ4 梯度隔离训练（即在每个局部层后停止梯度传播）是否能在保持性能的同时，实现高效且解耦的训练？
• RQ5 在 ViT 和 Swin 架构中，性能对局部层选择和监督尺度的选择有多敏感？

主要发现

• LocalMIM 仅使用 0.7 GPU 小时预训练 ViT-B，即在 ImageNet-1K 上达到 83.3% 的 top-1 准确率，相比先前的 MIM 模型显著降低了预训练成本。
• 在 1.1 GPU 小时的训练时间下，LocalMIM 在 Swin-B 上达到 83.8% 的准确率，优于或匹配现有方法，尽管训练时间更短。
• Swin 风格的层划分方式 [2,4,10,12] 与从细到粗的多尺度监督 [28²,14²,7²,7²] 在 ViT 上表现最佳。
• 梯度隔离训练在 ViT-B 上达到 83.0% 准确率，在 Swin-B 上达到 83.7%，几乎与全局反向传播相当，验证了局部监督的有效性。
• 多尺度监督优于单尺度监督，且从细到粗的监督策略比从粗到细更有效，与生物和架构先验一致。
• 可视化结果表明，LocalMIM 的自注意力机制在无需任务特定监督的情况下，能有效区分前景物体和语义一致的区域。