[论文解读] FastMIM: Expediting Masked Image Modeling Pre-training for Vision

Jianyuan Guo, Kai Han+4 more|arXiv (Cornell University)|Dec 13, 2022

Advanced Neural Network Applications被引用 4

一句话总结

FastMIM 通过使用低分辨率输入（例如 128×128）并重建方向梯度直方图（HOG）特征而非原始 RGB 像素，加速了掩码图像建模（MIM）的预训练过程。该方法将显存消耗和训练时间减少约 5 倍，同时在使用 ViT-B/Swin-B 主干网络时，在 ImageNet-1K 上实现了 83.8%/84.1% 的 top-1 准确率，达到或超过先前方法，且收敛速度显著更快。

ABSTRACT

The combination of transformers and masked image modeling (MIM) pre-training framework has shown great potential in various vision tasks. However, the pre-training computational budget is too heavy and withholds the MIM from becoming a practical training paradigm. This paper presents FastMIM, a simple and generic framework for expediting masked image modeling with the following two steps: (i) pre-training vision backbones with low-resolution input images; and (ii) reconstructing Histograms of Oriented Gradients (HOG) feature instead of original RGB values of the input images. In addition, we propose FastMIM-P to progressively enlarge the input resolution during pre-training stage to further enhance the transfer results of models with high capacity. We point out that: (i) a wide range of input resolutions in pre-training phase can lead to similar performances in fine-tuning phase and downstream tasks such as detection and segmentation; (ii) the shallow layers of encoder are more important during pre-training and discarding last several layers can speed up the training stage with no harm to fine-tuning performance; (iii) the decoder should match the size of selected network; and (iv) HOG is more stable than RGB values when resolution transfers;. Equipped with FastMIM, all kinds of vision backbones can be pre-trained in an efficient way. For example, we can achieve 83.8%/84.1% top-1 accuracy on ImageNet-1K with ViT-B/Swin-B as backbones. Compared to previous relevant approaches, we can achieve comparable or better top-1 accuracy while accelerate the training procedure by $\sim$5$ imes$. Code can be found in https://github.com/ggjy/FastMIM.pytorch.

研究动机与目标

• 为解决视觉变换器中掩码图像建模（MIM）预训练的高计算成本和内存负担问题。
• 实现在多种视觉主干网络（包括 Swin Transformer 等层次化架构）上的高效、通用的预训练。
• 在不损失下游性能的前提下，减少训练时间和内存使用。
• 研究输入分辨率、特征表示（RGB 与 HOG）以及网络深度对 MIM 效率和准确率的影响。

提出的方法

• 使用低分辨率输入图像（例如 128×128 而非 224×224）对视觉主干网络进行预训练，以降低计算负载和内存使用。
• 将像素重建替换为方向梯度直方图（HOG）特征重建，以在分辨率降低的情况下仍能保留纹理和边缘信息。
• 采用对称的编码器-解码器架构，在预训练过程中保留所有输入 token，从而与多种主干架构兼容。
• 将解码器尺寸与所选编码器的特征图尺寸相匹配，以保持训练稳定性和性能。
• 引入 FastMIM-P，一种渐进式分辨率调度策略，可在预训练过程中逐步提高输入分辨率，以增强模型容量和迁移性能。
• 利用 HOG 特征对几何变换和分辨率变化的不变性，使其在低分辨率设置下比原始像素目标更具鲁棒性。

实验结果

研究问题

• RQ1 在 MIM 预训练过程中降低输入分辨率是否能显著加速训练而不损害下游性能？
• RQ2 在使用低分辨率输入时，用 HOG 特征重建替代原始像素是否能缓解纹理信息的损失？
• RQ3 在低分辨率 MIM 中，预测目标的选择（RGB 与 HOG）如何影响训练稳定性和收敛速度？
• RQ4 在预训练过程中，能否优先训练编码器的浅层网络，而舍弃深层网络以进一步加快训练速度？
• RQ5 在预训练过程中采用渐进式输入分辨率调度，能否提升高容量模型的迁移性能？

主要发现

• FastMIM 相较于 MAE 和 SimMIM，将预训练时间减少了约 5 倍，在使用 ViT-B 和 Swin-B 主干网络时，分别在 ImageNet-1K 上实现了 83.8% 和 84.1% 的 top-1 准确率。
• 与原始 RGB 像素重建相比，使用 HOG 特征重建在低分辨率输入下可显著降低预训练损失并提升重建质量，减少歧义，增强训练稳定性。
• 使用 HOG 目标时，不同输入分辨率之间的性能差距显著减小，表明 HOG 对分辨率变化具有更强的鲁棒性。
• 在预训练过程中舍弃编码器的最后几层不会损害微调性能，从而可在几乎不损失准确率的前提下实现更快的训练速度。
• 编码器的浅层在预训练中更为关键，其表征质量对下游迁移性能有显著影响。
• FastMIM-P（在预训练过程中逐步提高输入分辨率）进一步提升了迁移性能，尤其对高容量模型（如 Swin-L）效果显著。