AIRE-Prune: Energy-Based Pruning for SSMs

• AIRE-Prune is a pruning methodology that assigns each state an impulse-response energy score to quantify its long-run contribution within deep SSM layers.
• It employs a closed-form computation combined with cross-layer prefix normalization, enabling global thresholding and effective pruning decisions.
• Experimental results show that AIRE-Prune achieves significant parameter reduction and inference speedups with minimal accuracy loss across various benchmarks.

AIRE-Prune (Asymptotic Impulse-Response Energy for State PRUN(E)) is a structured, post-training pruning methodology for state space models (SSMs) that directly targets the long-run output-energy distortion associated with reducing the hidden state dimensionality of each layer. By assigning to every state a closed-form energy score reflecting its total contribution to the model’s asymptotic impulse response, and then normalizing these scores across layers, AIRE-Prune enables global, cross-layer pruning decisions that retain accuracy and stability while dramatically reducing parameter counts and compute requirements. This approach extends classical modal truncation, typically applied to single linear time-invariant (LTI) systems, to deep and stacked SSM architectures, emphasizing typical-case, rather than worst-case, response characteristics (Padhy et al., 31 Jan 2026).

1. Formal Definition of the Asymptotic Impulse-Response Energy (AIRE) Score

Consider a single discrete-time, linear, time-invariant state-space layer expressed as

$x[k+1] = A\,x[k] + B\,u[k],~~ y[k] = C\,x[k]$

where $A \in \mathbb{C}^{n \times n}$ is diagonalizable and typically taken as diagonal $A = \mathrm{diag}(\lambda_1,\ldots,\lambda_n)$ with $|\lambda_i|<1~\forall i$ (ensuring asymptotic stability). The unit impulse response at lag $k$ is $H_k = C A^k B$. The contribution of each state, or mode $i$, to the output is given by

$H_k = \sum_{j=1}^n C_{:,j}\,\lambda_j^k\,B_{j,:}$

The total asymptotic (infinite-horizon) output energy of mode $i$ is

$$E_i = \sum_{k=0}^{\infty} \|C_{:,i}\,\lambda_i^k\,B_{i,:}\|_F^2 = |C_{:,i}|^2\,|B_{i,:}|^2\,\sum_{k=0}^{\infty} |\lambda_i|^{2k} = \frac{|C_{:,i}|^2\,|B_{i,:}|^2}{1 - |\lambda_i|^2}$$

This closed-form expression, derived from properties of geometric series and the layer’s diagonal structure, quantifies each mode’s steady-state impulse-response energy.

2. Score Normalization and Cross-Layer Comparability

Impulse-response energies $E_i$ can vary greatly in scale across different layers or even within channels of a single layer, primarily due to variations in encoder/decoder gain. To enable a single global pruning threshold, AIRE-Prune applies the following normalization procedure for each layer $\ell$ with $n_\ell$ modes:

1. Compute the local energies $\{E_i^{(\ell)}\}_{i=1,...,n_\ell}$.
2. Sort in descending order $$E_{(1)}^{(\ell)} \geq E_{(2)}^{(\ell)} \geq \dots \geq E_{(n_\ell)}^{(\ell)}$$.
3. Compute prefix sums $S^{(\ell)}(i) = \sum_{j=1}^i E_{(j)}^{(\ell)}$.
4. Define the prefix-normalized (hazard-rate) score for the $i$-th ranked mode: $$s_{(i)}^{(\ell)} = \frac{E_{(i)}^{(\ell)}}{S^{(\ell)}(i) + \varepsilon}$$ with $\varepsilon > 0$ for numerical stability. Since $S^{(\ell)}(i)$ is increasing, $s_{(i)}^{(\ell)}$ decreases with $i$, concentrating high scores in large-energy (head) modes and low scores in the energy tails. Collecting all $\{s_{(i)}^{(\ell)}\}$ across layers yields a global, cross-layer ordering for thresholding.

3. Algorithmic Implementation and Complexity

AIRE-Prune operates post-training on a stack of $L$ SSM layers, each with diagonalized parameters $(A^{(\ell)}, B^{(\ell)}, C^{(\ell)})$ and state dimension $n_\ell$, using a global pruning ratio $p$ or score threshold $T$. The algorithm proceeds as follows:

1. For $\ell=1...L$, compute $E_i^{(\ell)}$ for each mode via the closed-form above.
2. In each layer, sort $E_i^{(\ell)}$, compute prefix sums, and calculate per-mode scores $s_{(i)}^{(\ell)}$.
3. Aggregate all scores into a global list of length $N = \sum_\ell n_\ell$.
4. For a target pruning ratio $p$, set $K_\mathrm{tot} = \lceil (1-p)N \rceil$ and select $T$ as the $K_\mathrm{tot}$-th largest global score; or, directly use a threshold $T$.
5. In every layer, retain the largest prefix of sorted modes with $s_{(i)}^{(\ell)} \geq T$, pruning the rest.
6. Output kept/pruned indices for all layers.

The computational complexity is:

• Closed-form score computation: $O(\sum_\ell n_\ell)$.
• Sorting per layer: $O(\sum_\ell n_\ell \log n_\ell)$.
• Global aggregation: $O(N \log N)$.

Relative to the cost of a forward or backward pass in deep SSMs, this overhead is negligible (Padhy et al., 31 Jan 2026).

4. Comparison to Classical Modal Truncation and Related Methods

Classical modal truncation for a single diagonal LTI system ranks states by per-mode $H_\infty$ norms $g_i = |C_{:,i} B_{i,:}|/(1 - |\lambda_i|)$, seeking to bound the worst-case gain. Balanced truncation (Gramian-based) provides strong worst-case error guarantees but destroys diagonal structure, making it less suitable for deep, stacked SSMs. Methods like LAST apply an $H_\infty$ metric, with layer-wise normalization for global pruning, promoting a conservative, worst-case-oriented approach.

AIRE-Prune fundamentally diverges by employing an $H_2$-type (total impulse-response) energy $E_i$, focusing on typical-case rather than worst-case responses. Its closed-form, per-mode, mode-separable scores eliminate the need to solve large Lyapunov equations, and prefix normalization naturally clusters low-energy modes for contiguous block pruning. Empirically, AIRE-Prune achieves significantly higher prune ratios at matched accuracy drops, suggesting that real-world sequential tasks rarely excite the worst-case system frequencies (Padhy et al., 31 Jan 2026).

5. Experimental Evaluation

Extensive experiments were conducted on benchmarks including Long Range Arena variants and Speech Commands, pruning multiple SSM variants (S5, S4D, Mamba/S6) in a one-shot fashion (with no retraining). For S5 on LRA, results averaged over all tasks are:

• Mean pruning ratio: 60.8%.
• Mean accuracy drop: 0.29 percentage points.

Within a ≤1pp accuracy loss budget, safe prune ratios were:

• ListOps: 20%
• Text: 80%
• Retrieval: 50%
• Image: 65%
• Pathfinder: 80%
• Path-X: 70%

Inference speedups (NVIDIA H100) ranged from 1.2× for low pruning to 2.9× for aggressive pruning; parameter-count reductions spanned 19% to 65%. Across all backbone architectures tested, AIRE-Prune delivered superior safe prune budgets and lower accuracy loss compared to baselines (Random, Magnitude, LAMP, Uniform/Global $H_\infty$, LAST), confirming its effectiveness and versatility (Padhy et al., 31 Jan 2026).

6. Practical Guidelines and Implementation Notes

Key operational recommendations for using AIRE-Prune include:

• One-shot pruning eliminates the need for post-pruning retraining or fine-tuning.
• Plotting accuracy versus global threshold $T$ enables "elbow search" to identify thresholds just prior to rapid accuracy decline, maximizing compute efficiency for given performance.
• At aggressive thresholds, discrimination by prefix-normalized score can identify and remove entire low-energy layers, simplifying deployment.
• The method is architecture-agnostic, applicable to multi-SISO, diagonal MIMO, bidirectional, and gated SSMs.
• The normalization parameter $\varepsilon$ should be small but not negligible (e.g., $10^{-6}$).
• Although grounded in $H_2$-type energy, an $H_\infty$-style worst-case error bound can be derived from the sum of pruned energies and corresponding pole magnitudes (see Appendix B.5 of (Padhy et al., 31 Jan 2026)).
• After identifying kept modes, the model should be rebuilt at reduced state dimension for optimal realization of speed and memory gains.

7. Summary and Impact

AIRE-Prune extends modal truncation concepts from classical control to the context of deep, multi-layer SSM stacks, shifting the emphasis from worst-case gain to long-run average energy via an $H_2$ criterion. Its use of closed-form, mode-separable pruning scores and a simple global normalization and thresholding rule permits post-training compression with minimal accuracy loss and measurable system-level speedups. Empirical evidence demonstrates that AIRE-Prune efficiently reveals and excises significant redundancy in SSM-based sequence models across multiple domains, supporting its adoption for practical memory and compute optimization in state space architectures (Padhy et al., 31 Jan 2026).