Head-Partitioning Self-Attention (HPSA)

• HPSA is a self-attention method that partitions the global attention workload across different heads to minimize redundant computations.
• It employs techniques such as causal distance bands, role-specific masks, and low-dimensional factorization to assign distinct roles to heads.
• This structured approach improves computational efficiency and performance by distributing attention calculations and promoting head specialization.

Head-Partitioning Self-Attention (HPSA) denotes self-attention mechanisms in which the global attention workload is explicitly structured across heads rather than redundantly recomputed by every head. In the strongest formulation, different heads cover disjoint or nearly disjoint subsets of the token–token dependency matrix such that the union recovers the full causal or full attention space; related formulations constrain heads by role-specific masks, shared low-dimensional factors, axial decompositions, token partitions, or runtime head chunks (Zhao et al., 12 Nov 2025, Wang et al., 2020, Kang et al., 2024). Taken together, these works suggest that HPSA is best understood as a family of design patterns for reducing cross-head redundancy, inducing specialization, and improving efficiency while preserving the query–key–value structure of self-attention.

1. Formal definition and relation to standard multi-head attention

In standard multi-head self-attention, all heads receive projections of the same sequence and compute attention over the same token–token grid. For input $X \in \mathbb{R}^{N \times d_{\text{model}}}$, head $h$ computes

$$\text{Attn}_h(Q_h, K_h, V_h)=\mathrm{softmax}\!\left(\frac{Q_hK_h^\top}{\sqrt{d_k}} + M\right)V_h,$$

and multi-head composition concatenates the per-head outputs before the output projection (Zhao et al., 12 Nov 2025). In this baseline, the core score computation scales as $O(H \cdot N^2 \cdot d_k)$, and each head attends over the full causal matrix (Zhao et al., 12 Nov 2025).

HPSA modifies this premise by assigning different heads, or different subsets of heads, distinct admissible attention domains or interaction roles. A generic HPSA condition is that the full attention graph is decomposed into head-associated regions with full coverage and controlled overlap: $$\bigcup_h \mathrm{region}_h = \text{full attention graph}, \qquad \mathrm{region}_h \cap \mathrm{region}_{h'} = \emptyset \ \text{or controlled}.$$ SPAttention makes this definition explicit by partitioning the causal distance spectrum into balanced, contiguous, non-overlapping bands, one per head, so that every causal pair $(i,j)$ is preserved exactly once across heads (Zhao et al., 12 Nov 2025). The role-guided masking approach provides a second, more heterogeneous formulation: the head set is partitioned into five fixed role-constrained heads and the remaining unconstrained heads, with role-specific masks added directly to the logits (Wang et al., 2020).

This framing distinguishes HPSA from two adjacent ideas. First, it is not identical to head pruning: pruning removes or gates heads, whereas HPSA assigns differentiated computational responsibility to heads. Second, it is not identical to generic efficient attention: some efficient mechanisms compress token interactions without structurally differentiating heads. The literature therefore uses HPSA both as a literal partition of the attention graph across heads and as a broader principle for structuring head-level computation (Zhao et al., 12 Nov 2025, Jiang et al., 2023).

2. Structural realizations

The literature instantiates head partitioning along several axes. Some mechanisms partition the attention graph directly across heads; others partition roles, factors, spatial axes, or runtime execution units, while preserving the notion that heads should not all perform the same computation.

Mechanism	What is partitioned	Defining property
SPAttention	Causal distance spectrum across heads	Balanced, contiguous, non-overlapping bands
Transformer-Guided-Attn	Head set into fixed roles plus free heads	Role-specific masks on logits
iMHSA	Cross-head interaction in low-dimensional factors	Query-less and key-less factorization
GlobalMind	Spatial/temporal modules and axial directions	GAS with row-wise or column-wise global attention
MHPA	Spatial tokens, shared across heads	Intra-partition and inter-partition aggregation
HCMS	Runtime head chunks	Communication–computation overlap

SPAttention is the clearest direct HPSA construction. For sequence length $N$ and $H$ heads, it defines per-head distance bands by

$$W_{\text{base}}=\left\lfloor \frac{N}{H}\right\rfloor,\quad R=N\bmod H,\quad W_h=W_{\text{base}}+\mathbb{I}(h<R),\quad S_h=h\cdot W_{\text{base}}+\min(h,R),$$

and then allows head $h$ to attend only to keys whose causal distance satisfies $h$0 (Zhao et al., 12 Nov 2025). The result is disjointness per query and completeness over heads.

The role-guided formulation partitions heads functionally rather than geometrically. In each layer, one head is assigned to each of RareW, Seprat, DepSyn, MajRel, and RelPos, while the remaining heads are regular heads (Wang et al., 2020). The mask enters additively before softmax: $h$1 Here partitioning is not by distance band but by admissible linguistic or positional structure.

Interactive Multi-Head Self-Attention (iMHSA) is directly informative for HPSA-style design, although it does not explicitly split heads into named groups. Its key move is to treat each head’s attention matrix as a feature and to enable cross-head interaction after decomposing attention into low-dimensional query-less and key-less factors,

$h$2

with learned head-mixing matrices acting on the factors rather than on full $h$3 attention maps (Kang et al., 2024). This suggests an HPSA variant in which heads coordinate through a shared low-dimensional bottleneck rather than through disjoint masks.

GlobalMind introduces a different form of structuring. It does not define explicit per-head groups within a module, but it partitions attention at the module level into row-wise or column-wise global spatial attention through Global Axial Segmentation (GAS), and into temporal interaction modules whose $h$4, $h$5, and $h$6 are drawn from $h$7, $h$8, and $h$9, respectively (Hu et al., 2023). Multi-head outputs are then fused by a $$\text{Attn}_h(Q_h, K_h, V_h)=\mathrm{softmax}\!\left(\frac{Q_hK_h^\top}{\sqrt{d_k}} + M\right)V_h,$$0 convolution: $$\text{Attn}_h(Q_h, K_h, V_h)=\mathrm{softmax}\!\left(\frac{Q_hK_h^\top}{\sqrt{d_k}} + M\right)V_h,$$1 This is an implicit head-structuring mechanism: the axial choice and temporal/spatial separation differentiate the representational roles available to heads.

DualFormer’s Multi-Head Partition-wise Attention (MHPA) is not literal head partitioning, because all heads share the same spatial partitions, but it is a partitioned self-attention design that serves as a blueprint for HPSA generalization (Jiang et al., 2023). Tokens are clustered into $$\text{Attn}_h(Q_h, K_h, V_h)=\mathrm{softmax}\!\left(\frac{Q_hK_h^\top}{\sqrt{d_k}} + M\right)V_h,$$2 partitions, intra-partition and inter-partition aggregation are computed, and the resulting global interaction feature replaces dense token-wise attention.

HCMS extends the idea into systems design. It partitions the head dimension itself into chunks,

$$\text{Attn}_h(Q_h, K_h, V_h)=\mathrm{softmax}\!\left(\frac{Q_hK_h^\top}{\sqrt{d_k}} + M\right)V_h,$$3

and treats each chunk as an independent unit for all-to-all communication, local attention, and reverse all-to-all, without changing the attention equations (Yuan et al., 2 Jul 2026).

3. Efficiency mechanisms and complexity regimes

A central motivation for HPSA is that standard multi-head attention repeats a dense $$\text{Attn}_h(Q_h, K_h, V_h)=\mathrm{softmax}\!\left(\frac{Q_hK_h^\top}{\sqrt{d_k}} + M\right)V_h,$$4 computation across $$\text{Attn}_h(Q_h, K_h, V_h)=\mathrm{softmax}\!\left(\frac{Q_hK_h^\top}{\sqrt{d_k}} + M\right)V_h,$$5 heads. Different HPSA instantiations reduce cost by partitioning the attention domain, compressing it, or restructuring execution.

SPAttention reduces the core attention computation from $$\text{Attn}_h(Q_h, K_h, V_h)=\mathrm{softmax}\!\left(\frac{Q_hK_h^\top}{\sqrt{d_k}} + M\right)V_h,$$6 to $$\text{Attn}_h(Q_h, K_h, V_h)=\mathrm{softmax}\!\left(\frac{Q_hK_h^\top}{\sqrt{d_k}} + M\right)V_h,$$7 by distributing the causal matrix across heads instead of recomputing it $$\text{Attn}_h(Q_h, K_h, V_h)=\mathrm{softmax}\!\left(\frac{Q_hK_h^\top}{\sqrt{d_k}} + M\right)V_h,$$8 times (Zhao et al., 12 Nov 2025). The paper characterizes this as transforming multi-head attention from $$\text{Attn}_h(Q_h, K_h, V_h)=\mathrm{softmax}\!\left(\frac{Q_hK_h^\top}{\sqrt{d_k}} + M\right)V_h,$$9 independent $O(H \cdot N^2 \cdot d_k)$0 computations into a single, collaborative $O(H \cdot N^2 \cdot d_k)$1 computation, reducing complexity by a factor of $O(H \cdot N^2 \cdot d_k)$2 (Zhao et al., 12 Nov 2025). Because the masks are contiguous and regular, the method is described as hardware-friendly and compatible with kernels such as FlashAttention or FlexAttention (Zhao et al., 12 Nov 2025).

iMHSA addresses a different cost center: naïve cross-head interaction on full attention matrices incurs $O(H \cdot N^2 \cdot d_k)$3 extra operations. Its decomposition moves interaction to $O(H \cdot N^2 \cdot d_k)$4 and $O(H \cdot N^2 \cdot d_k)$5 factors, yielding the summarized complexity

$O(H \cdot N^2 \cdot d_k)$6

with output computed in reverse order as

$O(H \cdot N^2 \cdot d_k)$7

so that no full $O(H \cdot N^2 \cdot d_k)$8 matrix is materialized (Kang et al., 2024). The same decomposition also changes memory growth from quadratic to linear in $O(H \cdot N^2 \cdot d_k)$9 when $$\bigcup_h \mathrm{region}_h = \text{full attention graph}, \qquad \mathrm{region}_h \cap \mathrm{region}_{h'} = \emptyset \ \text{or controlled}.$$0 (Kang et al., 2024).

GlobalMind reduces full-image global attention by redefining the sequence through GAS. Full 2D self-attention over $$\bigcup_h \mathrm{region}_h = \text{full attention graph}, \qquad \mathrm{region}_h \cap \mathrm{region}_{h'} = \emptyset \ \text{or controlled}.$$1 pixels has complexity $$\bigcup_h \mathrm{region}_h = \text{full attention graph}, \qquad \mathrm{region}_h \cap \mathrm{region}_{h'} = \emptyset \ \text{or controlled}.$$2, while row-wise GAS reduces the quadratic factor to $$\bigcup_h \mathrm{region}_h = \text{full attention graph}, \qquad \mathrm{region}_h \cap \mathrm{region}_{h'} = \emptyset \ \text{or controlled}.$$3 and column-wise GAS to $$\bigcup_h \mathrm{region}_h = \text{full attention graph}, \qquad \mathrm{region}_h \cap \mathrm{region}_{h'} = \emptyset \ \text{or controlled}.$$4 (Hu et al., 2023). The paper states that the sequence length becomes only $$\bigcup_h \mathrm{region}_h = \text{full attention graph}, \qquad \mathrm{region}_h \cap \mathrm{region}_{h'} = \emptyset \ \text{or controlled}.$$5, not $$\bigcup_h \mathrm{region}_h = \text{full attention graph}, \qquad \mathrm{region}_h \cap \mathrm{region}_{h'} = \emptyset \ \text{or controlled}.$$6, which makes global attention feasible on large hyperspectral images (Hu et al., 2023).

MHPA replaces dense token–token attention with LSH-based partitioning and intra-/inter-partition aggregation. The resulting per-layer complexity is $$\bigcup_h \mathrm{region}_h = \text{full attention graph}, \qquad \mathrm{region}_h \cap \mathrm{region}_{h'} = \emptyset \ \text{or controlled}.$$7, and with channel-to-spatial downsampling by factor $$\bigcup_h \mathrm{region}_h = \text{full attention graph}, \qquad \mathrm{region}_h \cap \mathrm{region}_{h'} = \emptyset \ \text{or controlled}.$$8, it becomes $$\bigcup_h \mathrm{region}_h = \text{full attention graph}, \qquad \mathrm{region}_h \cap \mathrm{region}_{h'} = \emptyset \ \text{or controlled}.$$9 (Jiang et al., 2023). This is a distinct efficiency route: the partitioning is over tokens rather than heads, but it shows how partition structure can compress attention without explicit $(i,j)$0 matrices.

HCMS illustrates that head partitioning can be a systems primitive even when the mathematical attention rule is unchanged. Its time model is

$(i,j)$1

with optimal chunk count

$(i,j)$2

and an upper-bound speedup of $(i,j)$3 as chunk count increases and overhead vanishes (Yuan et al., 2 Jul 2026). The method is recommended when the communication ratio satisfies $(i,j)$4 (Yuan et al., 2 Jul 2026).

These results show that HPSA efficiency does not depend on a single algorithmic trick. The cited work uses at least four distinct routes: direct head-wise graph partitioning, low-rank cross-head factorization, axis redefinition, and runtime head chunking.

4. Empirical behavior: specialization, redundancy, and performance

The empirical record across NLP, vision, hyperspectral imaging, and long-context language modeling is consistent on one point: heads are not equally useful, and explicit structuring often improves the quality–efficiency trade-off.

SPAttention reports that head diversity is near zero under standard attention but increases by ~300× under its distance-band partitioning; attention entropy per head is lower, with an approx. 20% reduction (Zhao et al., 12 Nov 2025). In controlled experiments with 8 heads, 8 layers, $(i,j)$5, and $(i,j)$6, its average performance across HellaSwag, Winogrande, COPA, and STEM is $(i,j)$7 at $(i,j)$8 tok/s, compared with $(i,j)$9 at $N$0 tok/s for standard dense attention (Zhao et al., 12 Nov 2025). On OLMoE 0.25B–1.75B at sequence length 1024, the dense baseline average is $N$1 and SPAttention reaches $N$2; on OLMoE 1B–7B at sequence length 4096, the dense baseline average is $N$3 and SPAttention reaches $N$4, with reported ≈2× speedup in training throughput (Zhao et al., 12 Nov 2025).

iMHSA provides complementary evidence that structured head interaction can be beneficial when made affordable. On ImageNet-1K with ViT-T/16, softmax attention yields $N$5 top-1 at $N$6 GFLOPs, while iMHSA reaches $N$7 at $N$8 GFLOPs; on ViT-S/16, softmax reaches $N$9 and iMHSA $H$0; on ViT-T/8 with 784 tokens, softmax reaches $H$1 at $H$2 GFLOPs and iMHSA $H$3 at $H$4 GFLOPs (Kang et al., 2024). Its ablation also isolates the computational role of decomposition: interaction only, without decomposition, increases runtime from $H$5 ms to $H$6 ms and FLOPs from $H$7G to $H$8G in the ViT-Tiny experiment, while decomposition only lowers runtime but degrades accuracy (Kang et al., 2024). The paper further reports that interaction increases variance and decreases similarity between heads across depth, indicating enhanced head diversity (Kang et al., 2024).

Role-guided head partitioning improves both classification and translation. Across six text classification datasets, the guided Transformer matches or exceeds the best baseline on all datasets except MR, where it ties PSAN; for example, it reaches $H$9 on CR, $$W_{\text{base}}=\left\lfloor \frac{N}{H}\right\rfloor,\quad R=N\bmod H,\quad W_h=W_{\text{base}}+\mathbb{I}(h<R),\quad S_h=h\cdot W_{\text{base}}+\min(h,R),$$0 on TREC, and $$W_{\text{base}}=\left\lfloor \frac{N}{H}\right\rfloor,\quad R=N\bmod H,\quad W_h=W_{\text{base}}+\mathbb{I}(h<R),\quad S_h=h\cdot W_{\text{base}}+\min(h,R),$$1 on SST (Wang et al., 2020). On WMT 2016 English–German, the role-guided Transformer reaches $$W_{\text{base}}=\left\lfloor \frac{N}{H}\right\rfloor,\quad R=N\bmod H,\quad W_h=W_{\text{base}}+\mathbb{I}(h<R),\quad S_h=h\cdot W_{\text{base}}+\min(h,R),$$2 BLEU versus $$W_{\text{base}}=\left\lfloor \frac{N}{H}\right\rfloor,\quad R=N\bmod H,\quad W_h=W_{\text{base}}+\mathbb{I}(h<R),\quad S_h=h\cdot W_{\text{base}}+\min(h,R),$$3 for the baseline Transformer (Wang et al., 2020). Role ablations reduce performance in every case, with MajRel producing the largest average drop (Wang et al., 2020).

GlobalMind’s ablations show that both spatial and temporal interactive modules matter. Across five hyperspectral datasets, the baseline without GlobalM and GlobalD is improved by adding GlobalM, improved further by adding GlobalD, and the full GlobalMind achieves the best performance in almost all metrics (Hu et al., 2023). The visual analysis described in the paper reports fewer false positives, better separation of change versus unchanged regions, and more complete large change regions when both modules are present (Hu et al., 2023).

The pruning literature supplies the inverse perspective: if HPSA is absent, a large fraction of heads can be redundant. In encoder self-attention for EN–RU translation, pruning 38 out of 48 encoder heads causes only a $$W_{\text{base}}=\left\lfloor \frac{N}{H}\right\rfloor,\quad R=N\bmod H,\quad W_h=W_{\text{base}}+\mathbb{I}(h<R),\quad S_h=h\cdot W_{\text{base}}+\min(h,R),$$4 BLEU drop on WMT, and pruning 44 of 48 heads causes only ~$$W_{\text{base}}=\left\lfloor \frac{N}{H}\right\rfloor,\quad R=N\bmod H,\quad W_h=W_{\text{base}}+\mathbb{I}(h<R),\quad S_h=h\cdot W_{\text{base}}+\min(h,R),$$5 BLEU drop on OpenSubtitles (Voita et al., 2019). The heads that survive last are the specialized positional, syntactic, and rare-word heads (Voita et al., 2019). This is not itself an HPSA layer, but it provides direct empirical support for the claim that structuring heads by role is meaningful rather than merely decorative.

5. Analytical and theoretical perspectives

Recent theory formalizes head populations as a distribution over head parameters rather than as a fixed finite list. In a simplified single-layer causal multi-head self-attention model, the empirical law of $$W_{\text{base}}=\left\lfloor \frac{N}{H}\right\rfloor,\quad R=N\bmod H,\quad W_h=W_{\text{base}}+\mathbb{I}(h<R),\quad S_h=h\cdot W_{\text{base}}+\min(h,R),$$6 heads is

$$W_{\text{base}}=\left\lfloor \frac{N}{H}\right\rfloor,\quad R=N\bmod H,\quad W_h=W_{\text{base}}+\mathbb{I}(h<R),\quad S_h=h\cdot W_{\text{base}}+\min(h,R),$$7

and the multi-head logits are the empirical average of single-head features (Huan et al., 9 Jun 2026). In the infinite-head limit, the state variable becomes a probability measure $$W_{\text{base}}=\left\lfloor \frac{N}{H}\right\rfloor,\quad R=N\bmod H,\quad W_h=W_{\text{base}}+\mathbb{I}(h<R),\quad S_h=h\cdot W_{\text{base}}+\min(h,R),$$8, with mean-field logits

$$W_{\text{base}}=\left\lfloor \frac{N}{H}\right\rfloor,\quad R=N\bmod H,\quad W_h=W_{\text{base}}+\mathbb{I}(h<R),\quad S_h=h\cdot W_{\text{base}}+\min(h,R),$$9

and cross-entropy risk

$h$0

Training then follows a nonlinear Wasserstein gradient-flow equation

$h$1

with heads interacting only through the shared softmax residual field (Huan et al., 9 Jun 2026).

This mean-field view provides a rigorous interpretation of head partitioning as a distribution of capacity across parameter space. Global minimizers satisfy the support condition

$h$2

so optimal solutions may place mass on one point, on several discrete points, or on a manifold of minimizers (Huan et al., 9 Jun 2026). The theory therefore distinguishes head collapse from head specialization. Dirac stationary measures correspond to complete collapse onto a single prototype head; these can be exponentially stable under a Wasserstein strong-monotonicity condition, but they can also be translation-unstable, in which case the flow pushes mass away from the single-head configuration (Huan et al., 9 Jun 2026).

Finite-head analysis complements this continuum picture. The same paper proves a static approximation bound of order $h$3 between optimal finite-head risk and optimal mean-field risk, and a propagation-of-chaos result comparing finite-head SGD to the limiting PDE over finite time horizons (Huan et al., 9 Jun 2026). A plausible implication is that large-head HPSA architectures can be analyzed as particle approximations to a measure-valued optimization problem, rather than only as discrete architectural heuristics.

The pruning literature gives a more operational analytic lens. By attaching stochastic Hard Concrete gates $h$4 to each head and optimizing

$h$5

it partitions heads into active and inactive subsets during training (Voita et al., 2019). Although this is a sparsification framework rather than a partitioned attention graph, it shows that head-level structure can be studied through explicit head variables rather than solely through aggregate layer behavior.

6. Limitations, misconceptions, and research directions

A recurrent misconception is that head partitioning necessarily removes information. SPAttention explicitly argues the opposite for its own construction: no causal connection is globally dropped, because every pair $h$6 with $h$7 is preserved in at least one head (Zhao et al., 12 Nov 2025). The trade-off is architectural rigidity. Its partition is fixed, distance-based, and hyperparameter-free, which means a head assigned to long-range distances cannot attend to short-range positions at all (Zhao et al., 12 Nov 2025).

A second misconception is that any partitioned attention mechanism is literal head partitioning. The literature is more nuanced. iMHSA is HPSA-style because it structures cross-head interaction through low-dimensional factors, but it does not explicitly assign named head groups (Kang et al., 2024). GlobalMind partitions spatial and temporal processing largely at the module level through GAS, not by defining disjoint head groups inside a single attention operator (Hu et al., 2023). MHPA partitions spatial tokens and shares those partitions across heads, so it is partition-wise attention rather than literal head-wise partitioning (Jiang et al., 2023). HCMS partitions heads only as a scheduling unit and leaves the attention math unchanged (Yuan et al., 2 Jul 2026).

The design space also contains hard constraints with clear failure modes. Role-guided masks are fixed, deterministic, and parser-dependent; if the parse is wrong or the predefined role is unhelpful, a constrained head can become partially wasted (Wang et al., 2020). In iMHSA, interaction without decomposition is computationally prohibitive, while decomposition without interaction degrades accuracy (Kang et al., 2024). GlobalMind’s GAS assumes a favorable axial decomposition and may require different row/column choices for different aspect ratios (Hu et al., 2023). HCMS yields limited benefit when $h$8, and too many chunks increase overhead through event synchronization and kernel launches (Yuan et al., 2 Jul 2026).

The research directions proposed in the cited work converge on a common agenda. SPAttention suggests learned or dynamic partitioning, multi-axis partitioning, content-aware HPSA, hybrid HPSA plus global heads, layer-wise variation, and soft exclusivity (Zhao et al., 12 Nov 2025). iMHSA identifies better landmark selection and optimization of head count as open directions (Kang et al., 2024). GlobalMind suggests that axial choices can be adapted to aspect ratio and target structure (Hu et al., 2023). The mean-field analysis points toward multi-layer extensions, non-causal or bidirectional attention, relaxed compactness assumptions, and sharper convergence criteria for realistic Transformer parameterizations (Huan et al., 9 Jun 2026).

In aggregate, these directions indicate that HPSA is evolving along three coupled fronts: algorithmic partitioning of the attention graph, representational coordination or separation among heads, and systems-level partitioning of head execution. The unifying premise is that multi-head attention need not be a collection of $h$9 interchangeable dense operators over the same domain. Instead, heads can be assigned differentiated regions, factors, roles, or schedules, with the goal of preserving coverage while reducing redundancy and sharpening specialization.