Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

AIRE-Prune: Asymptotic Impulse-Response Energy for State Pruning in State Space Models

Apurba Prasad Padhy, Fernando Camacho, Saibal Mukhopadhyay

TL;DR

AIRE-Prune introduces a post-training, layer-adaptive pruning method for diagonal and diagonalizable state-space models that minimizes long-run output-energy distortion. It assigns each mode a closed-form energy score $E_i = \frac{\|C_{:,i}\|_2^2 \|B_{i,:}\|_2^2}{1-\lvert\lambda_i\rvert^2}$ and normalizes these scores within each layer to enable global cross-layer pruning decisions, effectively extending modal truncation to deep SSM stacks. Empirically, it achieves substantial compression on S5-style SSMs with an average of $60.8\%$ state removal and only $0.29\%$ accuracy loss without retraining across the Long Range Arena and Speech Commands benchmarks, while delivering meaningful inference speedups ($1.2\times$ to $2.9\times$) and parameter reductions. The method outperforms baselines based on worst-case $H_\infty$ criteria by exploiting an energy-based (typical-case) perspective, and its layer-normalized, globally thresholded design yields robust head-tail separation and occasional whole-layer removals, enabling hardware-friendly, structured reductions.

Abstract

State space models (SSMs) often sacrifice capacity, search space, or stability to offset the memory and compute costs of large state dimensions. We introduce a structured post-training pruning method for SSMs -- AIRE-Prune (Asymptotic Impulse-Response Energy for State PRUN(E)) -- that reduces each layer's state dimension by directly minimizing long-run output-energy distortion. AIRE-Prune assigns every state a closed-form asymptotic impulse-response energy-based score, i.e., the total impulse-response energy it contributes over an infinite horizon (time), and normalizes these scores layer-wise to enable global cross-layer comparison and selection. This extends modal truncation from single systems to deep stacks and aligns pruning with asymptotic response energy rather than worst-case gain. Across diverse sequence benchmarks, AIRE-Prune reveals substantial redundancy in SISO and MIMO SSMs with average pruning of 60.8%, with average accuracy drop of 0.29% without retraining, while significantly lowering compute. Code: https://github.com/falcon-arrow/AIRE-Prune.

AIRE-Prune: Asymptotic Impulse-Response Energy for State Pruning in State Space Models

TL;DR

AIRE-Prune introduces a post-training, layer-adaptive pruning method for diagonal and diagonalizable state-space models that minimizes long-run output-energy distortion. It assigns each mode a closed-form energy score and normalizes these scores within each layer to enable global cross-layer pruning decisions, effectively extending modal truncation to deep SSM stacks. Empirically, it achieves substantial compression on S5-style SSMs with an average of state removal and only accuracy loss without retraining across the Long Range Arena and Speech Commands benchmarks, while delivering meaningful inference speedups ( to ) and parameter reductions. The method outperforms baselines based on worst-case criteria by exploiting an energy-based (typical-case) perspective, and its layer-normalized, globally thresholded design yields robust head-tail separation and occasional whole-layer removals, enabling hardware-friendly, structured reductions.

Abstract

State space models (SSMs) often sacrifice capacity, search space, or stability to offset the memory and compute costs of large state dimensions. We introduce a structured post-training pruning method for SSMs -- AIRE-Prune (Asymptotic Impulse-Response Energy for State PRUN(E)) -- that reduces each layer's state dimension by directly minimizing long-run output-energy distortion. AIRE-Prune assigns every state a closed-form asymptotic impulse-response energy-based score, i.e., the total impulse-response energy it contributes over an infinite horizon (time), and normalizes these scores layer-wise to enable global cross-layer comparison and selection. This extends modal truncation from single systems to deep stacks and aligns pruning with asymptotic response energy rather than worst-case gain. Across diverse sequence benchmarks, AIRE-Prune reveals substantial redundancy in SISO and MIMO SSMs with average pruning of 60.8%, with average accuracy drop of 0.29% without retraining, while significantly lowering compute. Code: https://github.com/falcon-arrow/AIRE-Prune.
Paper Structure (75 sections, 1 theorem, 40 equations, 5 figures, 5 tables)

This paper contains 75 sections, 1 theorem, 40 equations, 5 figures, 5 tables.

Table of Contents

  1. Introduction
  2. This work.
  3. Results.
  4. Related Work
  5. Model Order Reduction (MOR)
  6. Layer-Adaptive Pruning in Deep Networks
  7. State Pruning for Deep SSMs: LAST vs. This Work
  8. Background
  9. Stability of state space models
  10. Diagonal State Space Models (S5-style)
  11. High-level architecture.
  12. Parameterization.
  13. AIRE-Prune: Asymptotic Impulse-Response Energy for State Pruning in State Space Models
  14. State space model.
  15. Why energy as importance metric?
  16. ...and 60 more sections

Key Result

Lemma 1

For ${\bm{u}}\in\mathbb{C}^{m}$ and ${\bm{v}}\in\mathbb{C}^{n}$,

Figures (5)

  • Figure 1: AIRE-Prune: Asymptotic Impulse Response Local Energy Score
  • Figure 2: AIRE-Prune: Prefix-normalized global scoring across all the layers. (Yellow shade defines pruned states as they correspond to low magnitude score)
  • Figure 3: Comparison across baselines
  • Figure 6: Trade-off curves between pruning ratio and accuracy for pruned S5 models across tasks in the LRA benchmark. Baselines LAST, Uniform $H_\infty$, Global $H_\infty$ are refered from LAST
  • Figure 7: Layer-wise pruning ratio as we increase the global pruning threshold for S5 models across tasks in the LRA benchmark.

Theorems & Definitions (2)

  • Lemma 1: Frobenius norm of a rank-1 outer product
  • proof