Layer-Adaptive State Pruning for Deep State Space Models

We thank the reviewer for their feedback and suggestions based on their insightful paper summary. We appreciate their assessment that our method, which is closely coupled with the structure of DSSMs, is particularly interesting.

Weakness:

Since it is a pioneer work, more comparison method should be used, such as magnitude pruning.

We appreciate the reviewer's suggestion to include more comparison methods, such as magnitude pruning. As a common point, this has been thoroughly discussed in Response to Reviewer y832. In summary, we have newly defined magnitude-based state pruning methods and presented the extended experiment results to show clearer effectiveness and contributions. Notably, LAST outperformed other methods by a large margin, and all $\mathcal{H}\_{\infty}$ pruning methods consistently showed superior results over magnitude pruning methods.

Question:

What is the difference between pruning states and parameters?

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. In a DSSM layer, how the input and output are associated with a state is represented by the training parameters, which include the continuous-time system matrices $\mathbf{\Lambda}$, $\mathbf{B}$, and $\mathbf{C}$, which are discretized into $\overline{\mathbf{\Lambda}}$, $\overline{\mathbf{B}}$, and $\mathbf{C}$. The implementation of state pruning incorporates the identification of insignificant states and masking all parameters corresponding to the states, i.e., masking $\lambda_i$ from $\mathbf{\Lambda}$, the row vector $\mathbf{B}_i$ from $\mathbf{B}$, and the column vector $\mathbf{C}_i$ from $\mathbf{C}$. As compared in Section 4, this concept originates from model reduction studies for single SSMs, and LAST extends this to the deep-layered architecture with nonlinear functions.

To explicitly demonstrate the necessity of state pruning in SSMs, we compared the performance of unstructured random pruning and (structured) random state pruning using the same experimental setup in Response to Reviewer y832. For unstructured random pruning, we pruned randomly selected elements from the system matrices, obtaining the following results:

Despite the same amount of parameters being removed, it can be observed that the model completely loses its performance when pruning is performed in an unstructured manner.

This is because unstructured pruning can disrupt the integrity of the model's learned dynamics by altering subsystems into significantly different ones. In contrast, state pruning maintains the functionality of subsystems, leading to less performance degradation. This highlights the importance of considering the structure and relationships within the model when applying pruning techniques.

Y. He et al. Channel pruning for accelerating very deep neural networks. In The IEEE International Conference on Computer Vision, 2017.

We are glad the reviewer appreciated the novelty and effectiveness of the proposed method. As the reviewer suggested, we compare our method to three magnitude pruning methods that use magnitude-based criteria from traditional DNNs and state pruning granularity as in the proposed method.

State pruning is structured pruning, where all training parameters corresponding to the insignificant states are pruned based on a state importance metric. In traditional DNN, such as MLP, the transfer function of a layer is a weight matrix, $f(u)=Wu$, where each element $W_{ij}$ is a weight parameter. This allows evaluating the training parameters by their magnitude, $|W_{ij}|$, as the magnitude controls the input's influence on the output. Applying this to SSMs could involve unstructured magnitude pruning, which masks the high-magnitude elements of each system matrix according to the pruning ratio.

However, as discussed in Response to Reviewer xVFN, there is a fundamental disparity between unstructured and state pruning methods in SSMs, making it unfair to compare them. Thus, we define magnitude-based state pruning methods for SSMs and evaluate their ability to distinguish significant and insignificant states, thereby preserving overall model performance.

• Uniform Magnitude

  Every layer is pruned to have an identical pruning ratio by evaluating the importance of each state $i$ as $|\overline{\lambda}_i|||\mathbf{\overline{B}}_i||||\mathbf{C}_i||$. The pruning ratio is applied uniformly across all layers. While any $L_p$ norm can be used, we present the results using the $L_2$ norm as an example.

• Global Magnitude

  The same state importance metric as in Uniform Magnitude is used, but the comparison group is extended from intra-layer to inter-layer, ensuring the pruning ratio is met globally across the entire network.

• Layer-Adaptive Magnitude-based Pruning (LAMP)

  This method employs a metric of $\frac{|\overline{\lambda}_i|^2||\mathbf{\overline{B}}_i||^2||\mathbf{C}_i||^2}{\sum\_{j\geq i}|\overline{\lambda}_j|^2||\mathbf{\overline{B}}_j||^2||\mathbf{C}_j||^2}$ adapted from Lee et al. [2021], which originally used $\frac{W_i^2}{\sum\_{j\geq i} W_j^2}$ for a real-valued weight parameter $W$. The state indices in the denominator are assumed to be ordered based on their evaluation using the basic magnitude criteria.

Extending Table 1, we report newly conducted JAX experiments with 3 different seeds on an NVIDIA RTX 3090 as follows:

The table lists the accuracy loss of pruning methods for S5 models (Smith et al. [2023]), averaged across all 10 tasks. Values in parentheses exclude non-compressible tasks with zero pruning ratios and accuracy loss.

At the same pruning ratio, LAST significantly outperforms other methods by exhibiting the least accuracy loss.

Random state pruning results show that all methods can identify insignificant states, but all $\mathcal{H}\_{\infty}$ pruning methods surpass magnitude pruning methods. This implies the suitability of $\mathcal{H}\_{\infty}$ pruning for SSMs, which can be explained by the transfer functions of SSMs. In an SSM layer, the multiplicative mapping from input to output in the frequency domain is given by $\mathbf{Y}(z)=\mathbf{C}(zI-\mathbf{\overline{\Lambda}})^{-1}\mathbf{\overline{B}}\mathbf{U}(z)$, where $\mathbf{U}$ and $\mathbf{Y}$ are the $\mathcal{Z}$-transforms of input and output sequences $\mathbf{u}$ and $\mathbf{y}$, and $\mathbf{C}(zI-\mathbf{\overline{\Lambda}})^{-1}\mathbf{\overline{B}}$ is the transfer function matrix.

The importance of training parameters in SSMs can be evaluated by their influence in the frequency domain. Using the $\mathcal{H}\_{\infty}$ norm, which measures the maximum gain of the transfer function in the frequency domain, $\mathcal{H}\_{\infty}$ pruning optimizes the system by minimizing the maximum gain alteration, leading to better performance than magnitude pruning. Specifically, $\overline{\lambda}_i$'s importance in $\overline{\mathbf{\Lambda}}$ is measured by $(1-|\overline{\lambda}_i|)^{-1}$ in $\mathcal{H}\_{\infty}$ pruning, while magnitude pruning measures $|\overline{\lambda}_i|$. Additionally, LAST normalizes this importance by the total energy transmission of the layer, preventing extensive pruning of high energy-conveying layers and enabling stable global comparison and pruning.

J. T. Smith et al. Simplified state space layers for sequence modeling. In The International Conference on Learning Representations, 2023.

J. Lee et al. Layer-adaptive sparsity for the magnitude-based pruning. In The International Conference on Learning Representations, 2021.

We appreciate the reviewer's assessment that the paper contributes to new areas of DSSMs and the experiments are thorough to show the robustness and effectiveness of the proposed method. The reviewer's concerns were the practical impacts and related works of reducing state dimension.

Please use some statistics to show the importance of state dimension.

We present the average evaluation step speed and peak GPU memory usage when the state dimension is reduced to the maximum pruning ratio that resulted in less than 1% accuracy loss using our proposed method. The evaluation was conducted by implementing an actual decrease in model size on an NVIDIA RTX 3090. While there is some variance depending on the channel size of each model per task, reducing the state dimension improves efficiency in both computational and memory costs.

If the state dimension is significant, can you provide more related works on reducing the state dimension, maybe those related works are not from the perspective of pruning to reduce the state dimension, but I think it is meaningful to discuss them.

To alleviate the issue of large state dimensions, related research has evolved with two primary approaches: architectural and algorithmic.

1. Architectural approaches

One architectural approach is modulating states for parallelized inference. Specifically, an $h\cdot n_{SISO}$-th order system is modulated by $h$ systems, each with $n_{SISO}$ states, where $h$ is the number of channels. This means each $n_{SISO}$-th order system handles information from a specific channel. To address the inability to integrate information across channels, a channel mixing layer is introduced. This is referred to as, in our paper, the multi-SISO architecture (Gu et al. [2021], Gu et al. [2022a]), allowing for training $h\cdot n_{SISO}$ effective state spaces with reduced computational burden by parallelizing into $n_{SISO}$, thus enabling the development of many subsequent SSMs (Gupta et al. [2022], Gu et al. [2022b]).

In comparison, another architecture proposed by Smith et al. [2023], referred to as MIMO architecture, uses $n_{MIMO}\ll h\cdot n_{SISO}$ states to determine the output for a channel based on information from all input channels, making the mixing layer unnecessary. This parameterization allows for high performance with much smaller state dimensions than equivalent block systems in multi-SISO layers. For instance, in the Path-X task, which involves the longest sequences among LRA tasks, this model shows state-of-the-art performance. However, both architectures lack optimization methods for model size, leading to computational inefficiencies when the task does not require such a high model capacity.

• Relation to the pruning approach

  Our state pruning method can be applied to all these architectures, effectively reducing the model size to fit task or resource requirements while preserving the performance. It identifies and adaptively prunes the states with the least impact on performance based on layer-wise energy transmission. In multi-SISO architectures, since each state is assigned to a specific channel, the difference in importance among states may not be significant. However, our method still shows superior retention compared to other methods. In MIMO architectures, the difference in state importance becomes more pronounced, and the accuracy loss caused by pruning varies significantly depending on how a method evaluates the state importance. Consequently, our method effectively preserves the most significant states and demonstrates superior performance compared to other pruning methods, as shown in the table in Response to Reviewer y832.

2. Algorithmic approaches

One recent study has shifted from learning the element of decomposition of the transfer function in DSSMs to directly learning the coefficients of the transfer function (Parnichkun et al. [2024]). The motivation behind this approach is also to address the computational load caused by large state dimensions. The primary contribution of this approach is making computations even independent of the state dimension. However, as discussed in Appendix A, this approach has a weakness of only guaranteeing stability at initialization, unlike DSSMs, which can consistently guarantee stability due to their direct parameterization for the poles of the transfer function. Therefore, as a branch of DSSM research, our work contributes to enhancing the efficiency of DSSMs, which is dependent on the state dimension, by minimizing the state dimension.

A. Gu et al. Combining recurrent, convolutional, and continuous-time models with linear state space layers. In Advances in neural information processing systems, 2021.

A. Gu et al. Efficiently modeling long sequences with structured state spaces. In The International Conference on Learning Representations, 2022a.

A. Gupta et al. Diagonal state spaces are as effective as structured state spaces. In Advances in Neural Information Processing Systems, 2022.

A. Gu et al. On the parameterization and initialization of diagonal state space models. In Advances in Neural Information Processing Systems, 2022b.

J. T. Smith et al. Simplified state space layers for sequence modeling. In The International Conference on Learning Representations, 2023.

R. N. Parnichkun et al. State-free inference of state-space models: The transfer function approach, In International Conference on Machine Learning, 2024.