Layer-Adaptive State Pruning for Deep State Space Models

We appreciate the reviewer's careful review and further feedback on our rebuttal. We agree that comparing the proposed method with other pruning methods beyond magnitude-based pruning could provide a clearer validation of our work. While magnitude-based pruning is straightforward, it is an effective and strong baseline, as it continues to be widely used in studies on iterative pruning, pruning schedules, or pruning on large language models (Sun et al. [2024], K. Xu et al. [2023], M. Gupta [2024]). Notably, LAST and $\mathcal{H}\_{\infty}$ pruning significantly outperform it, highlighting its necessity in SSM. Additionally, the previous comparison can be viewed as one for one-shot pruning methods, which are particularly useful when dealing with scaled models compared to iterative methods. In response to the remaining reviewer's concern, we present the following results for a scaling-based pruning method using batch normalization statistics as another baseline. **Uniform Scaling-based Pruning** (Liu et al. [2017]) This method uses the scale $\gamma^{(l)}\in\mathbb{R}^h$ trained in the $l$th batch normalization layer as the criterion for channel pruning out of $h$ channels. Specifically, we used $\gamma^{(l-1)}$ to prune $\mathbf{B}^{(l)}$ and $\gamma^{(l)}$ to prune $\mathbf{C}^{(l)}$ in consideration of the input-output processing scheme of SSM. | Method | Average pruning ratio | Average accuracy loss $\downarrow$ | Importance | | - | - | - | - | | Uniform scaling-based | 33.00 (36.67) | 25.84 (28.71) | $\gamma_k^{(l-1)}$ and $\gamma_k^{(l)}$ for channel $k$ | Due to time constraints, we could not test pruning methods that involve retraining, such as using a designed loss function or tracking gradient flow. We promise to include comparison experiments with other well-designed pruning methods in the final version of the paper. Additionally, although we considered recent feature statistics-based pruning methods, we noticed that further research is needed, especially in observing specific patterns in feature vectors within SSM. As we are the first to explore pruning in SSM, designing new criteria for state pruning is akin to suggesting a new algorithm. We hope the reviewer considers our work a foundational study that suggests appropriate pruning criteria and directions for SSM. --- M. Sun et al. A simple and effective pruning approach for large language models. In *The International Conference on Learning Representations*, 2024. K. Xu et al. Efficient joint optimization of layer-adaptive weight pruning in deep neural networks. In *Proceedings of the IEEE/CVF International Conference on Computer Vision*, 2023. M. Gupta, et al. Is complexity required for neural network pruning? a case study on global magnitude pruning. In IEEE Conference on Artificial Intelligence, 2024. Z. Liu et al. Learning efficient convolutional networks through network slimming. In *Proceedings of the IEEE international conference on computer vision*, 2017.

Thank you for taking the time to read our rebuttal. Due to space constraints, we had to distribute our responses across several rebuttals. We appreciate your thorough review of all the details. ### **Correction of misunderstanding** We apologize for any confusion caused by our use of 'disparity' and for not clearly referencing the relevant response to the reviewer's question. To clarify, we used 'disparity' to describe the performance gap between unstructured pruning and state pruning (a type of structured pruning), which is evident in our response. As shown in the previous rebuttal, structured pruning (state pruning) outperforms unstructured pruning even in randomly pruned cases, which clearly shows the effect of using state pruning granularity by excluding the other factors that might affect the performance. Thus, what we intended to deliver to Reviewer y832 was that **the presented table contains only state pruning**, omitting unstructured pruning due to the significant performance difference between unstructured/structured versions of every pruning method. In another way, our intention may be better conveyed by replacing 'fundamental' with 'trivial.' Therefore, that part does not conflict with the explanation that state pruning is a type of neural network pruning but is tailored for SSM. ### **Additional comments to the previous rebuttal** The equations used in our explanation are based on the standard form of system matrices in SSM, consistent with the notation in our submission. To assist the reviewer's understanding, we define key terminology and a detailed illustration of state pruning. Pruning granularity refers to the specific pattern or structure of parameters to be pruned, such as individual weights, entire neurons, channels, or states. For example, channel pruning prunes all parameters corresponding to an insignificant channel. Likewise, state pruning prunes all parameters corresponding to an insignificant state. For a DSSM layer's training parameters $\mathbf{\Lambda}$, $\mathbf{B}$, and $\mathbf{C}$, such that $x_{k+1}=\mathbf{\Lambda}x_k+\mathbf{B}u_k$ and $y_k=\mathbf{C}x_k$, the specific parameters for a state $i$ are the $i$th diagonal element $\lambda_i$ of $\mathbf{\Lambda}$, the $i$th row vector $\mathbf{B}_i$ of $\mathbf{B}$, and the $i$th column vector $\mathbf{C}_i$ of $\mathbf{C}$. For instance, if the 2nd state is identified as insignificant in a DSSM layer with the state dimension of 4, by masking the insignificant parameters to 0, the system matrices before and after pruning would be as follows: $x_{k+1}=\begin{bmatrix} \*&&& \\\ &0&& \\\ &&\*& \\\ &&&\*\end{bmatrix}x_k + \begin{bmatrix} \*&\*&\*&\* \\\ 0&0&0&0 \\\ \*&\*&\*&\* \\\ \*&\*&\*&\* \end{bmatrix} u_k$ $y_{k}=\begin{bmatrix} \*&0&\*&\* \\\ \*&0&\*&\* \\\ \*&0&\*&\* \\\ \*&0&\*&\* \end{bmatrix}x_k$ In this sense, we would like to reiterate that we responded to the question 'What is the difference between pruning states and parameters?' by 'As in channel pruning (He et al. [2017]), the term of state pruning represents its pruning granularity. That is, the difference arises in the unit of parameters being pruned at once, and the act of pruning training parameters follows the same approach as traditional pruning methods'.