[论文解读] Unified Visual Transformer Compression

Shixing Yu, Tianlong Chen+6 more|arXiv (Cornell University)|Mar 15, 2022

Advanced Neural Network Applications被引用 28

一句话总结

UVC 在预算约束下联合优化 ViT 的剪枝、块跳过和知识蒸馏，在 ImageNet 上对 DeiT 和 T2T-ViT 主干实现了显著的 FLOPs 减少，且精度损失极小。

ABSTRACT

Vision transformers (ViTs) have gained popularity recently. Even without customized image operators such as convolutions, ViTs can yield competitive performance when properly trained on massive data. However, the computational overhead of ViTs remains prohibitive, due to stacking multi-head self-attention modules and else. Compared to the vast literature and prevailing success in compressing convolutional neural networks, the study of Vision Transformer compression has also just emerged, and existing works focused on one or two aspects of compression. This paper proposes a unified ViT compression framework that seamlessly assembles three effective techniques: pruning, layer skipping, and knowledge distillation. We formulate a budget-constrained, end-to-end optimization framework, targeting jointly learning model weights, layer-wise pruning ratios/masks, and skip configurations, under a distillation loss. The optimization problem is then solved using the primal-dual algorithm. Experiments are conducted with several ViT variants, e.g. DeiT and T2T-ViT backbones on the ImageNet dataset, and our approach consistently outperforms recent competitors. For example, DeiT-Tiny can be trimmed down to 50\% of the original FLOPs almost without losing accuracy. Codes are available online:~\url{https://github.com/VITA-Group/UVC}.

研究动机与目标

• 通过降低计算成本而不牺牲准确性，促进视觉变换器的高效部署。
• 提出一个在单一端到端优化中结合多种压缩技术的统一框架。
• 在全局资源预算下实现剪枝、跳过和蒸馏的自动平衡。
• 在 ImageNet 上对流行的 ViT 变体（DeiT 和 T2T-ViT）进行评估，并与最先进的剪枝方法进行比较。

提出的方法

• 将 ViT 压缩表述为一个带约束的端到端优化问题。
• 在蒸馏损失下共同学习网络权重、按层的剪枝掩码/比率和跳过配置。
• 在每个 transformer block 内使用两级组稀疏性来剪枝注意力头和 MLP 神经元。
• 为跨 transformer 块的跳过配置引入二进制门控变量，以实现块跳过。
• 通过原-对偶算法求解极小极大优化问题。
• 提供参考实现及在 DeiT 与 T2T-ViT 主干上的实验结果。

实验结果

研究问题

• RQ1 可以在资源预算下，针对 ViT blokc 内以及跨块的跳过配置，联合优化剪枝吗？
• RQ2 将剪枝、跳过和知识蒸馏结合在一起，是否优于单一方法的 ViT 压缩？
• RQ3 在不同的 FLOPs 预算下，DeiT 和 T2T-ViT 能实现哪些精度/延迟权衡？
• RQ4 UVC 与现有的 Patch 递减和剪枝方法在 ImageNet 上的表现如何？

主要发现

Model	Method	Top-1 Acc. (%)	FLOPs (G)	FLOPs remained (%)
DeiT-Tiny	Baseline	72.2	1.3	100.0
DeiT-Tiny	SViTE	70.12	0.99	76.31
DeiT-Tiny	PatchSlimming	72.0	0.7	53.8
DeiT-Tiny	UVC	71.8	0.69	53.1
DeiT-Tiny	HVT	69.7	0.64	49.23
DeiT-Tiny	UVC (distilled)	71.3	0.64	49.23
DeiT-Tiny	UVC (other setting)	70.6	0.51	39.12
DeiT-Small	Baseline	79.8	4.6	100.0
DeiT-Small	SViTE	79.22	3.14	68.36
DeiT-Small	PatchSlimming	79.4	2.6	56.5
DeiT-Small	HVT	78.0	2.40	52.2
DeiT-Small	UVC	78.82	2.32	50.41
DeiT-Base	Baseline	81.8	17.6	100.0
DeiT-Base	SViTE	82.22	11.87	66.87
DeiT-Base	VTP	80.7	10.0	56.8
DeiT-Base	PatchSlimming	81.5	9.8	55.7
DeiT-Base	UVC	80.57	8.0	45.50
T2T-ViT-14	Baseline	81.5	4.8	100.0
T2T-ViT-14	PoWER	79.9	3.5	72.9
T2T-ViT-14	UVC	80.4	2.90	60.4
T2T-ViT-14	UVC (lower)	79.6	2.47	51.5
T2T-ViT-14	UVC (higher)	78.9	2.11	44.0

• UVC 可以在 DeiT-Tiny 上将 FLOPs 降低约 50%，在蒸馏/非蒸馏设置下精度损失极小（0.3%/0.9%）。
• UVC 在相似或更大 FLOPs 减少的情况下，在准确率方面优于若干现代 ViT 压缩方法（如 SViTE、PatchSlimming、HVT、IA-RED 2）。
• 对于 DeiT-Base，UVC 实现约 45% 的 FLOPs 下降，精度下降在可接受范围内，在 some 基线不易稳定训练的情况下仍具竞争力。
• 应用于 T2T-ViT-14 时，UVC 实现显著的 FLOPs 节省（如 44–60% 区间），精度下降温和，超过了一些基于 Patch 的减少方法。
• 消融实验表明，联合优化跳过、剪枝和蒸馏的权衡优于任一单一组件或逐步应用。
• UVC 的代码对外开源。
• （文本中的注释）UVC 在 ImageNet 的多种 ViT 变体上表现出强劲的性能。