Group-Aware SSM Elastification

• Group-aware SSM elastification is a technique that partitions SSM heads and state channels into groups to maintain structured broadcasting and robust long-range dependency modeling.
• It employs importance ranking and group-consistent masking to prune only within predefined groups, ensuring stable inference and effective model compression.
• The approach underpins hybrid model compression pipelines and enables nested submodel extraction, achieving high performance retention, faster inference, and significant training token savings.

Group-aware State Space Model (SSM) elastification refers to a class of structured width adaptation and pruning techniques for modern sequential models, particularly those employing Mamba and Mamba2 SSM blocks. These methods strategically partition SSM heads and state channels into predefined groups and enforce masking or pruning only within those groupings, thereby preserving broadcast, convolutional, and recurrence semantics critical to long-range dependency modeling and stable inference. Group-aware SSM elastification has emerged as a crucial building block in model compression pipelines for hybrid LLMs, as well as in the construction of “many-in-one” nested models enabling the extraction of multiple submodels from a shared set of weights (Taghibakhshi et al., 15 Apr 2025, Taghibakhshi et al., 20 Nov 2025, Karuvally et al., 7 Jul 2025).

1. Theoretical Motivation and Necessity

In selective SSMs such as Mamba and Mamba2, the timestep update operates as

$$h_t = A \cdot h_{t-1} + B \cdot x_t,\qquad y_t = C^\top \cdot h_t + D \cdot x_t$$

with head and group-specific parameterization for projections $B$, $C$, $A$, and $D$ (Taghibakhshi et al., 15 Apr 2025). The architectural design mandates that input $x_t$ and associated projections are broadcast groupwise over $g$ groups of $m_h$ heads (each with dimensionality $m_d$). Arbitrary pruning or reordering across group boundaries disrupts this structured broadcasting, leading to state propagation corruption and destroyed long-range information flow. Standard elastification—masking heads or neurons independently—causes catastrophic accuracy losses due to such violations. Group-aware SSM elastification is formulated to maintain invariance under group partitions, guaranteeing the model’s computations and memory structures remain valid at all sizes.

2. Formal Definition and Algorithmic Framework

2.1 Group Structures and Pruning Constraints

The $m_h$ heads in each Mamba layer are partitioned into $G$ disjoint groups $\mathcal{G}_1,\ldots,\mathcal{G}_G$ each containing $m_h/G$ heads. Pruning permutations $\mathcal{P}$ are restricted: they can only modify head orderings within each $\mathcal{G}_g$, never across groups (Eq. 8 in (Taghibakhshi et al., 15 Apr 2025)). State-channel groups are similarly blocked; projections such as $W_B$, $W_C$ stack $G$ state-segments of dimension $d_s$ each, forming contiguous “state-channel groups” (Taghibakhshi et al., 20 Nov 2025).

2.2 Importance Ranking and Mask Construction

Structured pruning is driven by head and channel importance scores, often computed via activations from the $W_x$ projection:

• Head-channel aggregation: $s_{d} = \| \sum_{b,l} s_{(b,l),d} \|_2$
• Within-group head selection: $f_h = \| s_{h,\mathcal{D}_{\text{top}}} \|_2$ over the $k_d$ most significant channels
• Mask application: Within each group $\mathcal{G}_g$, keep the $k_g$ heads with highest $f_h$, and only those channels in $\mathcal{D}_{\text{top}}$; other rows/columns are zeroed or pruned.

The resulting mask matrices enforce two constraints (Taghibakhshi et al., 20 Nov 2025):

1. Group-wise consistency: All heads within $\mathcal{G}_g$ share the same mask pattern at the channel level.
2. State-segment contiguity: Pruning of state channels proceeds only in contiguous blocks (entire state-segments).

2.3 Optimization Objective

Group-aware elastification is cast as:

$$\min_{\text{mask} \in \mathcal{P}} \; L_{\text{val}}(\theta \circ \text{mask}) + \lambda|\text{mask}|_0$$

subject to mask feasibility—only group-consistent pruning is permitted. In practice, priors are set via $k_g$ (per-group heads) and $k_d$ (channels), and optimized by forward-score ranking (Taghibakhshi et al., 15 Apr 2025).

3. Integration Into Compression and Elastic Model Pipelines

3.1 Hybrid Architecture Compression (Nemotron-H)

In the Nemotron-H pipeline (Taghibakhshi et al., 15 Apr 2025), group-aware SSM pruning is interleaved with:

• FFN neuron and embedding channel pruning (using forward-pass scoring)
• Depth pruning by layer-wise KL-divergence
• Lightweight and extended knowledge distillation phases
• Architecture search over layer count, embedding width, FFN width, SSM head and channel counts

The compressed Nemotron-H 4B model preserves ≥96% of the parent 8B performance and achieves double the inference speed, with up to 40-fold reduction in required training tokens compared to independent training of baseline models.

3.2 Nested and Many-in-One Model Construction (Nemotron Elastic)

In the Nemotron Elastic framework (Taghibakhshi et al., 20 Nov 2025), group-aware SSM elastification enables embedding multiple nested submodels (e.g., 12B, 9B, 6B) inside a single weight checkpoint. Routers trained with Gumbel–Softmax gating dynamically select group-consistent head/channel masks per model budget. Zero-shot extraction of any submodel simply slices weights as dictated by the hard (argmax) router outputs; no additional tuning or repair is needed.

Empirically, this enables 7x token reduction and ~43% deployment memory savings compared to state-of-the-art alternatives, with equivalent or superior accuracy on reasoning benchmarks.

4. Mathematical and Implementation Details

The mathematics of group-aware masking require that the pruning operator can be expressed as:

$$m_{\mathrm{mamba}}[\phi(h,c)] = \mathbb 1[h \leq h^* \;\land\; c \leq c^*]$$

with group-wise consistency for all $h,h' \in \mathcal{G}_g$:

$$m_{\mathrm{mamba}}[\phi(h,c)] = m_{\mathrm{mamba}}[\phi(h',c)] \quad \forall c$$

Mask gating for state-segments proceeds along block-level indices. During training, routers generate mask indices by importance ranking (from normalized MSE-based, forward-sensitivity, or activation-based metrics), and the pipeline applies these to participating weight matrices: $W_x$, $W_z$, $W_B$, $W_C$, $W_{dt}$.

In the Nemotron Elastic system, these dynamic masks are jointly coordinated with FFN and embedding masking, and depth gating, within a two-stage curriculum—first with uniform, then with long-context-skewed budget sampling, all guided by distillation to a frozen full teacher model (Taghibakhshi et al., 20 Nov 2025).

5. Algebraic and Expressive Extensions

Beyond architecture compression, group-aware SSM elastification takes on an algebraic dimension via the Adaptive Unitary SSM (AUSSM) construction (Karuvally et al., 7 Jul 2025). Here, the “group-awareness” refers to the ability of input-adaptive, skew-symmetric state-transition matrices $A_t = f(u_t)$ to instantiate unitary group actions (cyclic, permutation, and solvable automata). AUSSM blocks interleaved with conventional Mamba blocks achieve recognition of any solvable regular language at finite precision—an extension unachievable for time-invariant, real-valued SSMs. The resulting recurrence is both group-aware (expressing structured cyclic group actions) and elastic (frequencies are smoothly modulated by the input), implemented in CUDA-efficient separable convolutional form for scalable training.

6. Empirical Effectiveness and Ablation

The empirical foundation for group-aware SSM elastification is robust:

• Nemotron-H (group-aware pruning): 4B model achieves ≥96% of 8B baseline, 2x inference speed, 40x token efficiency, outperforming Phi-4-Mini, Qwen-2.5, and Zamba-2 on benchmarks such as MMLU, GSM8K, and HumanEval (Taghibakhshi et al., 15 Apr 2025).
• Nemotron Elastic: Nested 12B/9B/6B submodels deliver reasoning performance within 0.1–0.5 points of larger individually trained models, using only 110B tokens across all submodels and incurring no runtime penalty for elastification (Taghibakhshi et al., 20 Nov 2025).

Ablations demonstrate that abandoning group constraints (naive head pruning) degrades LM loss by 30–50% pre-distillation, introducing instability and accuracy drop-off. Group-aware masking ensures both the correctness of SSM recurrence and numerically stable weight sharing across variants.

Approach	Accuracy Retention	Latency/Speedup	Training Cost Savings
Group-aware SSM elast.	≥96% parent model	2–3× faster	>40× (tokens)
Naive SSM pruning	≤70% (catastrophic)	N/A	<2×
Elastic Nested pipeline	≤0.5 pt baseline drop	Inherited	7× (tokens)

7. Significance and Role in Modern LLMs

Group-aware SSM elastification is now a foundational technique for producing highly compressed, high-accuracy hybrid LLMs and enabling practical deployment of multi-budget nested models. Its principled enforcement of group semantics under elastification is empirically essential for both performance and stability. Combined with end-to-end router-based curriculum training and knowledge distillation, it underpins the latest generation of reasoning LLMs optimized for efficiency, modularity, and deployment flexibility (Taghibakhshi et al., 15 Apr 2025, Taghibakhshi et al., 20 Nov 2025, Karuvally et al., 7 Jul 2025).