Multi-Head Selective State-Space Models

• MHSSM is a neural architecture that integrates state-space models with multi-head modularization and selective gating to capture diverse dynamic patterns.
• It employs parallel SSMs with adaptive softmax gating to efficiently mix spatial, temporal, and channel information across tasks.
• MHSSM demonstrates state-of-the-art performance in imaging, language modeling, and point cloud analysis while reducing computational cost and enhancing trainability.

A Multi-Head Selective State-Space Model (MHSSM) is a state-space neural architecture that couples the linear dynamical modeling capabilities of state-space models (SSMs) with the architectural advantages of multi-head modularization and selective gating mechanisms. This hybridization matches the expressive and trainable subspaces of multi-head attention, yet leverages the inductive bias and linear-time efficiency of SSMs. MHSSM has emerged as a versatile backbone in fields such as vision, speech, medical imaging, point cloud learning, and large language modeling.

1. Formalism and Core Architectural Principles

MHSSM generalizes the standard SSM, where a sequence $X = [x_1, ..., x_n]$ is mapped to an output $Y = [y_1, ..., y_n]$ through an input-dependent interaction operator $W_{ij}(X)$:

$y_i = \sum_{j=1}^n W_{ij}(X)x_j.$

In MHSSM, $W_{ij}$ is decomposed as a sum across $H$ heads, each contributing a structured channel or direction:

$W_{ij}(X) = \sum_{h=1}^H \alpha_{ij}^{(h)}(X) H_h,$

where $H_h$ is a fixed or learned basis (often parameterized as the impulse response of an SSM), and $\alpha_{ij}^{(h)}(X)$ is a data-driven coefficient (e.g., softmax- or MLP-gated, or convolutional in sequence distance). Each head runs an independent SSM (possibly with input- or position-dependent parameters), and the outputs are weighted and fused, thus enabling the model to attend to multiple dynamic modes or spatial/temporal directions in parallel (Ghodsi, 17 Dec 2025, Safari et al., 22 Dec 2025).

2. Canonical Instantiations and Variants

Attention-SSM Unification and Head-Count Theorem

The unifying framework of (Ghodsi, 17 Dec 2025) demonstrates that:

• A single-head factorized (attention-style) model can only approximate a rank-1 span in the space of linear maps;
• To fully represent a k-dimensional SSM kernel family (i.e., $$k = \text{rank}_{int} = \dim \text{span}\{W(\tau)\}_{\tau}$$), at least $H = k$ heads are necessary and sufficient. Thus, MHSSM combines H parallel SSMs with selective weighting to achieve expressivity equivalent to multi-head attention, but grounded in dynamical system composition.

Multi-Directional and Selective Head Routing

In computer vision and spatial tasks, MHSSMs are architected with heads scanning in multiple directions: horizontal, vertical, and diagonal (and variants, such as anti-diagonal and reversals for "all-around" context). For example, in MRI super-resolution (Safari et al., 22 Dec 2025) and image restoration (Lin et al., 27 Jun 2025), per-head SSMs operate on 1D serializations defined by spatial scan patterns. Channel grouping is used to ensure constant parameter and computational cost as the number of scan patterns grows, with fusion performed via concatenation and a linear projection.

Selection of per-sample or per-patch importance weights for each head is achieved by a softmax gating mechanism applied to the concatenation of global SSM outputs, enabling adaptive, patch-wise routing, as detailed in plant counting (He et al., 2024).

Mixture-of-Experts and MoSE Emulation

MHSSM can also be realized via a sparse mixture-of-experts (MoE) paradigm, as in MossNet (Tuli et al., 30 Oct 2025): for each token or position, a gating network softly (or sparsely) selects among multiple SSM experts, each parameterizing independent (or low-rank-shared) state evolution. Analytically, this recovers the machinery of multi-head attention at the sequence operator level (query-key-value structure) in the limit of linear SSM flows.

Hybrid Selective Token/Channel Mixing

MHSSM generalizes beyond spatial direction and MoE to include mutually orthogonal parametrizations, e.g., mixing along tokens and channels via bidirectional SSM recurrences per head (Behrouz et al., 2024), with head-wise gating, projection, and fusion steps.

3. Implementation Details and Computational Properties

The core computational pattern in MHSSM is the parallel scan/recurrent update:

$$h_t^{(h)} = \bar{A}^{(h)}_t h_{t-1}^{(h)} + \bar{B}^{(h)}_t x_t^{(h)}; \quad y_t^{(h)} = C^{(h)} h_t^{(h)} + D^{(h)} x_t^{(h)}$$

where $\bar{A}_t$, $\bar{B}_t$ may be input-dependent (selective S6/Mamba style), and $x_t^{(h)}$ denotes per-head input embeddings. Head outputs are typically concatenated along the channel dimension and projected via a learned linear map.

• The total computational complexity is $O(H T N)$ for $T$ tokens/positions, $H$ heads, and SSM state-dim $N$ per head.
• By grouping (e.g., splitting $\lambda C$ channels into $n$ heads of $C' = (\lambda C)/n$), total FLOPs and parameters remain invariant as the head count (and thus scan patterns) increases (Lin et al., 27 Jun 2025, Safari et al., 22 Dec 2025).
• Adaptive gating (softmax or hard top-k, as in MoE models) incurs negligible overhead.
• Associative parallel scan algorithms are leveraged for efficient GPU implementation (Behrouz et al., 2024).

Pseudocode is standardized across visual and sequence tasks; for each head, input is projected, optionally depthwise convolved (injecting locality), SSM-scanned along a head-specific sequence, fused via gating or attention-style routing, and projected into the final output representation.

Table 1: Comparison of MHSSM Variants in Recent Literature

Paper	Input Dim / Domain	Head Scan/Type	Fusion/Routing
(Safari et al., 22 Dec 2025)	MRI/Patches (H×W×C)	HV, D 2D-scans	Concat + Linear, no gating
(He et al., 2024)	Images (counting)	H, V, D, A	Patchwise softmax gate
(Lin et al., 27 Jun 2025)	Images (restoration)	All-around (HVDA, rev)	Channel grouping, gating branch
(Tuli et al., 30 Oct 2025)	Text/Lang Modeling	MoE SSMs ("Temporal heads")	Top-k router softmax (Switch-style)
(Behrouz et al., 2024)	Tokens/Channels	Token/channel select	Weighted averages per layer
(Qu et al., 26 Jul 2025)	Point cloud (3D)	Shuffle-serialized	Concat + Linear

4. Empirical Outcomes and Task-Specific Performance

MHSSMs enable state-of-the-art or highly competitive performance in diverse domains with significantly improved efficiency compared to Transformer-style attention or dense SSM baselines.

• MRI Super-Resolution (Safari et al., 22 Dec 2025): MHSSM yields SSIM = 0.951, PSNR = 26.9 dB (brain), outperforming GAN, Transformer, SSM, and diffusion baselines with only 0.9M params, 57 GFLOPs—a 99.8%/97.5% reduction in parameters/computation vs. Res-SRDiff.
• Image Restoration (Lin et al., 27 Jun 2025): EAMamba (MHSSM) achieves 31–89% reduction in FLOPs versus 2DSS baselines, maintaining or improving PSNR/SSIM.
• Plant Counting (He et al., 2024): Adaptive multi-head selective SSM yields best MAE (4.6), outperforming both single-direction and non-adaptive multi-head variants.
• 3D Point Cloud Classification (Qu et al., 26 Jul 2025): Increasing SSM heads systematically improves ModelNet40 accuracy (peak 93.96% at 6 heads). Further increase yields diminishing returns.
• Language Modeling (Tuli et al., 30 Oct 2025): MossNet (MHSSMoE) consistently surpasses SSM and Transformer models at fixed parameter budget; scaling expert count increases perplexity and downstream accuracy gains.
• Speech Recognition (Fathullah et al., 2023): Replacing Transformer self-attention with MH-SSM improves LibriSpeech WER by 0.4–0.5pp absolute, and adding attention block (“Stateformer”) matches or outperforms Conformer.

5. Selective Routing, Specialization, and Theoretical Foundations

Selective gating mechanisms in MHSSM extend standard multi-head modularity to allow for dynamic, sample- or patch-dependent information routing:

• Per-patch softmax gates distribute head importance, encouraging specialization without explicit regularization (He et al., 2024).
• Switch-style load balancing (auxiliary loss) is used for expert utilization in large-scale sparse routing (Tuli et al., 30 Oct 2025).
• Adaptive merging across layers and directions enables the model to attend preferentially to the most informative axis or temporal context per token, supported by ablation showing intermediate configurations are suboptimal.

The (Ghodsi, 17 Dec 2025) Equivalence (Head-Count) Theorem formalizes the necessity and sufficiency of $H = k$ heads for representing a k-dimensional SSM kernel subspace. Furthermore, the Gradient Highway Result demonstrates that whereas stable SSM Jacobians decay exponentially with sequence distance, attention-style mixing (including MHSSM gating) offers direct input-output gradient pathways, improving trainability.

A plausible implication is that MHSSM architectures can be tuned for both algebraic expressivity (operator span, by increasing heads) and efficient gradient propagation (via selective gating), closing the gap between attention mechanisms and dynamical SSMs in both theory and practice.

6. Implementation, Hyperparameters, and Limitations

Key hyperparameters include number of heads $H$, SSM state dimension per head $N$, channel group size, and scan patterns (directional, all-around, or task-specialized). Empirical guidance from point cloud and vision benchmarks suggests moderate $H$ (e.g., 4–6) is often optimal, balancing diversity and channel fragmentation (Qu et al., 26 Jul 2025, Lin et al., 27 Jun 2025).

MHSSM blocks are modular components placed within UNet backbones, Transformer-like encoder-decoders, or hierarchical stacks, always retaining linear or quasi-linear time/space complexity in input size.

Practical limitations include:

• For very high $H$, over-fragmentation of channel space can degrade accuracy.
• Complexity-control via gating and grouping is required to avoid linear growth of computational cost with the number of scan patterns or experts.
• Soft selection can, in principle, be replaced with hard routing (MoE), but trade-offs in sparsity and hardware efficiency persist.
• For tasks with extreme long-range dependencies, explicit attention blocks or auxiliary mechanisms may still be beneficial (Fathullah et al., 2023).

7. Outlook and Research Directions

MHSSM advances a unified paradigm that interpolates between SSM-derived recurrence and attention-derived modularity, with proven expressivity-complexity and gradient propagation trade-offs (Ghodsi, 17 Dec 2025). Open research avenues include:

• Investigation of combinatorial scan pattern selection for arbitrary spatial/temporal structures,
• Dynamic or learned head/basis adaptation per layer,
• Deeper integration with MoE paradigms and scaling laws in ultra-large foundation models,
• Theoretical characterization of selective dynamics and convergence when the gating network and SSM parameters are jointly optimized.

As empirical and theoretical progress continues, MHSSM is likely to serve as a canonical design point for efficient, expressive sequence and spatial modeling in both specialized domains and general-purpose architectures.