Low-latency vision transformers via large-scale multi-head attention

2025-06-30

AI Summary

To address computational redundancy, high inference latency, and accuracy limitations in Vision Transformers (ViTs), this paper proposes the Large-Scale Multi-Head Attention (LS-MHA) mechanism. Through per-head efficiency analysis, we reveal the collaborative behavior and sparse activation patterns of attention heads in classification tasks, identifying a small subset of dominant heads sufficient for label discrimination. Building on this insight, we introduce a soft committee fusion strategy that integrates multiple equivalent compact architectures and incorporate lightweight convolutional layers to accelerate early feature extraction. Evaluated on CIFAR-100, our compact Convolution-Transformer model achieves improved classification accuracy and reduced FLOPs without sacrificing precision, while significantly lowering inference latency. This work is the first to systematically uncover functional specialization among ViT attention heads and establishes a transferable attention optimization paradigm for efficient visual modeling.

Technology Category

• Computer Vision: Large Vision Models
• Machine Learning: Hardware-aware ML
• Search and Optimization: Non-convex Optimization

Application Category

• Web Mining and Content Analysis: Large pretrained models with web data
• Security and Privacy: Large-scale security measurements
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

The emergence of spontaneous symmetry breaking among a few heads of multi-head attention (MHA) across transformer blocks in classification tasks was recently demonstrated through the quantification of single-nodal performance (SNP). This finding indicates that each head focuses its attention on a subset of labels through cooperation among its SNPs. This underlying learning mechanism is generalized to large-scale MHA (LS-MHA) using a single matrix value representing single-head performance (SHP), analogous to single-filter performance in convolutional neural networks (CNNs). The results indicate that each SHP matrix comprises multiple unit clusters such that each label being explicitly recognized by a few heads with negligible noise. This leads to an increased signal-to-noise ratio (SNR) along the transformer blocks, thereby improving classification accuracy. These features give rise to several distinct vision transformer (ViT) architectures that achieve the same accuracy but differ in their LS-MHA structures. As a result, their soft committee yields superior accuracy, an outcome not typically observed in CNNs which rely on hundreds of filters. In addition, a significant reduction in latency is achieved without affecting the accuracy by replacing the initial transformer blocks with convolutional layers. This substitution accelerates early-stage learning, which is then improved by subsequent transformer layers. The extension of this learning mechanism to natural language processing tasks, based on quantitative differences between CNNs and ViT architectures, has the potential to yield new insights in deep learning. The findings are demonstrated using compact convolutional transformer architectures trained on the CIFAR-100 dataset.

Problem

• Improving classification accuracy via large-scale multi-head attention
• Reducing latency by combining convolutional and transformer layers
• Extending learning mechanisms to NLP tasks for new insights

Innovation

• Large-scale multi-head attention (LS-MHA) for improved SNR
• Convolutional layers replace initial transformer blocks to reduce latency
• Soft committee of distinct ViT architectures enhances accuracy