MOHSA: Overlapped MHSA in Vision Transformers

• MOHSA is a variant of MHSA that introduces structured overlap between adjacent attention heads to enable direct cross-head information exchange.
• It employs configurable overlap schemes, including fixed and layer-wise increasing or decreasing patterns, to optimize token mixing in vision transformers.
• Empirical results demonstrate Top-1 accuracy improvements of up to 7.4% with only a modest increase in parameters and FLOPs across various vision benchmarks.

Multi-Overlapped-Head Self-Attention (MOHSA) is an architectural variant of the Multi-Head Self-Attention (MHSA) mechanism employed in Transformer models, particularly Vision Transformers (ViTs). MOHSA introduces a structured overlap between adjacent attention heads, enabling direct information exchange during attention computation, in contrast to the hard partitioning of query, key, and value (Q, K, V) representations in standard MHSA. This approach is designed to enhance token interaction and inter-head representational mixing with minimal computational overhead, yielding measurable gains on vision tasks across multiple architectures and benchmark datasets (Zhang et al., 2024).

1. Rationale for Overlapping Attention Heads

Traditional MHSA computes attention in $h$ independent heads, each operating on a disjoint $d_h = d / h$-dimensional chunk of Q, K, V, where $d$ is the model dimension. This head-wise decorrelation fosters diversity, but completely suppresses cross-head information flow at the point of attention, relying on later linear projections for any cross-head synthesis. In image recognition settings, this may hinder the model’s ability to aggregate diverse spatial or semantic relationships, which are often distributed across heads.

MOHSA addresses this by overlapping adjacent heads’ subspaces. For a given head $i$, its $Q'_i$, $K'_i$, and $V'_i$ are augmented with a weighted contribution from the immediate neighboring heads, allowing each head’s contextual window to “bleed over” and directly draw information from adjacent subspaces. This results in richer, more entangled representations, with only a small increase in parameter count and floating-point operations.

2. Mathematical Description

Let $X\in\mathbb{R}^{T\times d}$ denote a sequence of $T$ tokens. Projecting to query, key, and value yields:

$Q = X W^Q,\quad K = X W^K,\quad V = X W^V,$

with $d_h = d / h$0. In MHSA, $d_h = d / h$1, $d_h = d / h$2, $d_h = d / h$3 are partitioned along the feature axis:

$d_h = d / h$4

with analogous splits for $d_h = d / h$5 and $d_h = d / h$6. MOHSA instead defines, for overlap fraction $d_h = d / h$7:

$d_h = d / h$8

subject to boundary conventions (e.g., for $d_h = d / h$9, $d$0). Each head thus aggregates its own and immediate neighbors’ representations.

The per-head attention is computed as:

$d$1

and the output concatenation and projection:

$d$2

where $d$3 projects the result back to model dimension.

Alternatively, the overlap can be parametrized as concatenating portions $d$4 and linearly projecting to $d$5 for each head.

3. Overlap Schemes and Layerwise Scheduling

MOHSA enables various paradigms for controlling the overlap per layer. Let $d$6 denote the number of overlap features at layer $d$7, so $d$8. The explored paradigms include:

• Fixed 1: $d$9 for all layers
• Fixed Half: $i$0 for all layers
• inc-0($i$1): $i$2, then increment by 1 every $i$3 layers (0-indexed)
• inc-1($i$4): $i$5, then increment every $i$6 layers (1-indexed)
• dec-0($i$7)/dec-1($i$8): symmetric decrease schemes

For example, inc-0(2) across 12 ViT layers yields $i$9. These variants allow for both fixed and depth-varying overlap to be studied, with empirical results showing that models with increasing or decreasing overlaps across depth can outperform those using constant overlap, especially in deeper architectures (Zhang et al., 2024).

4. Implementation Details and Computational Complexity

A MOHSA block proceeds as follows (condensed pseudocode):

$Q'_i$7

Relative to standard MHSA, the per-head compute grows as $Q'_i$0, where $Q'_i$1 is the average overlap. For typical overlaps $Q'_i$2, resulting in overheads of 5–15% for FLOPs and parameter count.

5. Empirical Evaluation

MOHSA was benchmarked on multiple Transformer backbones and datasets:

• Models: ViT-Tiny, ViT-Small (12 heads, 12 layers), CaiT-xxs12/24 (4 heads, 12/24 layers plus class layers), Swin-Tiny (staged heads, 12 total layers)
• Datasets: CIFAR-10, CIFAR-100, Tiny-ImageNet, ImageNet-1k
• Training: 100 epochs + 20 warmup; AdamW optimizer; batch sizes: 128 (CIFAR/Tiny-INet), 512 (ImageNet); initial LR 5e-4 (ViT/CaiT), 2.5e-4 or 5e-4 (Swin-Tiny)

Top-1 accuracy was the principal evaluation metric, with parameter count and FLOPs also reported.

A summary of results:

Dataset	Model	Orig Acc (%)	MOHSA Acc (%)	ΔAcc	Params Δ	FLOPs Δ
CIFAR-10	ViT-Tiny	84.31	85.65	+1.34	+0.0M	+0.0G
CIFAR-10	ViT-Small	86.53	87.30	+0.77	+0.6M	+0.3G
CIFAR-10	CaiT-xxs12	77.04	80.04	+3.00	+0.0M	+0.0G
CIFAR-100	ViT-Tiny	61.88	63.01	+1.13	+0.0M	+0.0G
CIFAR-100	CaiT-xxs24	50.94	55.94	+5.00	+0.1M	+0.1G
Tiny-INet	ViT-Tiny	50.79	52.40	+1.61	+0.2M	+0.1G
Tiny-INet	CaiT-xxs24	42.46	49.87	+7.41	+0.4M	+0.3G
ImageNet	ViT-Tiny	69.25	70.31	+1.06	+0.5M	+0.2G
ImageNet	CaiT-xxs24	70.56	74.26	+3.70	+0.5M	+0.3G

Absolute Top-1 gains of 0.2–7.4% were observed, with sub-10% increases in parameters and FLOPs (Zhang et al., 2024).

6. Model Behavior and Ablation Insights

Key insights from ablation studies:

• Overlapping the value ($Q'_i$3) projections yields much of the benefit, but overlapping all ($Q'_i$4) is optimal.
• Small fixed overlaps ($Q'_i$5 or $Q'_i$6) provide measurable gains, but schedules where the overlap dimension increases or decreases with layer depth (“inc”/“dec”) can further improve accuracy, especially in deeper models.
• Deeper architectures (e.g., CaiT-xxs24) derive the greatest benefit, with performance boosts of 3–7% without additional model depth or head count. This suggests that head overlap facilitates better information mixing and gradient flow in deeper stacks.
• Learning curves show that MOHSA models match the early-phase convergence of MHSA baselines but surpass them after several epochs, indicating improved learning stability and final capacity.

7. Applicability, Extensions, and Research Directions

MOHSA functions as a drop-in substitution for MHSA in any Transformer encoder, with immediate applicability to vision models. Several further lines of investigation are proposed:

• End-to-end learnable overlaps, with per-head or per-layer coefficients
• Non-local head interaction schemes, such as block-circulant or global mixing
• Porting the approach to domains beyond vision, including language, speech, and graph Transformer models
• Synergy with other head-interaction variants, such as re-attention or talking-heads architectures

The operational simplicity, minimal computational cost, and consistent performance improvement position MOHSA as a general mechanism for enhancing tokenwise and cross-head information propagation within attention-based architectures (Zhang et al., 2024).