AMAP: Automatic Multihead Attention Pruning by Similarity-Based Pruning Indicator.

Area of Science:

• Computer Vision
• Machine Learning
• Artificial Intelligence

Background:

• Transformers achieve high performance in vision tasks but suffer from quadratic computation complexity in self-attention.
• Linear attention offers a linear complexity alternative, improving the computation-performance tradeoff.
• Automatic pruning optimizes model structure for resource constraints but faces challenges with multi-head attention channel mismatch.

Purpose of the Study:

• To propose an automatic pruning method for linear attention mechanisms in vision transformers that resolves channel mismatch issues.
• To enhance computational efficiency and performance of transformer models for vision tasks.

Main Methods:

• Integrated channel similarity-based weights into the pruning indicator to preserve informative channels within each attention head.
• Adjusted the pruning indicator to ensure even channel removal across all heads, preventing channel mismatch.
• Incorporated a reweight module to mitigate information loss and introduced an effective pruning indicator initialization for linear attention.

Main Results:

• Achieved a 30% reduction in FLOPs (Floating Point Operations Per Second) for the FLattenTransformer on ImageNet-1K with near-lossless performance.
• Demonstrated a 1.96% accuracy gain over the DeiT-B model while reducing FLOPs by 37%.
• Showcased a 1.05% accuracy increase over the Swin-B model with a 10% FLOPs reduction.

Conclusions:

• The proposed automatic pruning method effectively optimizes linear attention mechanisms in vision transformers.
• The method successfully addresses channel mismatch, leading to significant computational savings without compromising performance.
• Outperforms existing state-of-the-art efficient models and recent pruning techniques in terms of efficiency and accuracy.