Asymmetric Masked Distillation for Pre-Training Small Foundation Models

📄 arXiv: 2311.03149v2 📥 PDF

作者: Zhiyu Zhao, Bingkun Huang, Sen Xing, Gangshan Wu, Yu Qiao, Limin Wang

分类: cs.CV

发布日期: 2023-11-06 (更新: 2024-04-01)

备注: Accepted by CVPR 2024

🔗 代码/项目: GITHUB

💡 一句话要点

提出不对称掩码蒸馏以高效预训练小型基础模型

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 不对称掩码蒸馏 小型基础模型 视觉变换器 自监督学习 模型压缩 特征对齐 计算机视觉 预训练

📋 核心要点

  1. 现有的大型基础模型在计算性能上存在高昂的计算成本，限制了其在资源受限环境中的应用。
  2. 本文提出的不对称掩码蒸馏（AMD）框架，通过不对称掩码策略优化小型模型的预训练过程，提高了模型的上下文理解能力。
  3. 实验表明，AMD在多个数据集上均取得了显著的性能提升，尤其是在分类任务中表现出色，验证了其有效性和适用性。

📝 摘要（中文）

自监督基础模型在计算机视觉领域展现出巨大潜力，尤其是在掩码自编码的预训练范式下。然而，现有的大型基础模型往往导致高计算成本。本文聚焦于预训练相对较小的视觉变换器模型，以便高效适应下游任务。我们提出了一种新的不对称掩码蒸馏（AMD）框架，通过不对称掩码策略，使教师模型能够以较低的掩码比例获取更多上下文信息，而学生模型则保持较高的掩码比例。我们设计了定制的多层特征对齐机制，以规范学生模型的预训练。实验结果表明，AMD在IN1K数据集上使用ViT-B模型达到了84.6%的分类准确率，并在Something-in-Something V2数据集上实现了73.3%的分类准确率，相较于原始VideoMAE模型提升了3.7%。

🔬 方法详解

问题定义：本文旨在解决大型基础模型在计算成本和资源消耗上的高昂问题，限制了其在实际应用中的广泛使用。现有方法在小型模型的预训练上缺乏有效的策略，导致性能不足。

核心思路：提出的不对称掩码蒸馏（AMD）框架，通过设计不对称的掩码策略，使教师模型能够获取更多上下文信息，同时保持学生模型的高掩码比例，从而提升小型模型的预训练效果。

技术框架：AMD框架包括教师模型和学生模型的构建，教师模型使用较低的掩码比例以获取更多信息，而学生模型则使用较高的掩码比例。通过多层特征对齐机制，确保教师和学生模型之间的有效信息传递。

关键创新：最重要的创新在于不对称掩码策略的引入，使得教师模型和学生模型在信息获取上存在差异化设计，从而提升了小型模型的学习能力和性能。

关键设计：在模型设计中，采用了定制的多层特征对齐机制，确保教师和学生模型在不同层次上的特征能够有效对齐。此外，损失函数的设计也针对性地优化了学生模型的学习过程。

📊 实验亮点

在IN1K数据集上，AMD框架使用ViT-B模型达到了84.6%的分类准确率；在Something-in-Something V2数据集上，分类准确率为73.3%，相比于原始VideoMAE模型提升了3.7%。这些结果表明，AMD在小型模型的预训练中具有显著的性能提升，验证了其有效性。

🎯 应用场景

该研究的潜在应用领域包括计算机视觉中的图像分类、视频理解等任务，尤其适用于资源受限的设备和场景。通过高效的预训练策略，AMD框架能够为小型模型提供更强的性能支持，推动其在实际应用中的广泛使用，未来可能在智能设备和边缘计算等领域产生重要影响。

📄 摘要（原文）

Self-supervised foundation models have shown great potential in computer vision thanks to the pre-training paradigm of masked autoencoding. Scale is a primary factor influencing the performance of these foundation models. However, these large foundation models often result in high computational cost. This paper focuses on pre-training relatively small vision transformer models that could be efficiently adapted to downstream tasks. Specifically, taking inspiration from knowledge distillation in model compression, we propose a new asymmetric masked distillation (AMD) framework for pre-training relatively small models with autoencoding. The core of AMD is to devise an asymmetric masking strategy, where the teacher model is enabled to see more context information with a lower masking ratio, while the student model is still equipped with a high masking ratio. We design customized multi-layer feature alignment between the teacher encoder and student encoder to regularize the pre-training of student MAE. To demonstrate the effectiveness and versatility of AMD, we apply it to both ImageMAE and VideoMAE for pre-training relatively small ViT models. AMD achieved 84.6% classification accuracy on IN1K using the ViT-B model. And AMD achieves 73.3% classification accuracy using the ViT-B model on the Something-in-Something V2 dataset, a 3.7% improvement over the original ViT-B model from VideoMAE. We also transfer AMD pre-trained models to downstream tasks and obtain consistent performance improvement over the original masked autoencoding. The code and models are available at https://github.com/MCG-NJU/AMD.