用于预训练小型基础模型的非对称掩码蒸馏
计算机视觉与模式识别
2024-04-02 v2
摘要
得益于掩码自编码的预训练范式，自监督基础模型在计算机视觉中展现出巨大潜力。规模是影响这些基础模型性能的主要因素。然而，这些大型基础模型往往带来高计算成本。本文聚焦于预训练相对较小的视觉 transformer 模型，以便高效适配下游任务。具体而言，受模型压缩中知识蒸馏的启发，我们提出一种新的非对称掩码蒸馏（AMD）框架，通过自编码预训练相对较小的模型。AMD 的核心在于设计一种非对称掩码策略，其中教师模型以较低掩码率看到更多上下文信息，而学生模型仍保持高掩码率。我们设计了教师编码器与学生编码器之间的定制化多层特征对齐，以正则化学生 MAE 的预训练。为证明 AMD 的有效性与通用性，我们将其应用于 ImageMAE 和 VideoMAE 以预训练相对较小的 ViT 模型。AMD 使用 ViT-B 模型在 IN1K 上取得 84.6% 的分类准确率。且 AMD 在 Something-in-Something V2 数据集上使用 ViT-B 模型取得 73.3% 的分类准确率，较 VideoMAE 原始 ViT-B 模型提升 3.7%。我们还将 AMD 预训练模型迁移至下游任务，相较原始掩码自编码获得一致的性能提升。代码与模型见 https://github.com/MCG-NJU/AMD。
引用
@article{arxiv.2311.03149,
title = {Asymmetric Masked Distillation for Pre-Training Small Foundation Models},
author = {Zhiyu Zhao and Bingkun Huang and Sen Xing and Gangshan Wu and Yu Qiao and Limin Wang},
journal= {arXiv preprint arXiv:2311.03149},
year = {2024}
}
备注
Accepted by CVPR 2024