[论文解读] Exploring The Role of Mean Teachers in Self-supervised Masked Auto-Encoders

Youngwan Lee, Jeffrey Willette+3 more|arXiv (Cornell University)|Oct 5, 2022

Domain Adaptation and Few-Shot Learning被引用 7

一句话总结

本文提出RC-MAE，一种自监督学习方法，通过集成指数移动平均（EMA）教师模型来增强掩码自编码器（MAE），该教师模型通过基于特征相似性的梯度调整，充当条件动量正则化器。与MAE及最先进自蒸馏方法相比，该方法实现了更快的收敛速度、更低的显存占用，以及在下游视觉任务中更高的鲁棒性和性能。

ABSTRACT

Masked image modeling (MIM) has become a popular strategy for self-supervised learning~(SSL) of visual representations with Vision Transformers. A representative MIM model, the masked auto-encoder (MAE), randomly masks a subset of image patches and reconstructs the masked patches given the unmasked patches. Concurrently, many recent works in self-supervised learning utilize the student/teacher paradigm which provides the student with an additional target based on the output of a teacher composed of an exponential moving average (EMA) of previous students. Although common, relatively little is known about the dynamics of the interaction between the student and teacher. Through analysis on a simple linear model, we find that the teacher conditionally removes previous gradient directions based on feature similarities which effectively acts as a conditional momentum regularizer. From this analysis, we present a simple SSL method, the Reconstruction-Consistent Masked Auto-Encoder (RC-MAE) by adding an EMA teacher to MAE. We find that RC-MAE converges faster and requires less memory usage than state-of-the-art self-distillation methods during pre-training, which may provide a way to enhance the practicality of prohibitively expensive self-supervised learning of Vision Transformer models. Additionally, we show that RC-MAE achieves more robustness and better performance compared to MAE on downstream tasks such as ImageNet-1K classification, object detection, and instance segmentation.

研究动机与目标

• 理解基于EMA蒸馏的自监督学习中学生与教师之间的动态交互机制。
• 探究均值教师模型在视觉Transformer的像素级掩码自编码（MIM）中的功能角色。
• 开发一种更高效、更有效的自监督预训练方法，相比现有自蒸馏方法显著降低计算与显存开销。
• 提升在图像分类以及密集预测任务（如目标检测与实例分割）上的下游性能。

提出的方法

• 提出RC-MAE，一种改进的MAE，通过引入EMA教师模型为学生模型的重建补丁提供一致性目标。
• 使用重建损失 $\mathcal{L}_{\text{r}}$ 比较学生输出与原始掩码补丁。
• 应用一致性损失 $\mathcal{L}_{\text{c}}$ 将学生预测与教师在可见补丁上的输出对齐。
• 采用EMA更新规则维护一个缓慢更新的教师网络，其动量基于当前输入与历史输入之间的特征相似性进行调整。
• 通过线性模型分析梯度动态，表明当特征相似时，教师通过移除先前方向来有条件地校正梯度。
• 在ViT-B和ViT-L主干网络上，基于ImageNet-1K、COCO及鲁棒性基准对方法进行实证验证。

实验结果

研究问题

• RQ1 EMA教师如何影响像素级掩码自编码中的梯度更新？
• RQ2 除了提供目标外，教师在学生优化过程中的功能角色是什么？
• RQ3 将EMA教师简单集成到MAE中，是否能相比现有自蒸馏方法实现更快收敛与更低显存占用？
• RQ4 教师的条件梯度校正是否能提升视觉任务的鲁棒性与下游性能？
• RQ5 教师的作用能否在理论上解释为一种条件动量正则化形式？

主要发现

• RC-MAE在800个周期后于ImageNet-1K分类任务上达到83.4%的top-1准确率，与MAE在1600个周期时的性能相当，表明收敛速度显著提升。
• RC-MAE相比最先进的自蒸馏方法（如iBOT、MSN和BootMAE）显著降低显存占用与计算成本，仅需将可见补丁输入教师网络。
• 在ImageNet-C、ImageNet-A、ImageNet-R和ImageNet-Sketch上，RC-MAE展现出一致的鲁棒性提升，mCE降低1.8个百分点（49.9 vs. 51.7），在IN-A和IN-R上准确率最高提升1.2个百分点。
• RC-MAE的注意力图显示，其对查询补丁附近的图像相关区域具有更清晰的关注，而iBOT和MSN则关注无关区域。
• 理论分析表明，教师模型充当条件动量正则化器，根据当前与历史特征表示之间的相似性调整梯度。
• 在线性模型与深层ViT架构中的实证验证均确认了条件梯度校正机制，支持所提出的理论分析。