Efficient Hybrid Language Model Compression through Group-Aware SSM Pruning
Authors: Ali Taghibakhshi
Sharath Turuvekere Sreenivas
Saurav Muralidharan
Marcin Chochowski
Yashaswi Karnati
Raviraj Bhuminand Joshi
Ameya Sunil Mahabaleshwarkar
ZIJIA CHEN
Yoshi Suhara
Oluwatobi Olabiyi
Daniel Korzekwa
Mostofa Patwary
Mohammad Shoeybi
Jan Kautz
Bryan Catanzaro
Ashwath Aithal
Nima Tajbakhsh
Pavlo Molchanov
Venue: NeurIPS 2025 poster ()
Keywords: Compression
Pruning
SSMs
Hybrid LLM
Primary area: deep_learning
TL;DR: The paper introduces a pruning and distillation method for hybrid LLMs, compressing Nemotron-H 8B to 4B with better accuracy and ~2× faster inference, advancing the efficiency-accuracy trade-off.

Abstract: Hybrid language models that combine Attention and State Space Models (SSMs) have been shown to achieve state-of-the-art accuracy and runtime performance. Recent work has also demonstrated that applying pruning and distillation to Attention-only models yields smaller, more accurate models at a fraction of the training cost. In this work, we explore the effectiveness of compressing Hybrid architectures. To this end, we introduce a novel group-aware pruning method for Mamba layers that preserves the structural integrity of SSM blocks and their sequence modeling capabilities. We combine this method with FFN, embedding dimension, and layer pruning, along with knowledge distillation-based retraining to obtain a unified compression recipe for hybrid models. Using this recipe, we compress the Nemotron-H 8B Hybrid model down to 4B parameters with up to $40\times$ fewer training tokens compared to similarly-sized models. The resulting model surpasses the accuracy of similarly-sized models while achieving $\sim2\times$ faster inference throughput, significantly advancing the Pareto frontier.

Minitron-SSM: Efficient Hybrid Language Model Compression through Group-Aware SSM Pruning
Ali Taghibakhshi, Sharath Turuvekere Sreenivas, Saurav Muralidharan" Marcin Chochowski, Yashaswi Karnati, Raviraj Joshi, Ameya Sunil Mahabaleshwarkar Zijia Chen, Yoshi Suhara, Oluwatobi Olabiyi, Daniel Korzekwa
Mostofa Patwary, Mohammad Shoeybi, Jan Kautz, Bryan Catanzaro Ashwath Aithal, Nima Tajbakhsh, Pavlo Molchanov
NVIDIA
fataghibakhsh, sharatht, sauravm, mchochouski, ykarnati, ravirajj, ameyasunilm, zijiac, ysuhara, oolabiyi, dkorzekwa, mpatwary, mshoeybi, jkautz, bcatanzaro, aaithal, ntajbakhsh, pmolchanoryYnnvddi..cmm
Abstract
Hybrid language models that combine Attention and State Space Models (SSMs) have been shown achieve state-of-the-a ahe accuracy and runtime performance. Re- cent work has also demonstrated that applying pruning and distillation to Attention- only models yields smaller, more accurate models fraction of the training cost. In this work, we explore the effectiveness of compressing Hybrid architec- tures. To this end, we introduce novel group-aware pruning method for Mamba layers that preserves the structural integrity of SSM blocks and their sequence modeling capabilities. We combine this method with FFN, embedding dimension, and layer pruning, along with knowledge distillation-be unified compression recipe for hybrid models. Using this recipe, we compress the Nemotron-H H 8B Hybrid model down to 4B parameters with up to 40x fewer training tokens compared to similarly-sin
retraining to obtain
models. The resulting model surpasses the accuracy of similarly-sized models while achieving ~ 2x faster inference throughput, significantly advancing the Pareto frontier.
Models on Hugging Face
Nemotron H----------------- ABB-Base-89
Nemotron-Hi ---4 118-Instruct-11
Introduction
Recent advances in language modeling have led to the development of hybrid architectures that
combine Transformer layers [25) with State Space Models (SSMs) [T] 7 These hybrid models leverage the complementary strengths of both approaches: Transformers excel at capturing global dependencies through self-attention mechanisms, while SSMs provide efficient sequence processing with ONV) scaling during training and ((1) cache size during inference. Mamba 7 in particular is a popular SSM designed for efficient sequence modeling with linear-tima complexity and support for long contexts and is often the preferred choice for non-attention layers in hybrid architectures.
Equal contribution.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Accuracy Inference Speed Frontier
Accuracy Training Cest
Nemortinn 88
Nametrun-H--8
Net
Nuerige Accurnty (AleyyyAlggen (Algred mode))
20
Aecuracy (heee modell
Qwen2.5
Dalzard
68-
Zambs2
Phi-4-Mini-08
50
68-
2x faster, 677
Ours Compresed
0men-2.5-38
Others 3P Compressed
Compression Path
Trained from
62
1/2
0.1
5.0
12.0
15.0
18.0
Throughput
model (tiltion tokend)
Figure 1: Comparison of NEMOTRON-H 4B model accuracy w.r.t. inference throughput (left), and training budget for the base model (right) to similarly-sized community models. Inference throughput measured an input and output sequence length of 65536 and 1024, respectively.
However, despite their improved efficiency, many hybrid LLMs remain incredibly large, often spanning billions of parameters this motivates the need for efficiently creating smaller hybrid models suitable for deployment in resource-constrainee environments.
Model pruning- the removal of redundant parameters while preserving accuracy- has recently emerged as a promising approach for compressing LLMs. In particular, methods that combine structured pruning (i.e., pruning of entire parameter blocks such as neurons, attention heads, etc.) with knowledge distillation [13] have proven effective at simultaneously reducing model memory footprint while improving runtime performance and accuracy (20].
While pruning techniques have been extensively studied for Transformer architectures 20 4 24].
their application to hybrid models remains significantly underexplored. Indeed, the optimization of SSMs is an active area of research, with various methods being explored, including post-training quantization [6.
Regarding pruning specifically, some early work on Mamba and SSMs includes Mamba- -hreddr 19 which removes the entire state space module from the Mamba layers, leaving only linear projections and convolution layer. concurrent study, Ghattas et al. [ propose a method for pruning Mamba architectures by focusing on three aspects: state space dimension reduction, Mamba head dimension pruning, and Mamba head merging.
To the best of our knowledge, no existing work on SSM/Mamba pruning presents a holistic com- pression strategy that simultaneously combines various aspects of SSM pruning with the pruning of other network components such as FFN neurons, embedding channels, and network depth; we believe such an approach essential for obtaining the best combination of runtime performance and model accuracy. In this paper, we introduce a novel group-aware pruning method for Mamba layers that targets both Mamba heads and head channels. We combine this method with FFN, embedding dimension, and layer pruning to obtain unified compression recipe for hybrid models that maximizes their accuracy and runtime performance.
This paper makes the following key contributions:
Introduces a group-aware pruning method for Mamba layers that preserves SSM block structure and sequence modeling capabilities.
Presents novel hybrid pruning recipe that effectively combines Mamba pruning with the pruning of other network components such as FFN neurons, embedding channels and layers.
Presents findings the sensitivity of Mamba block components to pruning, along with accuracy-throghhpur trade-offs when combined with pruning of other network components.
Utilizes the proposed hybrid pruning recipe to compress the Nemotron-H 8B model parameters through pruning and knowledge distillation. As shown in Figure I the resulting

---

Teacher Model
Scowing Activations
Mombo Loyar
Emboddiing
Altuntti Fmo
Actunt Fon
CM0.CM
Knowledge Distillation
124
Vins
Attentiion
Chr.0 002
0h CAA.C11...00
Pruning
Banking Scores
Figure 2: Overview of pruning and distillation for hybrid architectures. Starting from pretrained LLM, we first evaluate the importance Mamba heads and channels, FFN neurons, and embedding channels. We then rank them, trim the least important neurons, and distill the knowledge from the original LLM to the pruned model. Attention layers are not pruned since they amount to only 8% of the total number of layers.
model requires up 40x fewer training tokens compared to others in the same size range. It also achieves state-of-the-ar accuracy on benchmarks, and a 2x speedup inference throughput compared to similarly-sized models, significantly pushing the Pareto frontier.
2 Pruning & Distillation Methodology
We start the pruning procedure by computing the importance or sensitivity each network component; namely, Mamba heads and head channels, FFN neurons, embedding channels, and layers. To keep this phase lightweight, we adopt purely activation- brased strategy (requiring only forward propagation passes) for computing importance scores, similar Minitron [20]. Once scores are computed, sort the corresponding network components in decreasing order of importance while following any additional implementation constraints (discussed in more detail later in Section 2.1). We then prune away the network components with the lowest scores. Finally, the pruned model is distilled using the teacher model obtaín the final pruned model. The full procedure is illustrated in Figure 2
2.1 Mamba Pruning
Background State Space Models (SSMs). SSMs are class of sequence models that process inputs through hidden states evolving over time (7]. The general form of an SSM is given by:
h Ah-1 Br,
(1)
Yt C. + Dr,
Here, h4 represents the hidden state, 24 the input, Ye the output, and A, B. C, and are parameter matrices. The above equations describe linear time-invariant (LTI) SSMs, where the parameters remain constant across timesteps. The Mamba architecture 7 introduced selective SSM variant with time-varying parameters:
h Ash1-1 B
y: - DDz,
This selective mechanism allows the model to adapt dynamically to the input sequence, improving performance complex tasks. Mamba2 [7] builds upon the selective SSM framework and introduces several enhancements to improve efficiency and scalability. leverages the Structured State Space Duality (SSD), which connects SSMs and attention mechanisms through semi-separable matrix representations. This duality enables Mamba2 to combine the linear efficiency of SSMs with hardware- hardware-friendly quadratic computations typical of attention models.

---

Broadcast
02
8 Mamba
groups
H1
H8
H4
Broadcast
Broadcast
14
H5
space
statia space
H5
Intra-Group Permutation
HB
dimension of
dimension of
H6
H7
H7
HB
HB
9 Mamba heads
Figure 3: Mamba group structure visualization showing broadcasting and original B, computation. Colors represent distinct entries. The Figure illustrates how only within-group head permutations can preserve SSM semantics. As a counter example, if H3 and H8 were to swapped, the resulting B,r, would NOT be any permutation of the original (no permutation) BrZt.
To better understand importance estimation and pruning Mamba layers, we first dive into the forward pass of Mamba layer. The Mamba layer processes input through five distinct projection matrices Wz, Wa, WB. Wc, and Wa., following layer normalization. Let us denote the input to the Mamba layer X. Then, these projections generate intermediate matrices H WB (B((L(()), = We (LN(())) and de = Wa, (X)) which have the following dimensions, respectively; W., Wa E Rd. ((m, xma); WB. W..We We E Rdxx(gxd,), and Wa, E km,, LN denotes layer normalization, d, is the model embedding dimension (a.k.a hidden dimension). g is the number of Mamba groups, dg is the SSM state dimension, m. the number of Mamba heads, and m, is the Mamba head channels. The matrices x, B, and C' undergo causal convolution, as in 3 conv convid(x:) B conv d B), Ö convld(C) Consequently, the obtained matrices participate in the selective state space model (SSM) updates, ğ = SSM(2, B, Ć, D, dd). Here, A, D E mh are SSM learnable parameters corresponding to state transition and direct feed through, respectively (see Equations 3 and 4 The SSM output fed into gated normalization layer, which is then followed by output projection, y (mma x.. is the output projection.
W. (LN(X)).
I = (L.)))
Wo (RMSNNrmm 2)), where Wo
Group-Aware Head Permutation Constraints. Pruning requires scoring, sorting, and trimming neurons or heads of each layer, as shown in Figure 2 The FFN and embedding activations are permutation equivariant, i.e. for a permutation operator FFN or embedding layer and activation A. and input X we have L(X) = A = (X) = P(A). However, Mamba layers and activations are not permutation equivariant. As shown Figure 3 the B,z, operation from Eq. B involves reshaping B into B xRa and broadcasting it across E (M This broadcasting creates group-specific interaction patterns that constrain our pruning approach. As result, permuting heads across groups would alter the B, broadcast pattern, violating Eq, 3s group-wise computation Therefore, when sorting Mamba heads using activation scores, heads belonging
as shown by Brr,+ + (BP (BP(())))' we must preserve Mamba's group structure. Let Gg C {1, (1..... m denote the to group g. Any permutation P of heads must satisfy P(h) Gg vh Gy. In other words, Mamba heads and activations are permutation equivariant only for the aforementioned permutation operators.
similar constraint for permuting Mamba head channels ap- plies. For head channel pruning, maintain consistency across all heads through shared ranking. The state tensor E Rm,, xma xm, requires channel-wise permutations Pa to satisfy Pa(hi,j.)) Pa(hy j.k) Vi, i' (1, mx). meaning each channel index k either preserved or pruned uniformly across all heads.
Head Channel Consistency.
We factor out the sequence length and batch size simplify our description; the analysis remains valid without them.

---

Scoring and Ranking Methodology. The Mamba head and head channel ranking follows nested scoring procedure:
1. Head Channel Scoring: For each head channel d E .1..., m.)) we compute aggregate importance and the aggregation is over I., the sequence length, and B, the batch size. Aggregation metric used along and B dimensions mean and L2, respectively, following Minitron [20]. E I (m(m))maa contains raw activation scores, and s,, denotes the d-th column across all heads. We then select the top-ka channels, Diop topk (sd, ka).
scores sd Σ B. S.d where (X)
de(l..)ma)
2. Head Scoring: Using the pruned channels Dup. compute head importance scores, ʃn. = (hp ll2 vh E (1....,m)..
Group-Constraieed Ranking: Within each Mamba group gg. sort heads by their scores, Rg (f.)
NEC,
The final head ranking R is the concatenation of group-wise rankings, R = Ro [1: kgl, where
kg is the target head count per group and f denotes ordered concatenation.
The following algorithm provides concise walkthrough of Mamba head and head channel ranking:
Require: Activation scores
E Xma, target channels has target heads per group
[
Ensure: Head ranking R. channel ranking Dup
1: Compute channel scores: Su f [|s. AA
2: Diop f- top-ka indices of [sa]
3: Compute head scores: fr. + D..l.. Ah
4: for f to G do
5: Rg f- argsort-descendingg (/a Gg))
6: Ri f- first k, elements of Rg
7: end for
8: R + D Rya
After obtaining the Mamba heads and head channel neurons to keep, we trim the corresponding matrices, W <- W[R), for W E (Wz, W., Wo, WA, WD, Wa, convld).
2.2 Pruning FFN, Embedding, and Depth
Similar the Minitron approach for structured pruning of Transformers 20 activations produced by the FFN and LayerNorm layers to determine which neurons and embedding channels contribute least to the model's performance. Formally, for the i-th neuron feed-forward
we examine the
F(t) layer, we compute its importance score as: neuron ΣB.L X(Wi)T, where Wi refers to the i-th
row of the weight matrix W1 in the first linear projection of the FFN, is the input to the FFN layer, and ΣB,L denotes aggregation along the batch and sequence dimensions. Similarly, for the
i-th embedding channel, we compute: f-(1) ER., LN(X); where LN(X); represents the i-th dimension of the layer-normalized input. The embedding channel scores computed across all layers that utilize the embedding channel, including FFN, Mamba and Attention projection layers, and LayerNorm components. We mean and 1.2-norm aggregation along B and L. dimensions, respectively. After computing these scores, we sort them in descending order and keep the top-k neurons and embedding channels based the target compression ratio, pruning those with the lowest importance scores.
We analyze layer importance using Kullback-Leible divergence (KLD) between logits from a model with specific layer removed and the full model. This importance estimation small random subset of 256 samples to account for sample variability. The averaged layer(s) with lowest importance scores are subsequently removed. We provide a detailed analysis of importance distribution for the Nemotron-H 8B model in Section A.2
Depth Pruning.
FLAP Importance. FLAAP 2 a retraining-free structured pruning technique designed to measure the recoverability of a model's output feature map upon removing specific columns from weight
5

---

matrices. While we explored FLAP as an alternative importance metric. we observed that it didn't t offer any clear advantage over activation based importance (see Section A. for more details).
2.3 Architecture Search
Our compression strategy explores multiple axes within the 4B parameter budget through combina- torial pruning. Our search space includes depth reduction (removing 4-26 layers from the original 52-layer architecture) combined with width pruning of embedding channels (3072-4096), FFN di- mension (9984-2) (254),,504) Mamba heads (64- 128), and Mamba head channels (32-64). This multi-axis search space includes 125 candidate architectures meeting the parameter constraints.
Our search procedure follows these steps: (I) compute the zero-shot validation loss for all candidates on 1024 calibration samples, (2) select the top K architectures (22 in this paper) with the best loss values and perform lightweight knowledge distillation (KD) on them with 3.8B tokens, using the original 8B model as the teacher, and (3) select the top architecture candidate from step (2), using throughput and latency measurements for breaking ties, and perform extended knowledge distillation with ~ 380 tokens to obtain the final model (see Table 2 We note that step (2) is critical for getting a reliable ranking of architectural candidates, as also noted in prior work 20
2.4 Accuracy Recovery with Knowledge Distillation (KD)
To recover the accuracy lost due to pruning, the model undergoes continued training. Recent work has demonstrated that distilling knowledge [13] from the original model the pruned model outperforms conventional fine-tuning [23] 4] we thus adopt logit-based distillation for continued training, employing forward KL divergence (FKLD) loss exclusively during the accuracy recovery phase, following Minitron 20
Experiments and Results
To identify the optimal compression strategy for hybrid models, we conduct several ablation stud- ies evaluating the impact of pruning different components on accuracy and runtime performance. Our experiments, summarized in Section [3.]] reveal key insights and highlight differences from Transformer-only compression We then describe our main results for the NEMOTRON- -H 4B model
in Section 3.2
3.1 Analysis and Ablations
Depth vs Width Pruning. As shown in Table width-only pruning (#1) significantly outperforms depth-only pruning (#24) at 50% compression ratio (8B to 4B). Notably, a cepth--runed model with 36 layers (#25), despite having ~1.44x more parameters performs worse than the least accurate width-only pruned 4B candidate (#23, with 64 Mamba heads), demonstrating the critical role of depth in maintaining accuracy as also observed with Transformer-onj, models.
Finding 1: Width pruning considerably outperforms depth pruning in terms of accuracy.
Impact on Inference Throughput. Table shows that depth-only pruning (#24) provides the highest inference speedups. F (Appendix) ( presents the correlation between pruning various network components and performance metrics such as throughput, latency, and LM-loss for fixed 4B parameter count. We notice from the Figure that pruning Mamba components results in faster models compared to pruning FFN and embedding dimensions. Furthermore, we also compare the effects of pruning Mamba heads pruning head channels Figure [ (Appendix): we observe that the former yields better speed improvements than the latter within a given Mamba layer.
Finding 2: Impact throughput: Depth pruning Mamba FFN and Embedding.
'Note that due to space restrictions, we provide most of the Tables/Figures related to ablations in the Appendix.

---

Layers
Channel
Relutive Throughput
3072
12228
112
1,380
3328
9984
112
1.384
5472
12388
112
1. 1,388
1.02
3072
13056
112
1.389
1.04
3072
14592
128
1.393
1.11
10752
1.394
52
5072
11520
112
64
1 1.396
1.02
3072
13056
1.396
1.04
1.396
12
52
M072
12228
103
1.397
1403
5072
13056
1.397
1.03
13056
15
S2
3072
14592
96
56
1.398
1.05
3072
12288
128
1.398
1.08
9984
1.399
1.07
18
S2
3072
13824
96
58
1.401
1.05
3072
11520
128
1,402
11648
1,002
21
48
3072
13824
112
64
1.403
1.09
3428
1.403
1.08
3072
16128
1,411
1.07
24
26
4096
215504
128
64
1.533
1.31
4096
21504
1,400
1.12
8B parcel
52
4096
21504
128
0.74
Table 1: Model configurations with their corresponding LM validation loss after lightweight KD (sorted in increasing order), and relative inference throughput. Highlighted row shows the best (lowest) loss. All models have 4B parameters, except entries marked with *, which have more.
Impact on Accuracy. Table T shows that model depth (424) is most sensitive to accuracy, followed by Mamba heads (#23), while FFN and embedding dimensions have less impact. Further ablations isolating the pruning of Mamba heads and head channels show that pruning head channels leads to a greater accuracy loss (Figure B Given depth pruning's ' effect on inference speed, we explore a combined pruning strategy, starting with depth-only pruning followed by distillation to assess its limits. As shown in Figure we observe significant accuracy drops on math and coding benchmarks below 44 layers. We then apply width pruning both the 44- and 48-layer variants produce corresponding ~4B-sized models. However, we notice that the best depth-width pruned candidate (#7, 44 layers) still under-performe the width-only model (#1).
Combining depth and width pruning mitigates the degradation seen with depth-only pruning, but still under-performs width-only pruning.
Finding 3:
Mamba Layer Scoring. As described in Section 2.1 we choose the activations obtained from the Wa matrix for scoring the Mamba heads and head channels. We can alternatively get Mamba scores by considering the activations obtained from W₂ and Wo matrices. Table T the effect of selecting the Mamba activations from different parts of the Mamba layer. For different configurations, we notice that scoring the activations from We output often results in the best LM-loss.
Finding 4: Scoring Mamba heads using activations from the We matrix yields the best LM-loss.
Effect of Parameter Choice on Performance Metrics. As described in Section we obtain 125 checkpoints that fit the 4B parameter size constraint with differentt FFN channels, embedding dimensions, my (Mamba heads), and ma (Mamba head channels). For each checkpoint, we evaluate the LM-loss, time to first token, and throughput. To analyze the relationships between model parameters and performance metrics, we compute correlations and visualize them in Figure 7 Additionally, since all 125 models have the same total parameter count (4B), the model parameters exhibit negative correlations with one another. Figure 7 shows that in 4B models derived from Nemotron-H 8B, Mamba components positively correlate with latency and negatively with throughput and L.Mllos, indicating that pruning them improves inference speed and slightly degrades accuracy. In contrast, pruning embedding and FFN dimensions improves accuracy (lower LM-loss) but leads to slower models with increased latency and reduced throughput.
7

---

Finding 5: Pruning Mamba components improves inference speed but slightly degrades accuracy, while pruning embedding and FFN dimensions improves accuracy but slows down models.
Closer Look at Mamba Pruning. We analyze the sensitivity of pruning two axes in the Mamba layer- Mayeba heads (m₂) and Mamba head channels (ma))to various metrics, including accuracy, latency, and throughput. In this study, each axis was pruned in isolation while keeping the rest of the network unchanged, preserving the architecture of the Nemotron-H 8B model. The objective was to determine which axis is more favorable for optimization. As shown in Figure [ pruning Mamba heads (m,) consistently outperforms pruning Mamba head channels (ma) across all metrics. Specifically, reducing m, consistently yields lower LM-loss, reduced latency, and higher throughput, making Mamba heads particularly impactful and practical target for pruning.
Finding 6: Pruning Mamba heads (m²) consistently improves both accuracy and speed compared to Mamba head channels (ma).
Transformer vs. Hybrid Pruning. The above findings highlight the importance of choosing the right pruning axes in hybrid models to balance accuracy and efficiency. Unlike Transformer-on/ models where pruning attention heads is less common (20). hybrid architectures can tolerate some head pruning, seen with candidates #1 and #2 in Table This tolerance may stem from Mamba layers having significantly more heads (128) than self-attention layers (32).
Finding 7: Head pruning is better tolerated in Hybrid architectures compared Transformer-only.
3.2 Obtaining the Best Compressed Hybrid Model
For our final model, we select the candidate with the lowest LM validation loss in Table 0 which shows our architecture search results after short knowledge distillation. We note that although both candidates #1 and #2 have identical losses, candidate #1 is chosen for extended KD with 380B tokens due to its higher inference throughput, enabled by the reduction in Mamba heads.
Data and Training Hyperparameters. We use random sample from the Phase 3 data mixture employed for training Nemotron-H H models [5] for both importance estimation and knowledge distillation (KD). For importance estimation,
1024 samples with a sequence length of 8192. For KD, the batch size is 768, with a sequence length of 8192, a cosine decay learning rate schedule (starting at 1.6e-4 and decaying to 84-4), with 60-step linear warmup.
Alignment. We perform Supervised Fine-tuning with Knowledge Distillation (SFT-KD) 22 using the N 8B aligned model as the teacher, along with Reward-aware Preference Optimization (RPO) 1 and NeMo-Alignes [21]. The NEMOTRON-HE 4B base model is fine-tuned using supervision from the lop-k (100) logits of the teacher over two rounds of SFT-KD: the first round uses math and coding data, while the second round focuses on instruction-folle following and general chat data. The instruction-tunee model is then further aligned with two rounds of RPO.
Long Context Extension. To extend the context length of the aligned NEMOTRON- H 4B model, we perform SFT using data designed for long- mng-context understanding: this training data is derived by manipulating the general domain chat dataset from the second SFT-KD round during alignment. We concatenate conversation turns and introduce long-range dependencies by placing related turns far apart within the extended context. The context length is varied randomly between 128k and 512k tokens, ensuring the model learns to maintain coherence and understanding across longer sequences, enhancing its ability to process information beyond shorter context windows.
Evaluation Summary. Tables 2 to present accuracy comparisons between our compressed 4B hybrid model, other similar-sized community models, and the parent 8B hybrid model. As shown in Tables, our 4B model retains over 96% of the original 8B model's 's accuracy, including safety scores on Garak [8] and AEGIS [IO, while improving throughput by ~14x Compared other similarly sized community models, it delivers state-of-the of-the--t-ttttttt accuracy across knowledge, math, coding, commonsense reasoning, and reading comprehension tasks, despite being trained on up to -00 fewer tokens. It also achieves 2.2x higher throughput and 11.8 lower latency than the second-best Phi-4-4B model (Figures [ [ad The aligned version further leads in math, coding, instruction following, and tool-use tasks. To assess long-context capabilities, we use the RULER
8

---

Benchmarks (shols)
1.ama-Minitron 4B-Depth- Deptlb-Bass
LJam--Minitrnn 4B Widtth-Base'
Llama-3.1 3B- Buse"
Falcon-3 3B-Buse*
Zamba-2 2.7B- Base
Qwen-255 Base
NEMOTRON-H
Nemotrun-H
4B- Base"
SB-Base
ARC Challenge (0)
46.3
49.0
46.5
41.4
51.5
47.3
54.4
60.1
ARC Easy (0)
75.0
76.5
72.0
72.4
79.5
72.7
81.6
83.6
27.1
75.2
HellsSwag (0)
69.5
74.0
74.1
65.3
76.6
73.6
77.0
$12
HumanEval Eval (0, pass 81
23.2
34.4
26.8
39.6
25.0
37.8
89,8
57.3
Evall-rval-.o. (6, procee1)
18.9
20.7
24.4
32.3
21.3
33.5
55.5
53.7
paseee
47.9
65.0
pass
MMLU
57.8
59,8
56.3
568
65.6
68.1
72.7
OpenbookQA (AAQAA00 (0)
42.4
24.1
41.4
39.4
46.4
42.2
44.2
47.2
PIQA (0)
15.9
77.5
78.0
75.5
SO.4
78.8
79.4
82.2
RACE v.3 (0)
596
65.6
66.7
69.7
73.7
84.5
80,9
$4.0
Social IQA
TruthfuQQA
39,3
49.0
49.8
Winogrande (0)
68.8
70.4
69.5
65.0
74.3
68.4
71.3
76.3
Avange Tokens
49.6
53.1
51.1
$4.7
55.2
60.3
64.5
66.7
6.1T
0.17
VT
0.17
37
18T
0.3KT
1ST
Table 2: Accuracy comparison of our NEMOTRON-H 4B with other similarly sized base community models. indicates compressed models.
Benchnarkk (chots)
Phi-4-Mini 1--Amsrruc-122.
Owenn225
1.maa-33
Falcun-3 38-Instruct3331
Zambu-2 2.78--mntruci-ik
NEMOTRON-F 4B Instrudd-126kk
Nemotron-H SB: 128-
31--Iattrui--3k
38-Iastruct-114
MMLEE (0, geneniive)
61 8B
63.25
$7.30
$4.27
55.32
66.96
68
GSMSK (0)
$7.71
$3.32
73.47
77.86
66.26
88.03
90.4
MATE--S90 (0)
70.8
65.6
48.2
45.90
29.40
76.4
77.6
Hlumontival (0, pa
75:07
75.0
49
66.34
19.10
76.2
79.3
HomonEival (asct)1)
70.85
(0, pasctt 1)
MBPP: pass(i1)
14.78
(9)
Table 3: Accuracy comparison for instruction-tuned models. For IFEval, we report the average of prompt striet and instruction strict categories. For BFCL v2, report live overall accuracy. For MT-Bench, we use GPT-4-Turbo as the judge. indicates compressed models.
Context Length
Phi-4-Mini
Qwen-2.5
Llama-3.2
NEMOTRON-H 48-Instruct-1
Nemotron-H NB-Insttruct-129
3B-Insrruct-32k
3B-Instruct-
16,384
34.39
83.64
77.92
86.28
91.5
32,768
32.90
79.21
72.71
82.27
89.8
65,536
35.01
66.42
75.95
131,072
20.07
23.61
59.26
63,57
81.7
Table 4: Average RULER benchmark scores up to 128k context length for aligned NEMOTRON-E 4B and other instruction-tune -uuned models in similar size range. indicates compressed models.
As shown in Table our model demonstrates strong performance and achieves the highest scores at context lengths up to 128k tokens. Finally, Figure 6 (Appendix) compares latency and throughput across four models: Phi-4-Mini-4B, Qwen-2.5-3B Nemotron-H 8B, and NEMOTRON-H 4B (ours). Our model achieves the best performance on both axes delivering the fastest -irrt-to-ittt and highest dbroughput-effectivee
benchmark [1].
advancing the latency-throughput Pareto frontier. In summary, compression approach successfully produces model with state-of- the-art accuracy while significantly improving inference speed and reducing training costs.
Related Work
Structured Weight Pruning. Weight pruning is a powerful and well-known technique for reducing 26, [4]. In particular, structured pruning removes blocks of nonzero elements once
model size (20
from model weights, making it easier realize actual hardware speedups; examples of structured pruning techniques include neuron, attention head, convolutional filter, and depth pruning [20) [7]
27
[8 28 16). Recent work 20 has demonstrated that knowledge distillation
13] can be an
effective alternative to traditional fine-tuning for post-pruning accuracy recovery.
Hybrid Model Pruning. Recent work on SSM pruning has focused on Mamba architectures: Manba-Shredder [19 proposes removing the entire state space module from each Mamba layer, leaving only linear projections and convolution components; while effective for model size reduction, this approach severely impact sequence modeling performance. In parallel, Ghattas et al. [9]

---

Model
Garak Score
AEGIS Score
Nemotron-H 70.75
99.83
Nemotron-H-44
67.77
98.17
Table 5: Safety scores before and after compression.
introduce a granular pruning strategy, targeting state space dimension, Mamba head dimension, and head merging. Unfortunately, this approach has limited scalability, especially for large models requiring tensor parallelism; the proposed head merging also risks per-head information loss.
5
Conclusions
This paper has presented novel group-aware pruning method for Mamba layers along with unified hybrid pruning recipe that effectively combines Mamba pruning with FFN, embedding dimension and layer pruning. We have applied this unified recipe to the Nemotron-H 8B model to obtain NEMOTRON-E H 4B, a compressed hybrid language model that achieves state-of-the-ad accuracy and efficiency. Despite 50% reduction in model size, NEMOTRON-H 4B retains over 96% of the original 8B model's accuracy, while requiring up to 40x fewer training tokens. Notably, our compression method also maintains the model's safety performance. NEMOTRON-H 4B advances the accuracy-effiiinncy Pareto frontier, achieving ~x faster inference and 2.6% higher accuracy across diverse of tasks.
References
[1] Adler, Niket Agarwal, Ashwath Aithal, Dong Н Anh, Pallab Bhattacharya, Annika Brundyn, Jared Casper, Bryan Catanzaro, Sharon Clay, Jonathan Cohen, et al. Nemotron-4 340b technical report. arXiv preprint Xrx::22406. 11704, 2024.
[2] Yongqi An, Xu Zhao, Tao Yu, Ming Tang, and Jinqiao Wang. Fluctuation-hased adaptive structured pruning for large language models. In Proceedings of the AAAI Conference Artificial Intelligence, volume 38, pages 10865-10873, 2024.
[3] Salch Ashkboos, Maximilian L Croci, Marcelo Gennari do Nascimento, Torsten Hoefler, and James Hensman. Slicegpt: Compress large language models by deleting rows columns. In The Twelfth International Conference on Learning Representations, 2023.
[4] Akhiad Bercovich, Tomer Ronen, Talor Abramovich, Nir Ailon, Nave Assaf, Mohammad Dabbah, Ido Galil, Ammon Geifman, Yonatan Geifman, Izhak Golan, Netanel Haber. Ehud Karpas, Roi Koren, Itay Levy, Pavlo Molchanov, Shahar Mor. Zach Moshe, Najeeb Nabwani, Omri Puny. Ran Rubin, Itamar Schen, Ido Shahaf, Oren Tropp Omer Ullman Argov, Ran Zilberstein, and Ran El- ElnYa..v. Puzzle: Distillation-Bass NAS for Inference-Optimizee LLMs, 2024.
[5] Aaron Blakeman, Aarti Basant, Ahhinav Khattar, Adithya Renduchintala, Akhiad Bercovich, Aleksander Ficek, Alexis Bjorlin, Ali Taghihakhshi, Amala Sanjay Deshmukh, Ameya Sunil Mahabaleshwarkar, et al. Nemotron-h: A family of accurate and efficient hybrid mamha-transformer models. arXiv preprint TrXXi:2504.03022 2025.
[6] Hung- Yuch Chiang. Zhen-Has Liu, Chien Yu Wang, Tsung- Han Chen. Ting-An Wei, Wei-Chen Liu, Chih Mhth-Cheng Liu, Yu-Ching Chuang, Jhih-Cing Wei, Y.-C F. Wang, and Wei-Chao Chang. Quamba2: robust and scalable post-training quantization framework selective state space models. arXiv preprint arXiv:2405. 5077, 2024.
[7] Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.21066 2024.
[8] Leon Derczynski, Erick Galinkin, Jeffrey Martin, Subho Majumdar, and Nanna Inie, garak: A framework for security probing large language models, 2024.
[9] Tamer Ghattas. Michael Hassid, and Roy Schwartz. On pruning state-space 1ls.. arXiv preprint 250iv:2502. 1888. 2025.
[10] Shaona Ghosh, Prasoon Varshney. Erick Galinkin, and Christopher Parisien. Aegis: Online adaptive ai content safety moderation with ensemble of Ilm experts, 2024.
10

---

[11] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint urXiv:2312.00752, 2023.
[12] Yang He. Guoliang Kang, Xuanyi Dong. Yanwei Fu, and Yi Yang. Soft filter pruning for accelerating deep convolutional neural networks. arXiv preprint 10Xiv:188.006860 06866, 2018.
[13] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the Knowledge in a Neural Network. arXiv preprint arXiv: 02531,, 2015.
[14] Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste. Sparsity in Deep Learning: Pruning and growth for efficient inference and training in neural networks. arXiv preprint urXiv:2102.00554 2021.
[15] Cheng-Ping Hsich, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's 's the real context size your long-context language models?, 2024.
[16] Bo- Bo-Kyeong Kim, Geonmin Kim, Tae-Ho Kim, Thibault Castells, Shinkook Choi, Junho Shin, and Hyoung- Kyu Song. Shortened LLaMA: A simple depth pruning for large language models. In ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2024.
[17] Jian-Hao Luo, Jianxin Wu, and Weiyao Lin. Thinet: A filter level pruning method for deep neural network compression. In Proceedings of the IEEE international conference on computer vision, pages 5058-5066, 2017.
[18] Xin Men, Mingyu Xu, Qingyu Zhang, Bingning Wang, Hongyu Lin, Yanjie Lu, Xianpei Han, and Weipeng Chen. ShortGPT: Layers in Large Language Models are More Redundant Than You Expect, 2024.
[19] Pablo Muñoo, Jinjie Yuan, and Nilesh Jain. Mamba-shedder: Post-transformer compression for efficient selective structured state space models. arXiv preprint arXiv:2501. 17088 2025.
[20] Saurav Muralidharan, Sharath Turuvekere Sreenivas, Raviraj Joshi, Marcin Chochowski, Mostofa Patwary, Mohammad Shoeybi. Bryan Catanzaro, Jan Kautz, and Pavlo Molchanov. Compact language models via pruning and knowledge distillation. arXiv preprint arXiv:2407. 14679, 2024.
[21] Gerald Shen, Zhilin Wang, Olivier Delalleau, Jiaqi Zeng, Yi Dong, Daniel Egert, Shengyang Sun, Jimmy Zhang, Sahil Jain, Ali Taghibakhshi, Markel Sanz Ausin, Ashwath Aithal, and Oleksii Kuchaiev. Nemo- aligner: Scalable toolkit for efficient model alignment, 2024.
[22] Anna Shors, Sharath Sreenivas, Shengyang Sun, and Teny Kong. Data-efficient knowledge distillation for supervised fine-tuning with nvidia nemo aligner. lattps://developer www.iiiiiiciooooo wtt.ddi..com//bie coo/blog/////oooooooooccoooooooo data- efficient- knowledge mhtoleegg-ddssillattae distillation for cor-esueeeeeed fine- tuning- with nvidia nemo-
aligner/ May 2024.
[23] Sharath Turuvekere Sreenivas, Saurav Muralidharan, Raviraj Joshi, Marcin Chochowski, Ameya Sunil Mahabalcshwarkar, Gerald Shen, Jiaqi Zeng, Zijia Chen, Yoshi Suhara, Shizhe Diao, Chenhan Yu, Wei- Chun Chen, Hayley Ross, Oluwatobi Olabiyi, Ashwath Aithal, Oleksii Kuchaiev, Daniel Korzekwa, Pavlo Molchanov, Mostofa Patwary, Mohammad Shocybi, Jan Kautz, and Bryan Catanzaro. LLM Pruning and The Minitron Approach, 2024.
Distillation in Practice:
[24] Shengkun Tang, Oliver Sieberling, Eldar Kurtic, Zhiqiang Shen, and Dan Alistarh. Darwinlm: Evolutionary structured pruning of large language models. arXiv preprint arXiv:2502.0778) 2025.
[25] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[26] Wenxiao Wang, Wei Chen, Yicong Luo, Yongliu Long, Zhengkai Lin, Liye Zhang, Binbin Lin, Deng Cai, and Xiaofci Hc. Model compression and efficient inference for large language models: A survey, arXiv preprint 2rx:::2402. 09748, 2024.
[27] Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared llama: Accelerating language model pre-training via structured pruning. In The Twelfth International Conference on Learning Representations, 2023.
[28] Yifei Yang. Zouying Cao, and Hai Zhao. Laco: Large language model pruning via layer collapse. arXiv preprint 2402 2024.

---

A Technical Appendices and Supplementary Material
A.1 FLAP Importance for Hybrid Models
FLAP 2 a retraining-free structured pruning technique designed to measure the recoverability of model's s output feature map upon removing specific columns from weight matrices. FLAP quantifies the "fluctuation" of each input feature relative to a baseline using calibration data. Specifically, the FLAP importance score for a column is computed as the product of the squared norm of the column weights and the sample variance of the corresponding imput features across calibration samples.
We extend FLAP to the SSM layers in hybrid architectures by applying the metric to the activations serving as inputs to the output projection (Wo). Here, we compute the FLAP importance by assessing the variance in activations input to the Wo, weighted by the squared norms of the respective columns of the Wo. Mathematically, the extended FLAP importance metric for given column j weight matrix W in SSM layers can be defined as:
Sj = -W, Var(X,)
where ||2 denotes the squared norm of the column weights and Var(X;) represents the variance of the activations input to the output projection matrix of SSM layer across calibration samples.
We use the above-computed metric rank different heads within each group and remove the corresponding rows in the input projection matrix, the corresponding channels in the SSM convolution kernel, corresponding rows the A and D matrices of SSM, as well trimming the corresponding columns in the output projection matrix.
Table O shows that FLAP- importance estimation yields mixed results before lightweight KD across pruning strategies. After KD, performs on par with the L2-based approach when applied to candidate #1; it doesn't seem to offer any clear advantage, however.
Results.
Pruning Type
Configuration
L2 LM Loss
FLAP LM Loss
Bascline
No pruning
1.168
1.168
FEN 16384 FEN 11568
1.364
1.32
FEN
1.803
1.64
FFN 8192
2.281
1.95
Attention
ATT Heads 16
1.282
1.40
Mamba Heads = 112
1.305
1.73
Mamba (SSM)
Mamha Heads
2.150
4.59
Mamba Heads 64
9.040
11,21
#1 #1 + Lightweight KD
3.690
5.854
Mixed
1.380
1,380
Table 6: LM Joss comparison when pruning different model components using L2 and FLAP metrics. Baseline: 128 Mamba heads, 21,504 FFN size, 32 attention heads.
A.2
Depth Pruning
Figure shows the average importance scores for each layer in the Nemotron-H 8B Base model, with green, blue, and red dotted lines representing self-attention, FFN, and Mamba layers. As seen in previous work [5], the most important layers are concentrated at the model's start and end. Interestingly, the first attention layer is among the least important, while other attention layers are more critical than neighboring layers. "saw-like" pattern emerges where MLP layers are more important than adjacent Mamba layers in the middle of the network, though this reverses in the model's 's critical regions.
We experimented by pruning the least important layers (4, 8, 12, 16, and 26 layers), followed by distillation with 126B tokens. While core- knowlede benchmarks remained largely unaffected, tasks like math and coding showed significant performance degradation (Figure
12

---

Per lnyer Importance
tellnws;
221
muu
Acoumay or
Aieegande
pandi
humareval gready
prooyy
Layer
Figure 4: Layer importance measured as the KLD between logits of the full model and model with that layer removed, averaged over small training subset. Vertical dotted lines in- dicate layer types: self-attention (green), FFN (blue), and Mamba2 (red).
Number of lyyers
Figure 5: Accuracy drop relative the 8B model across progressively depth-only pruned variants (48, 44, 40, 36, and 26 layers). Each model is directly pruned from the 8B and distilled using 126B tokens.
A.3 Mamba Scoring
FEN Embedding Dim
Mamba
LM-Loss
Heads
Head Channels
W,
W;
Wo
12,288
13,056
3,072
112
56
3,59
6.61
5.30
14,592
3,072
96
56
4.68
7.09
10.01
12,288
128
13,824
3,072
128
48
5,99
6,01
9.47
Table Mamba scoring ablation. The zero-shot LM-loss for top 6 pruned models based on Mamba
scores calculated from activations of W,. Wa, and Wo. The W, activations result in the best zero- shot LM-loss most of the cases.
A.4 Throughput and Latency Comparisons
Figure 6 compares the throughput and latency of NEMOTRON-F H 4B to Phi-4 Mini 4B, Qwen-2.5 3B and Nemotron-H 8B.
Latenoy vs Throughput Frontier (Batch Size
Ph4-48
Relative Latency (Normalized to Min)
Owon-38
Nomotron-Hi 03
Nemoiror-Ht 48 (num)
Lower Beller
Higher Beller
0.5
0.6
07
0.8
0.9
1.0
Reletve Throughput (Normalized to Max)
Figure 6: Throughput and latency comparisons across four models: Phi-4-Mini-4B, Qwen-2.5-3B, Nemotron-H 8B, and NEMOTRON-H 4B (ours). Relative throughput and latency represents measured for an input and output context length of 65536 and 1024, respectively.
A.5 Model Parameter Correlations
This section details the relationships between different model components for our variants, illustrating the hyperparameter trade-offss required to maintain fixed parameter budget (Figure 7)
13

---

Furthermore, we present an ablation study justifying our choice to optimize Mamba heads (m²) over head channels (ma)), as m., pruning yields superior results across all metrics (Figure g
my mtunin (/Akza))
myg trpporg
ptaring (Jateney)
muatunin 16hros)
my pootnggTowoggp.00
m, pruving 3.thnas)
Pfuuiig Mamoa on LMiloss. Throughoul, and TTFT
betaer
Relato Labeoor Lower bebot
Parumatin Midee Convialiro 48 Moddlss Paanceeehaunccr Condistara Motea
Lidcce- Lower codee
9/8
FPN
-0.0
-M
487
Lekrey
Paltemana Machs
Partenlage Phating (76)
Figure 7: Left: Correlation matrix showing rela- tionships between performance metrics and model components FFN, embedding dimension (d,), and Mamba parameters (varying both heads my and head dimension ma) across 125 4B vari- ants with fixed depth (52 layers). Right: Model parameter correlations for fixed 4B parameter trade-off's where increasing one component reduces others.
Figure 8: Impact of pruning Mamba heads (m²) versus Mamba head channels (m,) isolation, with the rest of the network unchanged. Pruning my consistently outperforms ma pruning across LM loss, latency, and throughput-establisse it as the preferred target for optimization.
A.6 Generalizability to Mamba2
To evaluate the generalizability of our compression strategy to other models, apply it to the Mamba2 1.3B model 77]. We prune the model to 780M parameters via SSM and embedding pruning, and then subsequently train the prumed model 10.5B tokens. We compare
pruned 780M model
to the Mamba2 780M and 1.3B models trained from scratch on 300B tokens 7
As shown in Table our compressed 780M model, despite being trained on significantly fewer tokens (10.5B vs. 300B), outperforms the 780M model trained from scratch and achieves an average score comparable to the original 1.3B model. These results provide further insights into the generalizability of our compression method.
Table 8: Comparison of our compressed Mamba2 780M model against Mamba2 780M and 1.3B models trained from scratch (7). Despite being trained on significantly fewer tokens (10.5B vs. 300B), our compressed model achieves a better average score than the 780M baseline.
Benchmark
Mamba2 780M
Mamba2
Compressed 780M
are_challenge
28.6
33.2
34.2
arc arc_casy
54.7
60.6
commonsense_qq
19.6
20,9
26.4
hellaswag
54.7
59,9
50.4
openbookqa
36.4
37.0
34.2
piqa
72.1
73.5
race
21.8
24.8
32,5
social_iqa
41.0
42,9
41.1
truthfulqa _m2 mc2
38.1
36.1
38.9
winogrande
58.0
60.1
57.7
Avg
42.5
44.9
44.7
Tokens
300B
300B
10.5B
A.7 Compute Resources
All experiments were performed on 16-32 x NVIDIA DGX H100 nodes (8x H100 80GB) for short tumaround times.
14

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Experimental results provided in Section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Section 5 states the limitation of this work.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
15

---

Justification:
No theoretical proofs or results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide relevant information for reproducing our results in Section 3.2 and the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
16

---

Answer:
[No]
Justification:
We intend to release the model weights and code pending internal review.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ w/wpps.cc/publlie. cttccc//uuuu www.ppublc/guuucuccceeeccceeboee d ssi onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: Provided in Section 3.2 Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Training LLMs is extremely expensive; therefore unable to perform multiple runs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
17

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Detailed in Section A.7
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Question:
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
Generic algorithm for optimizing LLM training.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
18

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
No such release artifacts.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Assets properly credited.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, thpprd, wwteccode..comddaiio dattsatsssss has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
19

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
No new assets released.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
20

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: LLM is not used for the methodology, scientific rigorousnsss originality of the research.
Guidelines:
The answer NA means that the method development in this research does involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy for what should or should be described.
21


Reviews:
[{"review_id": "zng3nz6s3s", "reviewer": "Reviewer_2zWJ", "summary": "The paper proposes a group-aware pruning method for the Mamba layers, which preserves the structural integrity of SSM blocks and their sequence modeling capabilities. The authors combine this with embedding, FFN, and layer pruning to further reduce the parameter count, and they conduct extensive ablations to study the effect of pruning different components on accuracy, latency, and throughput — while keeping the total parameter count fixed at 4B. The proposed pruning method enables compression model by 50% with minimal drop in performance.", "questions": "1. How well does the proposed pruning strategy generalize for different family architectures? \n\n2. How does the pruning affect the model bias for hybrid LLMs?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "### Strengths\n1. Hybrid attention is critical for long context performance; however, pruning hybrid architecture is not well studied. This paper presents a compression strategy that combines SSM pruning with the pruning of other network components such as FFN neurons, embedding channels, and network depth. Authors also conducted an in-depth architectural search for pruning different components (different layers, depth vs width, etc). \n\n2. As demonstrated in Figure 3, the proposed group-aware pruning method for the Mamba layer preserves the structural representation of the Mamba blocks and provides a simple method for pruning hybrid LLMs.\n\n3.  Extensive ablations are conducted in Section 3.1  to compare different pruning strategies and the impact of pruning different components on validation loss, latency and throughput (Table 1). \n\n4. As shown in Table 2, the proposed pruning strategy achieves 50% compression rate while retaining  over 96% of the original 8B model’s accuracy, while improving throughput by ∼1.4x – demonstrating the efficacy of the method.\n\n5. Pushes Pareto-frontier for throughput, though results are limited to the Nemotron family.\n\n### Weaknesses\n\n1. The major limitation of this work (as the authors also noted) is that all the experiments are conducted with the Nematron model; it is not clear if the observation and analysis will hold true for different model families. \n\n2. The model performance deteriorates after pruning, which is later recovered using distillation. \n\n3. There is no ablation on pruning heuristic used in sec 2.2. \n\n4. No discussion on the potential impact of pruning strategy proposed on the model bias, as prior work has highlighted that model compression exacerbates model bias.[1].\n\n[1]. Gonçalves and Strubell, Understanding the Effect of Model Compression on Social Bias in Large Language Models", "was_revised": true, "final_justification": "The authors have addressed most of the concerns and have included new experiments with different model families and analysis on model bias.", "rebuttal": "We thank the reviewer for their insightful questions and discussions. We will try our best to address the questions and concerns below:\n\n# Weaknesses \n## Reviewer: it is not clear if the observation and analysis will hold true for different model families\n\nTo help answer the question on whether our method generalizes to other hybrid models, we have applied our compression strategy to the original Mamba2 1.3B model [6], trimmed it down to 780M parameters via ssm and embedding pruning, and then further trained the pruned model on 10.5B tokens. We compare our pruned 780M model to the Mamba2 780M model trained from scratch on 300B tokens [6]. Given the short time limit of the rebuttal and limited compute, this was the largest model we could evaluate; however, we hope that these Mamba2 pruning results provide further insights into the generalizability of our method.\nAs shown in the table below, our compressed 780M model outperforms the model trained from scratch.\n\n| Benchmark        | Mamba2 780M | Mamba2 1.3B | Compressed 780M |\n|------------------|-------------|-------------|------------------|\n| arc_challenge    | 28.6        | 33.2        | 34.2             |\n| arc_easy         | 54.7        | 60.6        | 60.4             |\n| commonsense_qa   | 19.6        | 20.9        | 26.4             |\n| hellaswag        | 54.7        | 59.9        | 50.4             |\n| openbookqa       | 36.4        | 37.0        | 34.2             |\n| piqa             | 72.1        | 73.5        | 71.3             |\n| race             | 21.8        | 24.8        | 32.5             |\n| social_iqa       | 41.0        | 42.9        | 41.1             |\n| truthfulqa_mc2   | 38.1        | 36.1        | 38.9             |\n| winogrande       | 58.0        | 60.1        | 57.7             |\n| **Avg**          | **42.5**    | **44.9**    | **44.7**         |\n| **Tokens**       | **300B**    | **300B**    | **10.5B**        |\n\nWe will include these results in the supplementary materials.\n[6] Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through 320 structured state space duality. arXiv preprint arXiv:2405.21060, 2024.\n\n## Reviewer: model performance deteriorates after pruning\nThe considerable deterioration in model performance right after pruning, and before applying knowledge distillation, is expected and universal (e.g. [17]; this is because some of the heads, channels, embedding dimension, and MLPs that contribute to the final logits are removed from the model. However, this degradation, although major, is recoverable to a significant degree with minimal data and knowledge distillation, as shown in the paper.\n\n## Reviewer: ablation on pruning heuristic used in sec 2.2\nWe would like to point out that since the aggregation metric for FFN and Layernorms is not Mamba-dependent, we leverage the ablations from [17]. However, if the reviewer advises so, we would be happy to add this ablation to the appendix. Regarding the depth pruning and FLAP importance in Section 2.2, the ablations and heuristics are available in Section A.2 and A.1 of the appendix provided with the submitted manuscript. \n\n## Reviewer: potential impact of pruning strategy proposed on the model bias\nThank you for highlighting this important concern. We evaluated Garak and AEGIS scores for the compressed 4B model and its 8B parent, observing similar relative differences as in the evaluation benchmarks. This indicates the safety scores are as expected for the compressed model. We will include these results and the discussion in the supplementary materials.\n| Model              | Garak Score | AEGIS Score |\n|--------------------|-------------|-------------|\n| Nemotron-H-8B      | 70.75%      | 99.83%      |\n| Compressed-4B      | 67.77%      | 98.17%      |\n\n# Questions\n\n \n## Reviewer: How well does the proposed pruning strategy generalize for different family architectures?\n\nPlease refer to the Mamba2 1.3b to 780m compression discussed above\n\n## Reviewer: how pruning affects the model bias\n\nPlease refer to our answer on model safety and bias above."}, {"review_id": "FFtNhSgwyB", "reviewer": "Reviewer_D6BM", "summary": "This paper introduces a ​​compression framework​ for hybrid Transformer-State Space Models, the innovations include:\n- A group-aware pruning method for Mamba layers that preserves SSM block structure and sequence modeling capabilities.\n- A unified hybrid pruning recipe\n\nThe proposed method compresses the Nemotron-H 8B model to 4B parameters with much less training tokens and comparable performances.", "questions": "- How knowledge distillation affects the performances? Results with pruning only before distillation can be ablated to showcase the effectiveness of proposed pruning techniques.\n- I wonder the generalization ability of this method, what will the performance be if the proposed method is applied on other open-sourced models such as jamba.", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "strength\n- The paper is well written and well-motivated.\n- The 7 findings provide meaningful and comprehensive guidance for SSM-based model compression in the community. \n- The performance and efficiency of the final model is well evaluated and achieves strong performances.", "was_revised": true, "final_justification": "The paper proposes a meaningful for mamba-based model pruning. And my two questions are basically addressed by the rebuttal.", "rebuttal": "We thank the reviewer for their valuable insights and comments. We have tried our best to provide the responses for the questions as follows:\n\n# Questions \n\n## Reviewer: How knowledge distillation affects the performances? Results with pruning only before distillation can be ablated to showcase the effectiveness of proposed pruning techniques.\n\nThank you for this point. Assuming performance here refers to benchmark scores, we note that the pruned model initially performs near-randomly, but retains core capacity to recover through distillation. In terms of LM loss, it starts at 3.3 and improves to 1.15 after full knowledge distillation. \n\nWe analyze the impact of pruning along different axes on LM loss (without distillation) and speed in Figures 7 and 8. As noted in Section 3.1, we prune 125 candidates and select the best 22 based on LM loss post-pruning. However, our main selection metric across different pruning techniques/candidates is LM loss after short (6B token) distillation, as shown in Table 1, since better LM loss post-pruning does not always translate to better post-distillation performance. This behavior has been demonstrated in [17]. If advised, we will add the above details to the manuscript.\n\n[17] Saurav Muralidharan, Sharath Turuvekere Sreenivas, Raviraj Joshi, Marcin Chochowski, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro, Jan Kautz, and Pavlo Molchanov. Compact language models via pruning and knowledge distillation. arXiv preprint arXiv:2407.14679, 2024.\n\n\n## Reviewer: I wonder the generalization ability of this method, what will the performance be if the proposed method is applied on other open-sourced models such as jamba.\n\nTo help answer the question on whether our method generalizes to other Mamba based models, we have applied our compression strategy to the original Mamba2 1.3B model [6], pruned it down to 780M parameters via ssm and embedding pruning, and then further trained the pruned model on 10.5B tokens. We compare our pruned 780M model to the Mamba2 780M model trained from scratch on 300B tokens [6]. Unfortunately, due to the very short response period, we couldn’t apply our method to Jamba; however, we hope that these Mamba2 pruning results provide further insights into the generalizability of our method.\n\nAs shown in the table below, our compressed 780M model outperforms the model trained from scratch.\n\n| Benchmark        | Mamba2 780M | Mamba2 1.3B | Compressed 780M |\n|------------------|-------------|-------------|------------------|\n| arc_challenge    | 28.6        | 33.2        | 34.2             |\n| arc_easy         | 54.7        | 60.6        | 60.4             |\n| commonsense_qa   | 19.6        | 20.9        | 26.4             |\n| hellaswag        | 54.7        | 59.9        | 50.4             |\n| openbookqa       | 36.4        | 37.0        | 34.2             |\n| piqa             | 72.1        | 73.5        | 71.3             |\n| race             | 21.8        | 24.8        | 32.5             |\n| social_iqa       | 41.0        | 42.9        | 41.1             |\n| truthfulqa_mc2   | 38.1        | 36.1        | 38.9             |\n| winogrande       | 58.0        | 60.1        | 57.7             |\n| **Avg**          | **42.5**    | **44.9**    | **44.7**         |\n| **Tokens**       | **300B**    | **300B**    | **10.5B**        |\n\nWe will include these results in the supplementary materials.\n\n[6] Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through 320 structured state space duality. arXiv preprint arXiv:2405.21060, 2024."}, {"review_id": "IUrTuZ8feT", "reviewer": "Reviewer_WxVs", "summary": "This paper proposes a group-aware pruning method for the Mamba layer. It combines Mamba pruning with pruning of other network layers to produce smaller and more accurate models while reducing training costs.\nThis paper uses this method to compress the parameters of the Nemotron-H 8B hybrid model, reducing training tags, increasing throughput, and improving accuracy compared to models with the same parameters.", "questions": "1. Can the pruning method show similar effects on models other than Nemotron-H 8B?\n2. Comparison of the compressed model with the model of the same series with smaller parameters.\n3. Experiments and analysis of inference throughput.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n\nThis paper focuses on pruning and distillation, and the proposed pruning method is intuitively reasonable. A relatively detailed experiment and analysis are conducted, and 7 findings are proposed and corresponding experimental results are provided. The compressed model of Nemotron-H 8B also achieves improvements over the same parameter model in knowledge, mathematics, encoding, common sense reasoning, and reading comprehension tasks.\n\n**Weaknesses**\n\nThis paper only distills Nemotron-H 8B, which makes people wonder whether the pruning scheme can be extended. I hope the authors can provide more than one compressed model and show similar results. Does the Mamba layer limit the range of LLMs that can be pruned?\nIn the main experiment, the Nemotron-H-4B compressed model is only used for comparison with the non-Nemotron series baseline. In order to eliminate the interference of model quality, can you provide the comparison results of the compressed model with the model with smaller parameters in the same series? For example, compress Qwen-7B to 3B and compare the compressed model with Qwen-3B?\nThe author mentioned that the compressed model achieved 2 times faster inference throughput, and experiments and analysis of this conclusion are needed.", "was_revised": true, "final_justification": "Many of my concerns were addressed by the authors, and I encourage them to include new results and interpretations in the final paper.\n\nThank you for your additional analysis, we would like to raise our rating to 5", "rebuttal": "We thank the reviewer for their valuable insights and comments. We have tried our best to provide the responses for the questions as follows:\n# Weaknesses\n## Reviewer: Can the pruning method show similar effects on models other than Nemotron-H 8B?\nTo help answer the question on whether our method generalizes to other hybrid models, we have applied our compression strategy to the original Mamba2 1.3B model [6], pruned it down to 780M parameters via ssm and embedding pruning, and then further trained the pruned model on 10.5B tokens. We compare our pruned 780M model to the Mamba2 780M model trained from scratch on 300B tokens [6]. Given the short time limit of the rebuttal and limited compute, this was the largest model we could evaluate; however, we hope that these Mamba2 pruning results provide further insights into the generalizability of our method.\nAs shown in the table below, our compressed 780M model outperforms the model trained from scratch.\n\n| Benchmark        | Mamba2 780M | Mamba2 1.3B | Compressed 780M |\n|------------------|-------------|-------------|------------------|\n| arc_challenge    | 28.6        | 33.2        | 34.2             |\n| arc_easy         | 54.7        | 60.6        | 60.4             |\n| commonsense_qa   | 19.6        | 20.9        | 26.4             |\n| hellaswag        | 54.7        | 59.9        | 50.4             |\n| openbookqa       | 36.4        | 37.0        | 34.2             |\n| piqa             | 72.1        | 73.5        | 71.3             |\n| race             | 21.8        | 24.8        | 32.5             |\n| social_iqa       | 41.0        | 42.9        | 41.1             |\n| truthfulqa_mc2   | 38.1        | 36.1        | 38.9             |\n| winogrande       | 58.0        | 60.1        | 57.7             |\n| **Avg**          | **42.5**    | **44.9**    | **44.7**         |\n| **Tokens**       | **300B**    | **300B**    | **10.5B**        |\n\nWe will include these results in the supplementary materials.\n[6] Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through 320 structured state space duality. arXiv preprint arXiv:2405.21060, 2024.\n\n## Reviewer: Does the Mamba layer limit the range of LLMs that can be pruned?\nAdding Mamba layer pruning support expands the range of LLMs that can be compressed by providing more axes to prune the model along. Prior works [20] were limited to compression of LLMs with self attention and FFN layers only. Our addition of Mamba layer pruning extends this to support hybrid models (with Mamba layers) and pure Mamba models as well.\n\n## Reviewer: Comparison of the compressed model with the model of the same series with smaller parameters.\nWhile no 4B hybrid model with the same architecture trained from scratch exists, we address this concern with additional results on compressing Mamba 1.3B → 780M, included in the response to the weaknesses above.\n\n## Reviewer: Experiments and analysis of inference throughput\nInference and throughput were measured using the Megatron-LM repository (https://github.com/NVIDIA/Megatron-LM/tree/main/megatron ) for all the models. Figure 6 shows a relative 2x speedup for the compressed 4B model over the previous SOTA, Phi-4-Mini-4B. We additionally provide the table below with absolute numbers for the same setting (64k input, 1k output). \n\n| Model                 | Throughput (tokens/s/GPU) | Latency (time to first token/s) |\n|-----------------------|----------------------------|----------------------------------|\n| Phi-4-Mini-4B         | 131.7                      | 12.384                           |\n| Qwen-2.5-3B           | 204.0                      | 9.606                            |\n| Nemotron-H 8B         | 212.5                      | 9.524                            |\n| COMPRESSED-HYBRID 4B  | 287.5                      | 6.6                              |\n\nWe also note that the data on how different pruning axes affect latency and throughput are available in Figure 7 and 8 in the manuscript, and the relative inference throughput numbers for the candidates we provide in Table 1 is available in the last column of the same table.\n\n#Questions\n\n## Reviewer: Can the pruning method show similar effects on models other than Nemotron-H 8B? Comparison of the compressed model with the model of the same series with smaller parameters.\n\nPlease refer to our response above on the added study on compressing Mamba2 1.3b to 780m.\n\n## Reviewer: Experiments and analysis of inference throughput.\n\nPlease refer to our response in the final section addressing the identified weaknesses."}, {"review_id": "BO187jSMoS", "reviewer": "Reviewer_tXTL", "summary": "This paper explores width, depth, and head pruning for hybrid Mamba architectures. In particular, it proposes a novel method for pruning SSM (State Space Model) modules.", "questions": "I am confused about the **Group-Aware Head Permutation Constraints** and **Head Channel Consistency**. Prior work [1] suggests that the channel dimension in the Mamba architecture exhibits **permutation invariance**. Why, then, is it necessary to maintain **uniform pruning across all heads within a group**, rather than allowing permutation of heads across groups?\n\n[1] Hung-Yueh Chiang et al.,  Quamba2: A Robust and Scalable Post-training Quantization Framework for Selective State Space Models", "limitations": "The authors transparently acknowledge all of the limitations I was able to identify within the paper itself.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "### **Strengths**:\n\n- This is the first work, to the best of my knowledge, that conducts an in-depth analysis of different pruning strategies for the Mamba architecture. The pruning strategy is systematically explored across multiple architectural axes (depth, width, and heads).\n\n\n\n### **Weaknesses**:\n- The paper lacks sufficient **discussion of related work** and **stronger baseline methods**, such as [2]. Including these would improve the completeness of the evaluation.\n- In addition, the **experimental scope is limited**—the evaluation is conducted solely on **Nemotron-H-8B-Base**, which makes it difficult to assess the general applicability of the proposed method. Broader experiments on additional models are recommended.", "was_revised": true, "final_justification": "Thank you for your response. I believe your explanation has effectively resolved my confusion.", "rebuttal": "We thank the reviewer for their valuable insights and comments. We have tried our best to provide the responses for the questions as follows:\n# Weaknesses:\n\n## Reviewer: discussion of related work and stronger baseline methods\nAt submission, we were aware of only two Mamba pruning papers (refs. [7] and [16]), both cited. For transformer pruning (self-attention and FFN), we referenced [4], [17], and [20]. Regarding ref. [2] mentioned by the reviewer, assuming they’re referring to ref. [2] in the paper, we extended the FLAP methodology to Mamba in Appendix A1 and included results. We're happy to add any missing references if advised by the reviewer.\n\n## Reviewer: the evaluation is conducted solely on Nemotron-H-8B-Base\n\nTo help answer the question on whether our method generalizes to other hybrid models, we have applied our compression strategy to the original Mamba2 1.3B model [6], trimmed it down to 780M parameters via ssm and embedding pruning, and then further trained the pruned model on 10.5B tokens. We compare our pruned 780M model to the Mamba2 780M model trained from scratch on 300B tokens [6]. Given the short time limit of the rebuttal and limited compute, this was the largest model we could evaluate; however, we hope that these Mamba2 pruning results provide further insights into the generalizability of our method.\nAs shown in the table below, our compressed 780M model outperforms the model trained from scratch.\n\n| Benchmark        | Mamba2 780M | Mamba2 1.3B | Compressed 780M |\n|------------------|-------------|-------------|------------------|\n| arc_challenge    | 28.6        | 33.2        | 34.2             |\n| arc_easy         | 54.7        | 60.6        | 60.4             |\n| commonsense_qa   | 19.6        | 20.9        | 26.4             |\n| hellaswag        | 54.7        | 59.9        | 50.4             |\n| openbookqa       | 36.4        | 37.0        | 34.2             |\n| piqa             | 72.1        | 73.5        | 71.3             |\n| race             | 21.8        | 24.8        | 32.5             |\n| social_iqa       | 41.0        | 42.9        | 41.1             |\n| truthfulqa_mc2   | 38.1        | 36.1        | 38.9             |\n| winogrande       | 58.0        | 60.1        | 57.7             |\n| **Avg**          | **42.5**    | **44.9**    | **44.7**         |\n| **Tokens**       | **300B**    | **300B**    | **10.5B**        |\n\nWe will include these results in the supplementary materials.\n[6] Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through 320 structured state space duality. arXiv preprint arXiv:2405.21060, 2024.\n\n# Questions\nWe thank the reviewer for providing the related reference, and we will make sure to cite it in our manuscript under related works since it is directly related to making SSMs and hybrid models more efficient. \n## Reviewer: Group-Aware Head Permutation Constraints and Head Channel Consistency\nAs described in the paper, there are two constraints, both of which are directly due to the triton kernels implemented for mamba from https://github.com/state-spaces/mamba, which we use in our implementation.: \n\nConstraint 1: the constraints raised by the same channel permutation for all heads could be observed in multiple instances, a place where it is easy to follow is during the decoding of Mamba2. During token decoding, the latent SSM state, h, whose dimension is mamba head times mamba channels times SSM state dimension (m_h x m_d x d_s), is updated by Bx (equation 3 in the manuscript), and therefore, all the head channels in “x” should follow the same channel indexing in “h”.\n\nWe note that this constraint could be relaxed if one modifies all the operations to handle different rankings, but this needs further changes to the Triton kernel, and we leave it for future studies. To this point, while [1] cited by the reviewer is correct, it does not account for the implementation details of the Mamba kernels.\n\nAt the end of our response, we have also provided a simple snippet simulating a decoding step in SSMs, and showcase why inconsistent head permutation across heads can result in different outputs mathematically.\n\nConstraint 2: The constraints raised by groups that apply to the mamba heads, as shown in Figure 3 in the manuscript, is due to broadcasting of the SSM matrix “B”, whose dimension is the number of SSM groups times SSM state dimension (g x d_s), which is handled by replicating the d_s parameters for each group for the whole mamba heads assigned to that group in the input “x” matrix. Therefore, to allow heads to be permuted across different groups, one would need to implement a new Triton kernel to allow different broadcasting mechanisms. We leave this for future studies.\n\nWe would also like to note that, in many of the previous hybrid models, such as the one considered in reference [7], there is only one mamba group, so in that case, mamba heads can be permuted freely.\n\n```python # \nimport torch\nfrom einops import rearrange, repeat\nimport torch.nn.functional as F\nimport os\nimport random\nimport numpy as np\n\ndef set_all_seeds(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed) \n    np.random.seed(seed)  \n    torch.manual_seed(seed)  \n    torch.cuda.manual_seed(seed) \n    torch.backends.cudnn.deterministic = True  \n    torch.backends.cudnn.benchmark = False  \n\nseed = 42\nset_all_seeds(seed)\n\n#Toy example parameters\nd_state = 5\nd_inner_local=4*2\nngroups_local=2\nheaddim=4\nnheads=d_inner_local//headdim\n\nssm_state=torch.rand(nheads, headdim, d_state)\nx = torch.rand(d_inner_local)\n\nB = torch.rand(d_state*ngroups_local)\nC = torch.rand(d_state*ngroups_local)\ndt = torch.rand(nheads)\ndt_bias = torch.rand(nheads)\nA = torch.rand(nheads)\nD = torch.rand(nheads)\n\nB = rearrange(B, \"(g n) -> g n\", n=d_state) \nC = rearrange(C, \"(g n) -> g n\", n=d_state)\nB = repeat(B, \"g n -> (g h) n\", h=d_inner_local // ngroups_local)\nC = repeat(C, \"g n -> (g h) n\", h=d_inner_local // ngroups_local)\n\ndt = repeat(dt, \"h -> (h p)\", p=headdim)\ndt_bias = repeat(dt_bias, \"h -> (h p)\", p=headdim)\nA = repeat(A, \"h -> (h p) n\", p=headdim, n=d_state)\nD = repeat(D, \"h -> (h p)\", p=headdim)\n\n\ndef process(dt,dt_bias, A, D, x, ssm_state, B, C):\n    dt = F.softplus(dt + dt_bias.to(dtype=dt.dtype))\n    dA = torch.exp(torch.einsum(\"d,dn->dn\", dt, A))\n    \n    dB_x = torch.einsum('d,dn,d->dn', dt, B, x)\n    ssm_state.copy_(\n        ssm_state * rearrange(dA, \"(h p) n -> h p n\", p=headdim)\n        + rearrange(dB_x, \"(h p) n -> h p n\", p=headdim)\n    )\n    \n    y = torch.einsum(\n        \"dn,dn->d\",\n        rearrange(ssm_state, \"h p n -> (h p) n\", p=headdim),\n        C,\n    )\n    y = y + D * x\n    return y\n\n\n\nchannel_permutation = [1,0,3,2]\n\n# Change channel_permutation_for_two_heads to [1, 0, 3, 2, 4, 5, 6, 7] for inconsistent channel permutation to observe the differntt output\nchannel_permutation_for_two_heads = [1, 0, 3, 2, 5, 4, 7, 6]  # uncomment for CORRECT\n#channel_permutation_for_two_heads = [1, 0, 3, 2, 4, 5, 6, 7]  # uncomment for INCORRECT \n\nx_permuted = x[channel_permutation_for_two_heads]\nssm_state_permuted_channel = ssm_state[:,channel_permutation,:]\n\ny = process(dt,dt_bias, A, D, x, ssm_state, B, C)\ny_permuted = process(dt, dt_bias, A, D, x_permuted, ssm_state_permuted_channel, B, C)\nprint(abs(y_permuted[channel_permutation_for_two_heads]-y).sum()) # should be zero for correct and non-zero for incorrect permutation\n\n```"}]

Author rebuttal:

Meta-review: The paper seems to be one of the first in-depth studies on pruning hybrid SSM architectures, focusing on Mamba layers. The empirical results are quite strong: the authors demonstrate a 50% compression of the Nemotron-H 8B model down to 4B parameters, retaining most of the original model's accuracy. The reviewers noted extensive ablations presented in the paper. 

The most common concern was that the experiments were performed only on a single model, Nemotron-H 8B, making it difficult to assess if the method generalizes to other model families. The authors added new experiments during the rebuttal, which the reviewers found convincing. Another concern was model bias post-compression. Once again, additional experiments during the rebuttal were added on some safety benchmarks. Some clarity issues were also noted, which will hopefully be fixed in the camera-ready based on the provided feedback.
Decision: Accept (poster)