Compact Vision Transformer by Reduction of Kernel Complexity

TL;DR

This paper introduces KCR-Transformer, a compact vision transformer with channel selection guided by a novel theoretical generalization bound, reducing computational cost while maintaining or improving accuracy.

Contribution

The work proposes a new KCR-Transformer block with differentiable channel selection and provides a rigorous generalization bound, enabling effective channel pruning in vision transformers.

Findings

Reduces FLOPs of vision transformers significantly.

Maintains or improves prediction accuracy after pruning.

Achieves superior performance on various vision tasks.

Abstract

Self-attention and transformer architectures have become foundational components in modern deep learning. Recent efforts have integrated transformer blocks into compact neural architectures for computer vision, giving rise to various efficient vision transformers. In this work, we introduce Transformer with Kernel Complexity Reduction, or KCR-Transformer, a compact transformer block equipped with differentiable channel selection, guided by a novel and sharp theoretical generalization bound. KCR-Transformer performs input/output channel selection in the MLP layers of transformer blocks to reduce the computational cost. Furthermore, we provide a rigorous theoretical analysis establishing a tight generalization bound for networks equipped with KCR-Transformer blocks. Leveraging such strong theoretical results, the channel pruning by KCR-Transformer is conducted in a generalization-aware manner, ensuring that the resulting network retains a provably small generalization error. Our KCR-Transformer is compatible with many popular and compact transformer networks, such as ViT and Swin, and it reduces the FLOPs of the vision transformers while maintaining or even improving the prediction accuracy. In the experiments, we replace all the transformer blocks in the vision transformers with KCR-Transformer blocks, leading to KCR-Transformer networks with different backbones. The resulting TCR-Transformers achieve superior performance on various computer vision tasks, achieving even better performance than the original models with even less FLOPs and parameters.