Hallucinated-MHSA in Vision Transformers
- hMHSA is a modified ViT attention mechanism that computes half the attention maps directly and hallucinates the rest using lightweight convolutions.
- The approach reduces redundant computations, saving roughly 19% in FLOPs and up to a 25% parameter reduction while maintaining or improving accuracy.
- It uses a two-stage process—Intra-Head and Cross-Head Hallucination with 3×3 and 1×1 convolutions—to efficiently replicate attention information.
Hallucinated Multi-Head Self-Attention (hMHSA) is an architectural modification of the standard Multi-Head Self-Attention (MHSA) mechanism within the Vision Transformer (ViT) model family. Its central premise is to exploit the observed redundancy among attention heads in ViT by reducing the number of directly computed attention maps and regenerating (“hallucinating”) the remainder with lightweight convolutional operations. This approach achieves significant reductions in floating point operations (FLOPs) and parameters without degrading, and occasionally improving, model accuracy across multiple ViT backbones (Xu et al., 2023).
1. Background and Motivation
In canonical ViT architectures, every MHSA block computes separate attention maps by projecting input tokens into query and key spaces and correlating all possible pairs, resulting in quadratic complexity , where is the token count and is the embedding dimension. Empirical evaluation reveals substantial redundancy among attention heads, as evidenced by mean cosine similarities exceeding 50% between head pairs. This redundancy motivates the hypothesis that a subset of attention maps can be “hallucinated” from others, dramatically reducing compute requirements while maintaining (or even enhancing) model expressivity. The principal motivation for hMHSA is thus to address this inefficiency, targeting a reduction in head-level redundancy within self-attention computations (Xu et al., 2023).
2. Detailed Architectural Description
Given an input feature matrix , hMHSA is structured in the following sequence:
Real Attention Heads
- Project to half-width queries and keys: , .
- Project to values: .
- Reshape 0 and 1 into 2 heads, each with dimension 3.
- Compute scaled dot-product scores for each real head: 4 for 5.
Hallucinated Attention Heads
- Intra-Head Hallucination (IHH): 6 is reshaped into a pseudo-feature map of shape 7, with 8. Each channel (head) undergoes a 9 depth-wise convolution to locally refine its map.
- Cross-Head Hallucination (CHH): Output from IHH is processed with a 0 convolution across head channels, mixing information among the 1 hallucinated heads.
Aggregation and Output
- Concatenate the real and hallucinated pre-softmax maps along the head axis: 2.
- Apply softmax normalization to obtain 3.
- Compute 4 for each head, concatenate them, and re-project via 5 to obtain the final output.
This sequence allows hMHSA to compute only 6 attention heads via full 7 operations, with the remaining 8 generated via lightweight convolutions.
3. Mathematical Formulation
For reference, the standard MHSA block computes:
- 9, 0, 1, with each split into 2 heads of width 3.
- For 4, 5, 6, 7.
- Output: 8.
In hMHSA, only 9 heads use direct 0, while the other 1 are produced as:
- 2.
- Complete pre-softmax map: 3.
The hallucination mapping is composed of the depth-wise 4 convolution (IHH) and head-wise 5 convolution (CHH), both applied before softmax normalization.
4. Computational and Parametric Efficiency
A direct result of halving the number of expensive 6 products is a substantial reduction in computational cost. Specifically:
- Vanilla MHSA: 7
- hMHSA: 8
Since 9 in typical ViTs, the convolutional overhead is negligible, and the dominant 0 term is reduced by approximately 50%. Empirical results show overall self-attention FLOPs savings of 1 19%, and total ViT FLOPs decline by 10–20% depending on the backbone.
Parameter count is reduced as 2 and 3 each shrink from 4 to 5, yielding a 25% saving for these projections. The additional convolutional parameters for IHH and CHH are of order 6 and contribute marginally to the total.
5. Practical Integration and Hyperparameterization
hMHSA requires the number of heads 7 to be even. For backbones with odd 8, 9 is doubled and each head’s width halved, keeping the total dimensionality 0 constant. The mechanism is deployed by fully replacing MHSA modules with hMHSA in various ViT-derived models, including DeiT-T/S, PVTv2-b1, and NextViT-S. No modifications are made to the backbone’s training protocol. The new hyperparameters introduced are minimal and pertain to the IHH convolution (3×3 depth-wise, padding 1) and CHH convolution (1×1 along head dimension, output channels 1).
6. Empirical Evaluation and Redundancy Analysis
Extensive experiments confirm the computational advantage and retention of model quality:
- DeiT-T: FLOPs decrease from 1.26G to 1.02G (–19%), parameters from 5.72M to 4.68M (–18%), with Top-1 accuracy increasing from 72.2% to 72.9%.
- DeiT-S: analogous reductions (–19% FLOPs/–18.8% parameters), with a Top-1 accuracy gain of +0.3%.
- PVTv2-b1 and NextViT-S show similar trends, with no loss in accuracy.
Ablation studies on hallucination strategy demonstrate that employing both IHH and CHH is critical; skipping either yields accuracy degradation. The ordering IHH→CHH outperforms the reverse, indicating the benefit of local refinement before cross-head mixing.
Redundancy is quantitatively measured using Contribution Cosine Similarity (CCS) among heads. CCS drops in hMHSA compared to vanilla MHSA, indicating reduced head redundancy. Visual inspection corroborates that hallucinated maps exhibit diverse, nontrivial focus patterns, not simple replication of the real heads.
7. Generalization and Downstream Performance
hMHSA demonstrates transferability beyond ImageNet classification. On CIFAR-100 and Stanford-Dog, hMHSA in conjunction with the companion compact FFN module (cFFN) matches or marginally surpasses the accuracy of the backbone models. This suggests the approach generalizes well and does not compromise representation power on diverse tasks.
In summary, hMHSA effects a principled halving of high-cost attention computations in ViTs by combining direct computation and convolutional hallucination of attention heads. This architectural change delivers 10–20% savings in FLOPs and parameter count, accompanied by equivalent or improved accuracy over a range of canonical and hybrid ViT models (Xu et al., 2023).