17 天前

用于预训练小型基础模型的非对称掩码蒸馏

Zhiyu Zhao, Bingkun Huang, Sen Xing, Gangshan Wu, Yu Qiao, Limin Wang
用于预训练小型基础模型的非对称掩码蒸馏
摘要

自监督基础模型凭借掩码自编码(masked autoencoding)的预训练范式,在计算机视觉领域展现出巨大潜力。模型规模是影响这类基础模型性能的关键因素之一。然而,大规模基础模型通常伴随着高昂的计算成本。本文聚焦于预训练相对小型的视觉Transformer模型,旨在实现对下游任务的高效适应。具体而言,受模型压缩中知识蒸馏思想的启发,我们提出一种新型的非对称掩码蒸馏(Asymmetric Masked Distillation, AMD)框架,用于对小型模型进行基于自编码的预训练。AMD的核心在于设计一种非对称掩码策略:教师模型以较低的掩码率运行,能够获取更丰富的上下文信息;而学生模型则保持较高的掩码率,以增强其对缺失信息的重建能力。此外,我们设计了定制化的多层特征对齐机制,实现教师编码器与学生编码器之间的深度协同,从而有效正则化学生模型的掩码自编码预训练过程。为验证AMD的有效性与通用性,我们将其应用于ImageMAE和VideoMAE,对小型ViT模型进行预训练。实验结果表明,在ImageNet-1K(IN1K)数据集上,使用ViT-B模型的AMD方法达到了84.6%的分类准确率;在Something-in-Something V2数据集上,同样使用ViT-B模型,AMD方法取得了73.3%的分类准确率,较VideoMAE原始的ViT-B模型提升了3.7%。此外,我们将AMD预训练模型迁移到多个下游任务中,均取得了优于原始掩码自编码方法的一致性能提升。相关代码与模型已开源,地址为:https://github.com/MCG-NJU/AMD。