通过掩码微调弥合令牌剪枝与完整预训练之间的鸿沟

尽管Transformer在各种计算机视觉任务上取得成功，但它们承受着过高的内存与计算成本。一些工作提出动态视觉Transformer，通过剪枝冗余令牌来加速推理。改进令牌剪枝的一个关键是以训练良好的模型作为初始化，以实现更快收敛与更优性能。然而，当前基础模型通常采用全图训练，即使用完整图像作为输入并前向过程中保留整张特征图，这与逐步减少令牌的动态模型存在不一致，包括计算模式、信息量与令牌选择策略的不一致。受执行掩码与重建自监督任务的MAE启发，我们设计了掩码微调，通过对图像块掩码并基于剩余未掩码块预测图像类别标签，来弥合用作初始化预训练基础模型与基于令牌剪枝的动态视觉Transformer之间的鸿沟。在ImageNet上的大量实验表明，经掩码微调的基础模型获得了强遮挡鲁棒性与抗信息丢失能力。借助这一更好的初始化，Dynamic ViT取得了更高准确率，尤其在大令牌剪枝比例下（例如，基于DeiT的Dynamic ViT/0.8与Dynamic ViT/0.3分别为81.9% vs. 81.3%，以及62.3% vs. 58.9%）。此外，我们将方法应用于不同基于令牌剪枝的动态视觉Transformer、不同预训练模型及随机初始化模型，以证明泛化能力。