Visualizing the loss landscape of Self-supervised Vision Transformer

📄 arXiv: 2405.18042v1 📥 PDF

作者: Youngwan Lee, Jeffrey Ryan Willette, Jonghee Kim, Sung Ju Hwang

分类: cs.CV, cs.LG

发布日期: 2024-05-28

备注: NeurIPS 2023 Workshop: Self-Supervised Learning - Theory and Practice

💡 一句话要点

通过可视化损失 landscape 分析自监督 ViT 泛化能力

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 自监督学习 视觉Transformer 损失 landscape 掩码图像建模 泛化能力 优化 EMA教师模型

📋 核心要点

  1. 现有研究缺乏对MAE自监督ViT优于监督ViT的泛化能力原因的深入理解。
  2. 该论文通过可视化损失 landscape，从优化角度分析自监督ViT的泛化能力。
  3. 实验发现MAE-ViT损失 landscape 更平滑，RC-MAE通过EMA教师扩大凸性区域，加速收敛。

📝 摘要（中文）

掩码自编码器(MAE)作为一种具有代表性的自监督方法，在视觉Transformer的掩码图像建模中备受关注。尽管MAE显示出比从头开始的完全监督训练更好的泛化能力，但其原因尚未得到充分探索。重建一致性掩码自编码器(RC-MAE)在MAE中采用了一种指数移动平均(EMA)教师模型的自蒸馏方案，研究表明EMA教师在优化过程中执行条件梯度校正。为了进一步研究MAE训练的自监督ViT(MAE-ViT)具有更好泛化能力的原因，以及RC-MAE梯度校正的效果，我们通过可视化MAE和RC-MAE训练的自监督视觉Transformer的损失 landscape，并将其与监督ViT(Sup-ViT)进行比较。与之前基于分类任务损失的神经网络损失 landscape 可视化不同，我们通过计算预训练任务损失来可视化ViT的损失 landscape。通过损失 landscape，我们发现了两个有趣的观察结果：(1)MAE-ViT比Sup-ViT具有更平滑和更宽的整体损失曲率。(2)EMA教师允许MAE在预训练和线性探测中扩大凸性区域，从而加快收敛速度。据我们所知，这项工作是第一个通过损失 landscape 来研究自监督ViT的工作。

🔬 方法详解

问题定义：论文旨在解决自监督视觉Transformer（ViT）在掩码图像建模（MIM）任务中，为何能够获得比监督学习更好的泛化能力的问题。现有方法缺乏对这种泛化能力背后优化过程的深入理解，特别是对于MAE及其变体RC-MAE，其优越性能的原因尚不明确。

核心思路：论文的核心思路是通过可视化损失 landscape 来理解自监督ViT的优化过程。通过比较自监督（MAE和RC-MAE）和监督ViT的损失 landscape，分析其平滑度、宽度和凸性等特征，从而揭示自监督学习泛化能力更强的原因。论文假设更平滑、更宽的损失 landscape 对应于更好的泛化能力。

技术框架：论文的技术框架主要包括以下几个步骤：1) 使用MAE和RC-MAE对ViT进行预训练；2) 使用监督学习对ViT进行训练；3) 计算预训练和监督学习的损失；4) 可视化不同训练方式下的损失 landscape；5) 分析损失 landscape 的特征，如平滑度、宽度和凸性。其中，损失 landscape 的可视化是通过在模型参数空间中选取两个随机方向，并计算沿着这些方向的损失值来实现的。

关键创新：论文的关键创新在于首次将损失 landscape 可视化技术应用于自监督视觉Transformer的分析。与以往基于分类任务损失的 landscape 可视化不同，该论文基于预训练任务的损失进行可视化，从而更直接地反映了自监督学习的优化过程。此外，论文还揭示了EMA教师在RC-MAE中对损失 landscape 的影响，即扩大凸性区域，加速收敛。

关键设计：论文的关键设计包括：1) 使用预训练任务的损失（例如，图像重建损失）来计算损失 landscape；2) 使用随机方向来探索参数空间，并计算沿着这些方向的损失值；3) 使用EMA教师模型来指导RC-MAE的训练；4) 通过比较不同训练方式下的损失 landscape 的平滑度、宽度和凸性等特征，来分析自监督学习的泛化能力。

📊 实验亮点

实验结果表明，MAE-ViT的损失 landscape 比 Sup-ViT 更平滑和更宽，这解释了MAE-ViT 更好的泛化能力。此外，RC-MAE中的EMA教师能够扩大损失 landscape 的凸性区域，从而加快收敛速度。这些发现为理解自监督学习的优化过程提供了新的视角。

🎯 应用场景

该研究成果可应用于改进自监督视觉Transformer的训练策略，例如，通过设计更有效的梯度校正方法或损失函数，来获得更平滑、更宽的损失 landscape，从而提高模型的泛化能力。此外，该研究方法也可推广到其他自监督学习算法的分析和优化中，例如，对比学习等。

📄 摘要（原文）

The Masked autoencoder (MAE) has drawn attention as a representative self-supervised approach for masked image modeling with vision transformers. However, even though MAE shows better generalization capability than fully supervised training from scratch, the reason why has not been explored. In another line of work, the Reconstruction Consistent Masked Auto Encoder (RC-MAE), has been proposed which adopts a self-distillation scheme in the form of an exponential moving average (EMA) teacher into MAE, and it has been shown that the EMA-teacher performs a conditional gradient correction during optimization. To further investigate the reason for better generalization of the self-supervised ViT when trained by MAE (MAE-ViT) and the effect of the gradient correction of RC-MAE from the perspective of optimization, we visualize the loss landscapes of the self-supervised vision transformer by both MAE and RC-MAE and compare them with the supervised ViT (Sup-ViT). Unlike previous loss landscape visualizations of neural networks based on classification task loss, we visualize the loss landscape of ViT by computing pre-training task loss. Through the lens of loss landscapes, we find two interesting observations: (1) MAE-ViT has a smoother and wider overall loss curvature than Sup-ViT. (2) The EMA-teacher allows MAE to widen the region of convexity in both pretraining and linear probing, leading to quicker convergence. To the best of our knowledge, this work is the first to investigate the self-supervised ViT through the lens of the loss landscape.