Unified Visual Transformer Compression

TL;DR

This paper introduces a unified framework for compressing Vision Transformers by combining pruning, layer skipping, and knowledge distillation, achieving significant FLOPs reduction with minimal accuracy loss.

Contribution

It presents a novel end-to-end optimization approach that jointly learns weights, pruning ratios, and skip configurations for ViT compression.

Findings

DeiT-Tiny reduced to 50% FLOPs with minimal accuracy loss

Consistent outperformance over recent competitors on ImageNet

Effective joint optimization of multiple compression techniques

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}.