[论文解读] A Unified Pruning Framework for Vision Transformers

Hao Yu, Jianxin Wu|arXiv (Cornell University)|Nov 30, 2021

Advanced Neural Network Applications参考文献 43被引用 5

一句话总结

该论文提出UP-ViTs，一种用于视觉Transformer及其变体的统一结构化剪枝框架，可同时剪枝所有组件——自注意力、前馈网络、归一化层和卷积层——而无需改变模型架构。通过使用高效的重要性评分模块和渐进式剪枝策略，UP-ViTs在压缩模型上实现了最先进（SOTA）的准确率，例如UP-DeiT-T在与原始DeiT-T相同FLOPs下，ImageNet上的top-1准确率达到75.79%。

ABSTRACT

Recently, vision transformer (ViT) and its variants have achieved promising performances in various computer vision tasks. Yet the high computational costs and training data requirements of ViTs limit their application in resource-constrained settings. Model compression is an effective method to speed up deep learning models, but the research of compressing ViTs has been less explored. Many previous works concentrate on reducing the number of tokens. However, this line of attack breaks down the spatial structure of ViTs and is hard to be generalized into downstream tasks. In this paper, we design a unified framework for structural pruning of both ViTs and its variants, namely UP-ViTs. Our method focuses on pruning all ViTs components while maintaining the consistency of the model structure. Abundant experimental results show that our method can achieve high accuracy on compressed ViTs and variants, e.g., UP-DeiT-T achieves 75.79% accuracy on ImageNet, which outperforms the vanilla DeiT-T by 3.59% with the same computational cost. UP-PVTv2-B0 improves the accuracy of PVTv2-B0 by 4.83% for ImageNet classification. Meanwhile, UP-ViTs maintains the consistency of the token representation and gains consistent improvements on object detection tasks.

研究动机与目标

• 解决视觉Transformer缺乏有效且可泛化的剪枝方法的问题，同时保持结构一致性。
• 克服基于token的剪枝方法的局限性，后者会破坏空间结构并损害在下游任务中的迁移能力。
• 开发一种适用于原始ViT及其变体（包括具有卷积或层次设计的模型）的统一框架。
• 实现高准确率、低计算量的压缩ViT模型，使其在目标检测等密集预测任务中具有良好泛化能力。
• 通过联合剪枝所有组件并保持模型完整性，实现优于现有剪枝方法的性能表现。

提出的方法

• 设计一个高效的评估模块，使用基于梯度的幅度估计方法，计算预训练ViT组件中每个滤波器的重要性得分。
• 通过同时且统一地移除所有组件（包括多头自注意力、前馈网络、归一化层和卷积层）中的冗余通道，实施结构化剪枝。
• 提出一种新颖的多头自注意力层通道移除策略，以保持注意力头结构并维持表征一致性。
• 引入一种渐进式块剪枝方法，通过迭代移除最不重要的块，并用混合块替代，以维持性能。
• 在微调过程中使用知识蒸馏技术，以在剪枝后保持准确率，教师模型与原始架构相匹配。
• 在多个ViT变体（包括DeiT、PVTv2和Swin Transformer）上应用统一的剪枝流程，确保压缩后仍保持架构一致性。

实验结果

研究问题

• RQ1 能否设计一种统一的剪枝框架，在不改变其结构设计的前提下，同时压缩视觉Transformer的所有组件？
• RQ2 与基于token的剪枝相比，对所有层（包括FFN、MSA、归一化层和卷积层）进行结构化剪枝如何影响准确率和推理效率？
• RQ3 在剪枝过程中保持原始token表征是否能提升在目标检测等下游任务中的迁移能力？
• RQ4 一次性剪枝与渐进式剪枝策略在压缩ViT宽度和深度方面，其相对有效性如何？
• RQ5 与现有最先进压缩ViT相比，剪枝后的模型是否能在多样化的架构和任务中实现更好的泛化能力？

主要发现

• UP-DeiT-T在ImageNet上的top-1准确率达到75.79%，在相同FLOPs下比原始DeiT-T高出3.59%。
• 在相同计算成本下，UP-PVTv2-B0在ImageNet分类任务上的准确率比原始PVTv2-B0高出4.83%。
• 当将UP-DeiT-T压缩至10个块时，渐进式块剪枝的准确率（73.78%）高于单次块移除（73.20%）。
• 在相同任务中，对组件1（宽度）进行单次剪枝可达到66.95%的准确率，优于渐进式剪枝的65.05%。
• 由于保留了token表征，剪枝后的UP-ViTs模型在下游任务（包括目标检测）中保持了稳定的性能增益。
• 该框架在多种ViT变体（包括DeiT、PVTv2和Swin Transformer）上均表现出有效的泛化能力，展示了广泛的适用性和鲁棒性。