视觉Transformer瘦身:连续优化空间中的多维搜索
计算机视觉与模式识别
2022-04-26 v2 人工智能
机器学习
摘要
本文探讨了从视觉Transformer中寻得最优子模型的可行性,并引入了一个纯粹的视觉Transformer瘦身(ViT-Slim)框架。它能够在多个维度上以端到端方式从原始模型中搜索子结构,包括输入令牌、MHSA 和 MLP 模块,并具有最先进的性能。我们的方法基于可学习的统一 稀疏约束与预定义因子,以反映不同维度连续搜索空间中的全局重要性。通过单次训练方案,搜索过程极为高效。例如,在 DeiT-S 上,ViT-Slim 的搜索过程仅需约 43 GPU 小时,且搜索出的结构在不同模块中具有多样的维度,十分灵活。随后,根据运行设备上准确率-FLOPs 权衡的需求采用预算阈值,并执行重训练过程以获得最终模型。大量实验表明,我们的 ViT-Slim 在各种视觉Transformer上可压缩多达 40% 的参数和 40% 的 FLOPs,同时在 ImageNet 上提升约 0.6% 的准确率。我们还在若干下游数据集上展示了所搜索模型的优势。我们的代码可在 https://github.com/Arnav0400/ViT-Slim 获取。
引用
@article{arxiv.2201.00814,
title = {Vision Transformer Slimming: Multi-Dimension Searching in Continuous Optimization Space},
author = {Arnav Chavan and Zhiqiang Shen and Zhuang Liu and Zechun Liu and Kwang-Ting Cheng and Eric Xing},
journal= {arXiv preprint arXiv:2201.00814},
year = {2022}
}
备注
CVPR 2022. Code is available at https://github.com/Arnav0400/ViT-Slim