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

ICLR 2026 Conference Submission21955 Authors

Keywords: State Space Models, Pruning, Impulse Response, S5, State Pruning

TL;DR: State Pruning of Deep State Space Model (S5) using a closed form scoring method inspired by Asymtotic Impulse Response Energy.

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)ing ) — 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.

Primary Area: other topics in machine learning (i.e., none of the above)

Submission Number: 21955