Window-Based Multi-Head Self-Attention

• Window-Based Multi-Head Self-Attention is a Transformer strategy that restricts attention to predefined contiguous windows, reducing computational load while preserving key context.
• It employs varied window schemes, such as MW-MHA and MSWA, to integrate local details and global information by assigning distinct window sizes across heads and layers.
• Empirical evaluations demonstrate enhanced efficiency and performance in applications like audio masked autoencoding and language modeling compared to conventional full attention.

Window-Based Multi-Head Self-Attention (MHSA) denotes a family of Transformer attention strategies in which the receptive field for each attention head is limited to a window—i.e., a contiguous or otherwise structured subset—of the input sequence. In contrast to conventional MHSA, where each head computes attention across all input tokens, window-based variants lower computational and memory costs by confining attention calculation to these windows. Moreover, approaches such as Multi-Window Multi-Head Attention (MW-MHA) and Multi-Scale Window Attention (MSWA) permit different heads and/or layers to utilize distinct window sizes, thus integrating local and global context modeling within a single network layer (Yadav et al., 2023, Xu et al., 2 Jan 2025).

1. Mathematical Formulation of Window-Based Multi-Head Attention

Standard MHSA projects the input $X \in \mathbb{R}^{n \times d}$ into query, key, and value matrices for each head, then computes attention via a scaled dot-product over the full sequence. Window-based approaches diverge by restricting each head's attention scope to a subset of the sequence.

In MW-MHA, given $n$ input tokens and $h$ heads, each head $i$ attends over a window of $\mathrm{win}_i$ tokens: $$\mathrm{MWMHA}(Q, K, V) = \operatorname{Concat}(\mathrm{winHead}_1, ..., \mathrm{winHead}_h) W^O$$ where each

$$\mathrm{winHead}_i = \mathrm{WinAttention}(Q W^Q_i, K W^K_i, V W^V_i, \mathrm{win}_i)$$

defines windowed attention via:

• Reshaping $Q_i, K_i, V_i$ into $$(n/\mathrm{win}_i) \times \mathrm{win}_i \times d_k$$
• Computing standard attention within each window: $$\mathrm{Attention}(Q_{\mathrm{win}}, K_{\mathrm{win}}, V_{\mathrm{win}}) = \mathrm{softmax}\left(\frac{Q_{\mathrm{win}} K_{\mathrm{win}}^T}{\sqrt{d_k}}\right) V_{\mathrm{win}}$$
• Re-flattening the windowed outputs.

MSWA introduces per-head and per-layer window variation in standard sliding window attention, where each position $n$0 and head $n$1 at layer $n$2 attends over the set $n$3 with $n$4 varying by both head and depth: $n$5 This configuration supports both multi-granularity context mixing and progressive expansion of receptive field with layer depth (Xu et al., 2 Jan 2025).

2. Window Allocation and Head Partitioning Strategies

MW-MHA adopts a domain-appropriate, data-dependent window assignment. All non-trivial divisors of the token count ($n$6) are taken as candidate window sizes, with two additional heads assigned full global context. For example, with $n$7, head window sizes become $n$8 (Yadav et al., 2023).

MSWA partitions both heads and layers to introduce scale diversity:

• Across heads (MSWA-h): Each of $n$9 heads in a layer is assigned one of four window sizes: $h$0, where $h$1 is the layer-wise base window (Xu et al., 2 Jan 2025).
• Across layers (MSWA-l): $h$2 layers have base windows scheduled geometrically by depth: $h$3, allocating progressively larger attention context as depth increases.

This table illustrates typical window assignments in MSWA:

Layer Group	Base Window $h$4	Per-head Windows
Shallowest	$h$5	$h$6, $h$7, $h$8, $h$9 (per head)
...	...	...
Deepest	$i$0	$i$1, $i$2, $i$3, $i$4 (per head)

Such partitioning ensures all heads collectively span a range of locality and scope per layer and across depth.

3. Computational Complexity and Implementation Considerations

Window-based MHSA variants offer significant efficiency gains compared to global attention:

• Full MHSA: $i$5 operations per layer (with $i$6 tokens, $i$7 heads)
• Fixed-window MHSA: $i$8 (with uniform window size $i$9)
• MW-MHA: $\mathrm{win}_i$0

For MW-MHA, the dominant computation arises from the small number of global heads ($\mathrm{win}_i$1), typically amounting to approximately $\mathrm{win}_i$2 of the full quadratic cost. The remainder, using small windows, contribute linear or near-linear terms (Yadav et al., 2023).

MSWA organizes implementation around efficient grouped attention kernels, grouping heads that share window sizes to maximize batched CUDA execution. The window pattern is static and constructed at model initialization, with masking handled by standard kernels (e.g., FlashAttention, xFormers) (Xu et al., 2 Jan 2025).

Both approaches maintain the standard multi-head projection structure: $\mathrm{win}_i$3, $\mathrm{win}_i$4, $\mathrm{win}_i$5 for each head, and final output projection $\mathrm{win}_i$6.

4. Empirical Evaluation and Scaling Properties

MW-MHA was evaluated in the Multi-Window Masked Autoencoder (MW-MAE) architecture on ten HEAR benchmark audio tasks, using linear probing on learned representations. Results indicated that MW-MAE consistently outperformed standard MAE (identically configured) by $\mathrm{win}_i$7 to $\mathrm{win}_i$8 points in normalized cross-task score; superior scaling characteristics were observed both for reduced patch size (more tokens) and increased encoder/decoder depth. Notably, MW-MAE demonstrated improved data efficiency, losing only $\mathrm{win}_i$9 of relative performance under $$\mathrm{MWMHA}(Q, K, V) = \operatorname{Concat}(\mathrm{winHead}_1, ..., \mathrm{winHead}_h) W^O$$0 data compared to $$\mathrm{MWMHA}(Q, K, V) = \operatorname{Concat}(\mathrm{winHead}_1, ..., \mathrm{winHead}_h) W^O$$1 for pure MAE (Yadav et al., 2023).

MSWA was benchmarked for language modeling (Wikitext-103 and enwik8) and common-sense reasoning via Llama-7B fine-tuning. MSWA matched or closed the gap to global MHSA while incurring only $$\mathrm{MWMHA}(Q, K, V) = \operatorname{Concat}(\mathrm{winHead}_1, ..., \mathrm{winHead}_h) W^O$$2 cost. For Wikitext-103, MSWA achieved perplexity $$\mathrm{MWMHA}(Q, K, V) = \operatorname{Concat}(\mathrm{winHead}_1, ..., \mathrm{winHead}_h) W^O$$3 at relative cost $$\mathrm{MWMHA}(Q, K, V) = \operatorname{Concat}(\mathrm{winHead}_1, ..., \mathrm{winHead}_h) W^O$$4 vs. full MHSA with perplexity $$\mathrm{MWMHA}(Q, K, V) = \operatorname{Concat}(\mathrm{winHead}_1, ..., \mathrm{winHead}_h) W^O$$5 at cost $$\mathrm{MWMHA}(Q, K, V) = \operatorname{Concat}(\mathrm{winHead}_1, ..., \mathrm{winHead}_h) W^O$$6. In 3-shot common-sense reasoning, MSWA improved accuracy by $$\mathrm{MWMHA}(Q, K, V) = \operatorname{Concat}(\mathrm{winHead}_1, ..., \mathrm{winHead}_h) W^O$$7 (MSWA $$\mathrm{MWMHA}(Q, K, V) = \operatorname{Concat}(\mathrm{winHead}_1, ..., \mathrm{winHead}_h) W^O$$8 vs. SWA $$\mathrm{MWMHA}(Q, K, V) = \operatorname{Concat}(\mathrm{winHead}_1, ..., \mathrm{winHead}_h) W^O$$9) (Xu et al., 2 Jan 2025).

5. Attention Head Feature Analysis and Hierarchical Effects

MW-MHA and MSWA exhibited distinct attention head behaviors:

• Entropy & Distance: MW-MAE encoders, even with uniform MHSA, evolved per-head attention entropies and wider mean attention distances, indicating a broader mix of local and global sensitivity in learned representations.
• Canonical Correlation (PWCCA): In MW-MAE decoders, heads across layers sharing window sizes learned highly correlated feature subspaces, supporting a "decoupled feature hierarchy." Such specialization was less pronounced in standard MAE decoders (Yadav et al., 2023).

MSWA's layerwise window scaling was found critical: geometric (2×) progression outperformed both reversed and arithmetically increasing schedules, with ablations confirming that multi-head, multi-layer diversity, not simple window size inflation, generated the observed accuracy gains (Xu et al., 2 Jan 2025).

6. Applications, Limitations, and Future Directions

Window-based MHSA has been deployed in large-scale audio representation learning (MW-MHA in MW-MAE) and autoregressive language modeling (MSWA in LLMs). The primary advantage lies in integrating fine-grained local modeling with efficient representation of long-range or global context, without full quadratic complexity.

Limitations include:

• The lack of truly unbounded context capture except via designated global heads or external mechanisms (e.g., linear attention).
• Fixed window scale/grouping hyperparameters chosen manually.
• Uniform head dimensions across all windows, lacking per-head adaptation.

Future research directions include learnable or dynamic window schedules, adaptive content-based windowing, integration with sparse or state-space attention mechanisms (for unbounded contexts), and fine-tuning of window groupings beyond fixed four-way splits (Xu et al., 2 Jan 2025).

7. Summary Table: Key Features of MW-MHA and MSWA

Feature	MW-MHA (Yadav et al., 2023)	MSWA (Xu et al., 2 Jan 2025)
Variation Across Heads	Yes (distinct windows per head)	Yes (4× scale per layer)
Variation Across Layers	No (in MW-MAE; always in decoder)	Yes (progressive from shallow→deep)
Global Heads	Yes (2 heads global window)	No native global (combine if needed)
Complexity (per layer)	$$\mathrm{winHead}_i = \mathrm{WinAttention}(Q W^Q_i, K W^K_i, V W^V_i, \mathrm{win}_i)$$0 + linear terms	$$\mathrm{winHead}_i = \mathrm{WinAttention}(Q W^Q_i, K W^K_i, V W^V_i, \mathrm{win}_i)$$1
Empirical Margins	$$\mathrm{winHead}_i = \mathrm{WinAttention}(Q W^Q_i, K W^K_i, V W^V_i, \mathrm{win}_i)$$2–$$\mathrm{winHead}_i = \mathrm{WinAttention}(Q W^Q_i, K W^K_i, V W^V_i, \mathrm{win}_i)$$3 points vs. MAE	Closes gap to full MHSA ($$\mathrm{winHead}_i = \mathrm{WinAttention}(Q W^Q_i, K W^K_i, V W^V_i, \mathrm{win}_i)$$41 ppt)
Main Applications	Audio masked autoencoding	Autoregressive language modeling