A Comparative Analysis of Multi-Head Self-Attention and Axial Attention in Transformer Architectures

Abstract

Transformer architectures have become foundational models across natural language processing, computer vision, and multimodal learning, primarily due to the effectiveness of self-attention mechanisms. However, the quadratic computational complexity of standard self-attention limits scalability for high-dimensional inputs, particularly in vision tasks requiring high-resolution processing. Axial attention has emerged as a promising alternative that factorizes attention along individual axes to improve efficiency while maintaining global receptive fields. This paper presents a comprehensive comparative analysis of Multi-Head SelfAttention and Axial Attention, examining their theoretical foundations, computational properties, representational capacity, and practical trade-offs across domains. Through empirical evaluation on medical image segmentation using a hybrid Convolutional Neural Network-Transformer architecture, we demonstrate that while MHSA achieves marginally higher accuracy, Axial Attention provides substantial computational efficiency gains, enabling processing of higher resolution inputs with significantly reduced memory footprint. Our analysis provides practical guidelines for selecting appropriate attention mechanisms based on task requirements, computational constraints, and input characteristics.