Layer-Adaptive State Pruning for Deep State Space Models

TL;DR

This work addresses the inefficiency of large state dimensions in deep diagonal state-space models by introducing Layer-Adaptive STate pruning (LAST). LAST computes per-state H_infty-based importance scores and applies energy normalization to enable cross-layer, global pruning without retraining, while preserving stability via Hurwitz-parameterized diagonal SSMs. Across diverse benchmarks, LAST demonstrates substantial compressibility—averaging about 33% state reduction—with minimal accuracy loss (e.g., ~0.5%) in MIMO settings, and shows robust performance without retraining. The proposed approach provides a principled, transferable method to shrink state spaces in SSMs, enabling more efficient inference and training while maintaining stability guarantees.

Abstract

Due to the lack of state dimension optimization methods, deep state space models (SSMs) have sacrificed model capacity, training search space, or stability to alleviate computational costs caused by high state dimensions. In this work, we provide a structured pruning method for SSMs, Layer-Adaptive STate pruning (LAST), which reduces the state dimension of each layer in minimizing model-level output energy loss by extending modal truncation for a single system. LAST scores are evaluated using the $\mathcal{H}_{\infty}$ norms of subsystems and layer-wise energy normalization. The scores serve as global pruning criteria, enabling cross-layer comparison of states and layer-adaptive pruning. Across various sequence benchmarks, LAST optimizes previous SSMs, revealing the redundancy and compressibility of their state spaces. Notably, we demonstrate that, on average, pruning 33% of states still maintains performance with 0.52% accuracy loss in multi-input multi-output SSMs without retraining. Code is available at https://github.com/msgwak/LAST.

Figures (8)

• Figure 1: Illustration of LAST for two layers. Matrices are divided by lines on a per-state basis, and subsystems are sorted in descending order by their $\mathcal{H}_\infty$ norms. LAST scores are obtained by normalizing each $\mathcal{H}_\infty$ norm by the sum of all $\mathcal{H}_\infty$ norms in a layer when the states with lower $\mathcal{H}_\infty$ norms are excluded. Since LAST scores correlate with model-level output energy loss, we prune all parameters corresponding to states with low LAST scores.
• Figure 2: Efficiency-accuracy trade-off curves of pruned S5 models for tasks in LRA benchmark. LAST maintained accuracy better than other methods, Uniform $\mathcal{H}_\infty$ and Global $\mathcal{H}_\infty$ (LAST without energy normalization), demonstrating its superior ability to identify insignificant states.
• Figure 3: (Top) Evaluated state importance score and (Bottom) remaining state dimensions in an S5 model for Path-X task. The state indices are sorted by $\mathcal{H}_\infty$ scores, evaluated once for each conjugate pair.
• Figure 4: Remaining poles in $\mathbf{\Lambda}^{(6)}$ of an S5 model for Path-X task.
• Figure 5: Search space in the two-dimensional coefficient space for stability.