---
title: 'AIRE-Prune: Energy-Based Pruning for SSMs'
url: https://www.emergentmind.com/topics/aire-prune-asymptotic-impulse-response-energy-for-state-prun-e
type: topic
---

# AIRE-Prune: Energy-Based Pruning for SSMs

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 [2602.00534].

## 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 [2602.00534].

## 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 [2602.00534].

## 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 [2602.00534].

## 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 [2602.00534]).
- 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 [2602.00534].

Source: https://www.emergentmind.com/topics/aire-prune-asymptotic-impulse-response-energy-for-state-prun-e