WMSA-WBS: Efficient Wave Multi-Head Self-Attention with Wavelet Bottleneck.

Area of Science:

• Computer Science
• Artificial Intelligence
• Machine Learning

Background:

• Multi-head self-attention (MSA) is crucial for Vision Transformers (ViT) but struggles with local feature capture and high computational demands.
• Existing wavelet-enhanced ViT methods often isolate wavelet decomposition within attention layers, limiting comprehensive feature learning.

Purpose of the Study:

• To propose an integrated multi-head self-attention (WMSA) with a bottleneck enhancement structure (WBS), termed WMSA-WBS, to address the limitations of conventional MSA in ViT architectures.
• To achieve more efficient and comprehensive feature learning by co-designing wavelet-based frequency processing and bottleneck optimization.

Main Methods:

• Introduced a novel wavelet multi-head self-attention (WMSA) combined with a wavelet bottleneck structure (WBS).
• Developed a co-design approach integrating frequency processing and bottleneck optimization for enhanced feature learning.
• Focused on capturing global and local information across spatial, frequency, and channel domains with low computational complexity.

Main Results:

• WMSA-WBS effectively mitigates the limitations of conventional MSA in ViT models.
• The proposed module captures both global and local information efficiently, maintaining low computational complexity and memory usage.
• ViT models enhanced with WMSA-WBS demonstrated superior accuracy-complexity trade-offs in image classification, object detection, and semantic segmentation tasks.

Conclusions:

• WMSA-WBS offers a significant advancement in ViT architecture by improving feature learning efficiency and effectiveness.
• The integrated wavelet and bottleneck approach provides a better balance between performance and computational cost for vision tasks.
• This method presents a promising direction for developing more capable and efficient Vision Transformer models.