We thank the reviewer for their valuable insights and comments. We have tried our best to provide the responses for the questions as follows:

Weaknesses:

Reviewer: discussion of related work and stronger baseline methods

At 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.

Reviewer: the evaluation is conducted solely on Nemotron-H-8B-Base

To 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. As shown in the table below, our compressed 780M model outperforms the model trained from scratch.

We will include these results in the supplementary materials. [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.

Questions

We 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.

Reviewer: Group-Aware Head Permutation Constraints and Head Channel Consistency

As 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.:

Constraint 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”.

We 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.

At 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.

Constraint 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.

We 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.

import torch
from einops import rearrange, repeat
import torch.nn.functional as F
import os
import random
import numpy as np

def set_all_seeds(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed) 
    np.random.seed(seed)  
    torch.manual_seed(seed)  
    torch.cuda.manual_seed(seed) 
    torch.backends.cudnn.deterministic = True  
    torch.backends.cudnn.benchmark = False  

seed = 42
set_all_seeds(seed)

#Toy example parameters
d_state = 5
d_inner_local=4*2
ngroups_local=2
headdim=4
nheads=d_inner_local//headdim

ssm_state=torch.rand(nheads, headdim, d_state)
x = torch.rand(d_inner_local)

B = torch.rand(d_state*ngroups_local)
C = torch.rand(d_state*ngroups_local)
dt = torch.rand(nheads)
dt_bias = torch.rand(nheads)
A = torch.rand(nheads)
D = torch.rand(nheads)

B = rearrange(B, "(g n) -> g n", n=d_state) 
C = rearrange(C, "(g n) -> g n", n=d_state)
B = repeat(B, "g n -> (g h) n", h=d_inner_local // ngroups_local)
C = repeat(C, "g n -> (g h) n", h=d_inner_local // ngroups_local)

dt = repeat(dt, "h -> (h p)", p=headdim)
dt_bias = repeat(dt_bias, "h -> (h p)", p=headdim)
A = repeat(A, "h -> (h p) n", p=headdim, n=d_state)
D = repeat(D, "h -> (h p)", p=headdim)


def process(dt,dt_bias, A, D, x, ssm_state, B, C):
    dt = F.softplus(dt + dt_bias.to(dtype=dt.dtype))
    dA = torch.exp(torch.einsum("d,dn->dn", dt, A))
    
    dB_x = torch.einsum('d,dn,d->dn', dt, B, x)
    ssm_state.copy_(
        ssm_state * rearrange(dA, "(h p) n -> h p n", p=headdim)
        + rearrange(dB_x, "(h p) n -> h p n", p=headdim)
    )
    
    y = torch.einsum(
        "dn,dn->d",
        rearrange(ssm_state, "h p n -> (h p) n", p=headdim),
        C,
    )
    y = y + D * x
    return y



channel_permutation = [1,0,3,2]

# Change channel_permutation_for_two_heads to [1, 0, 3, 2, 4, 5, 6, 7] for inconsistent channel permutation to observe the differntt output
channel_permutation_for_two_heads = [1, 0, 3, 2, 5, 4, 7, 6]  # uncomment for CORRECT
#channel_permutation_for_two_heads = [1, 0, 3, 2, 4, 5, 6, 7]  # uncomment for INCORRECT 

x_permuted = x[channel_permutation_for_two_heads]
ssm_state_permuted_channel = ssm_state[:,channel_permutation,:]

y = process(dt,dt_bias, A, D, x, ssm_state, B, C)
y_permuted = process(dt, dt_bias, A, D, x_permuted, ssm_state_permuted_channel, B, C)
print(abs(y_permuted[channel_permutation_for_two_heads]-y).sum()) # should be zero for correct and non-zero for incorrect permutation

We thank the reviewer for their valuable insights and comments. We have tried our best to provide the responses for the questions as follows:

Weaknesses

Reviewer: Can the pruning method show similar effects on models other than Nemotron-H 8B?

To 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. As shown in the table below, our compressed 780M model outperforms the model trained from scratch.

We will include these results in the supplementary materials. [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.

Reviewer: Does the Mamba layer limit the range of LLMs that can be pruned?

Adding 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.

Reviewer: Comparison of the compressed model with the model of the same series with smaller parameters.

While 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.

Reviewer: Experiments and analysis of inference throughput

Inference 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).

We 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.

#Questions

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.

Please refer to our response above on the added study on compressing Mamba2 1.3b to 780m.

Reviewer: Experiments and analysis of inference throughput.

Please refer to our response in the final section addressing the identified weaknesses.

We thank the reviewer for their valuable insights and comments. We have tried our best to provide the responses for the questions as follows:

Questions

Reviewer: How knowledge distillation affects the performances? Results with pruning only before distillation can be ablated to showcase the effectiveness of proposed pruning techniques.

Thank 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.

We 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.

[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.

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.

To 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.

As shown in the table below, our compressed 780M model outperforms the model trained from scratch.

We will include these results in the supplementary materials.

[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.

We thank the reviewer for their insightful questions and discussions. We will try our best to address the questions and concerns below:

Weaknesses

Reviewer: it is not clear if the observation and analysis will hold true for different model families

To 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. As shown in the table below, our compressed 780M model outperforms the model trained from scratch.

We will include these results in the supplementary materials. [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.

Reviewer: model performance deteriorates after pruning

The 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.

Reviewer: ablation on pruning heuristic used in sec 2.2

We 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.

Reviewer: potential impact of pruning strategy proposed on the model bias

Thank 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.

Questions

Reviewer: How well does the proposed pruning strategy generalize for different family architectures?

Please refer to the Mamba2 1.3b to 780m compression discussed above

Reviewer: how pruning affects the model bias

Please refer to our answer on model safety and bias above.