Technical Summary: AIRE-Prune

Problem Statement

Deep State Space Models (SSMs) have demonstrated efficacy in modeling sequential data by compressing input history into internal states. However, the state dimension $n$ remains a primary driver of computational and memory costs. While modern architectures often utilize diagonal or diagonalizable parameterizations (e.g., multi-SISO or MIMO stacks like S5) to ensure stability and efficiency, these models frequently suffer from over-parameterization. Existing post-training mechanisms to optimize $n$ are limited. Prior approaches, such as unstructured weight pruning, do not directly reduce the state dimension. Other structured methods, like the recent LAST (Gwak et al., 2025), rely on worst-case frequency-domain gain ($H_\infty$) to score states. While effective, worst-case measures can be overly conservative for typical workloads, as they emphasize peak amplification that tasks may rarely excite, potentially leading to suboptimal compression ratios.

Methodology: AIRE-Prune

The authors introduce AIRE-Prune (Asymptotic Impulse-Response Energy for State PRUN(E)ing), a structured, post-training, layer-adaptive pruning method designed for diagonal (or diagonalizable) SSMs. Unlike previous methods that focus on worst-case gain, AIRE-Prune aligns pruning with asymptotic impulse-response energy, a typical-case criterion representing the total output energy a state contributes over an infinite time horizon.

Core Mechanism

1. Energy-Based Ranking:
   For a stable discrete-time diagonal SSM layer defined by $(\Lambda, B, C)$, where $\Lambda = \text{diag}(\lambda_1, \dots, \lambda_n)$ and $|\lambda_i| < 1$, the method calculates the total squared output energy generated by a unit impulse input for each state (mode) $i$.
   The closed-form score for the $i$-th mode is derived as:
   $$\text{EnergyScore}_{\text{local}}(x_i) = E_i = \frac{\|C_{:,i}\|_2^2 \|B_{i,:}\|_2^2}{1 - |\lambda_i|^2}$$
   This score integrates three factors:

   • Controllability: Represented by $\|B_{i,:}\|_2^2$ (ease of excitation).
   • Observability: Represented by $\|C_{:,i}\|_2^2$ (ease of measurement).
   • Damping/Memory: Represented by the pole $\lambda_i$; states with poles closer to the unit circle (slower decay) contribute more to long-run energy.

2. Cross-Layer Normalization:
   Since energy scales can vary significantly across different layers (e.g., due to encoder/decoder gains), the authors propose a prefix-normalized scoring strategy to enable global comparison:

   • Within each layer $\ell$, modes are sorted by their local energy $E^{(\ell)}_i$ in descending order.
   • A prefix sum $S^{(\ell)}_{(i)}$ is computed for the sorted energies.
   • The global score is defined as the "hazard-rate" ratio:
     $$\text{AIREPrune}(x^{(\ell)}_{(i)}) = \frac{E^{(\ell)}_{(i)}}{S^{(\ell)}_{(i)} + \epsilon}$$
     This normalization ensures that the scores are monotonically non-increasing within a layer, allowing for a single global threshold $\tau$ to be applied across the entire model.

3. Pruning Execution:
   Given a target pruning ratio or threshold, the algorithm retains the longest prefix of states in each layer where the normalized score exceeds $\tau$. This allows for layer-adaptive pruning where some layers may be heavily pruned or even entirely removed if their states fall below the global threshold, while others remain largely intact.

Key Contributions

• Closed-Form Asymptotic Energy Score: The paper provides a mathematically derived, closed-form score for every learned mode based on infinite-horizon output energy. This avoids the need for expensive numerical simulations or retraining.
• Extension of Modal Truncation: The work extends classical modal truncation (typically applied to single linear systems) to deep stacks of nonlinear SSMs, preserving the diagonal structure essential for efficiency.
• Typical-Case vs. Worst-Case: By shifting the importance metric from worst-case $H_\infty$ gain to total impulse-response energy, the method targets the aggregate behavior of the system rather than rare peak amplifications.
• Architecture Agnosticism: The method is applicable to both multi-SISO constructions and MIMO stacks (e.g., S5-style), as well as other SSM families like S4D and Mamba.

Experimental Results

The authors evaluated AIRE-Prune on the Long Range Arena (LRA) benchmark and the Speech Commands dataset using S5, S4D, and Mamba backbones.

• Compression Efficiency: Across diverse benchmarks, AIRE-Prune achieved an average state pruning of 60.8% with an average accuracy drop of only 0.29% without any retraining.
• Comparison to Baselines:
  • On the LRA suite, AIRE-Prune significantly outperformed baselines including Random, Uniform Magnitude, Global Magnitude, LAMP, Uniform $H_\infty$, Global $H_\infty$, and LAST.
  • Under a strict constraint of $\le 1$ percentage point (pp) accuracy loss, AIRE-Prune allowed for substantially higher pruning ratios than the strongest baseline (LAST). For example, on the Path-X task, AIRE-Prune safely pruned 70% of states compared to 30% for LAST.
  • On the Speech Commands task, 45% of states were pruned with negligible accuracy loss (96.40% vs. 96.43% for the full model).
• Inference Efficiency: The pruning translated directly to system-level gains. On an NVIDIA H100 GPU, the method yielded inference speedups ranging from 1.23$\times$ to 2.86$\times$ and parameter reductions of 19% to 65%, depending on the task and pruning ratio.
• Generalization: The method demonstrated robust performance on S4D and Mamba architectures, achieving up to 90% pruning on the Text task for S4D and 80% on the Image task for Mamba with minimal accuracy degradation.

Significance and Claims

The paper claims that AIRE-Prune reveals substantial redundancy in the state spaces of trained SSMs, which can be removed without retraining. By aligning pruning with asymptotic response energy rather than worst-case gain, the method achieves aggressive compression with negligible cost to task performance.

The authors highlight that the "step-like" accuracy drop profile observed with AIRE-Prune (where accuracy remains stable until a high pruning threshold is reached, then drops sharply) indicates a strong separation between important and unimportant states. This contrasts with baselines that show smooth degradation, suggesting their scoring methods mix important and unimportant states.

Practically, the method enables:

1. Global Budgeting: Setting a single target budget without per-layer tuning.
2. Block-Level Reduction: The ability to prune entire layers when their states fall below the threshold, converting fine-grained sparsity into block-structured reductions that are more hardware-friendly.
3. Training-Free Optimization: A post-training mechanism that requires no additional fine-tuning, making it immediately applicable to deployed models.

The work concludes that AIRE-Prune offers a principled, energy-based approach to model order reduction in deep SSMs, surpassing current state-of-the-art pruning strategies in both compression ratio and accuracy preservation.