Automatic Channel Pruning for Multi-Head Attention

TL;DR

This paper introduces an automatic channel pruning method tailored for multi-head attention in Transformers, effectively reducing computation while maintaining accuracy, applicable to both original and linear attention mechanisms.

Contribution

The proposed method uniquely addresses channel misalignment in multi-head attention pruning, incorporating similarity-based weights and equal channel removal across heads.

Findings

Outperforms previous models on ImageNet-1K in accuracy and efficiency.

Effective pruning of multi-head attention with minimal accuracy loss.

Applicable to both original and linear attention mechanisms.

Abstract

Despite the strong performance of Transformers, their quadratic computation complexity presents challenges in applying them to vision tasks. Automatic pruning is one of effective methods for reducing computation complexity without heuristic approaches. However, directly applying it to multi-head attention is not straightforward due to channel misalignment. In this paper, we propose an automatic channel pruning method to take into account the multi-head attention mechanism. First, we incorporate channel similarity-based weights into the pruning indicator to preserve more informative channels in each head. Then, we adjust pruning indicator to enforce removal of channels in equal proportions across all heads, preventing the channel misalignment. We also add a reweight module to compensate for information loss resulting from channel removal, and an effective initialization step for pruning indicator based on difference of attention between original structure and each channel. Our proposed method can be used to not only original attention, but also linear attention, which is more efficient as linear complexity with respect to the number of tokens. On ImageNet-1K, applying our pruning method to the FLattenTransformer, which includes both attention mechanisms, shows outperformed accuracy for several MACs compared with previous state-of-the-art efficient models and pruned methods. Code will be available soon.