统一的视觉 Transformer 压缩
机器学习
2022-03-17 v1 计算机视觉与模式识别
摘要
视觉 Transformer（ViTs）近年来受到广泛关注。即使没有卷积等定制化的图像算子，ViTs 在大规模数据上经过恰当训练后也能取得有竞争力的性能。然而，由于堆叠多头自注意力模块等原因，ViTs 的计算开销仍然过高。与压缩卷积神经网络方面丰富的文献和普遍的成功相比，视觉 Transformer 压缩的研究才刚刚兴起，且现有工作仅聚焦于压缩的一个或两个方面。本文提出了一种统一的 ViT 压缩框架，无缝集成了三种有效技术：剪枝、层跳过和知识蒸馏。我们构建了一个预算约束的端到端优化框架，在蒸馏损失下联合学习模型权重、逐层剪枝比例/掩码以及跳过配置。该优化问题随后使用原始对偶算法求解。实验在多个 ViT 变体上进行，例如在 ImageNet 数据集上的 DeiT 和 T2T-ViT 骨干网络，我们的方法持续优于近期竞争对手。例如，DeiT-Tiny 可被压缩至原始 FLOPs 的 50% 且几乎不损失精度。代码已在线提供：~\url{https://github.com/VITA-Group/UVC}。
引用
@article{arxiv.2203.08243,
title = {Unified Visual Transformer Compression},
author = {Shixing Yu and Tianlong Chen and Jiayi Shen and Huan Yuan and Jianchao Tan and Sen Yang and Ji Liu and Zhangyang Wang},
journal= {arXiv preprint arXiv:2203.08243},
year = {2022}
}
备注
Accepted by ICLR'22