AIRE-Prune: Asymptotic Impulse-Response Energy for State Pruning in State Space Models
Authors: Apurba Prasad Padhy
Fernando F. Camacho
Saibal Mukhopadhyay
Venue: ICLR 2026 Poster (2026)
Keywords: State Space Models
Pruning
Impulse Response
S5
State Pruning
Primary area: other topics in machine learning (i.e., none of the above)
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.

Published as a conference paper at ICLR 2026
AIRE-PRUNE: ASYMPTOTIC IMPULSE-RESPONT EN- ERGY FOR STATE PRUNING IN STATE SPACE MODELS
Apurba Prasad Padhy', Fernando Camacho² & Saibal Mukhopadhyay'
Georgia Institute of Technology, Atlanta, USA
2 Laboratory for Physical Sciences, College Park, Maryland, USA
a apadhy apadhy 9egatech edu, saibal mwkhopadhyay@eee. gatech.
2 fercamacho@l 1 ps. und. edu
ABSTRACT
State space models (SSMs) often sacrifice capacity, search space, stability offset the memory and compute costs large state dimensions. We introduce a structured post-training pruning method for SSMs AIRE- PrRneP (Asymptotic Impulse- Response Energy for State PRUN(E)ing ) - that reduces each layer's s state dimension by directly minimizing long-run output-energy distortion. AiRE- Prune assigns every state a closed-for form asymptotic impulse-respon score, i.c., the total resposs-response energy contributes over an infinite hori- zon (time), and normalizes these scores layer-wise to enable global cross-layer comparison and selection. This extends modal truncation from single systems 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 0.29% without retraining while significantly lowering compute. Code will be released: https://gihub..co//alccoo waw.:////wo////-/---.P. Prune
energy based
INTRODUCTION
Deep state space models (SSMs) have proven effective in modeling sequential data by optimally compressing input history into internal states (Gu et al., 2020 2021 2022c Gu & Dao 2023 Zhang et al. 2023 Parnichkun al. 2024). Alongside these advances, persistent challenge to train SSMs efficiently and stably without divergence. Building on classical linear systems the-
ory (Kallthh 1980), recent work has developed booolitt-guuranteeingg parameterizations (Gu et al.) 2022b), general architectural blueprints (Smith al., (023), and frequency-domate implementa-
tions that leverage transfer functions and FFT's for throughput (Gu et al. 2022a Zhang et al. 2023
Parnichkun et al. 2024).
A central driver of computation and memory in SSMs is the state dimension n. Since carly pro- posals, multiple single-input single-output (multi-SISO) construction has been widely adopted for scalable training: many small SISO subsystems are learned in parallel and then mixed through chan- nel projections [Gu et al., 2022b,a; Gu and Dao. 2023; Zhang et al., 2023; Parnichkun al., 2024]. Within this setting, diagonal (or diagonalizable) systems have been shown to match the accuracy of more general non-diagonal systems while remaining efficient [Gupta et al., 2022]. In paral- lel, multi-input multi-outpur (MIMO) SSMs (e.g.. S5-style stacks) explicitly exploit multi-channel structure, often attaining strong performance with smaller than equivalently sized multi-SISO blocks; this advantage is especially evident on long-horizon tasks such as Path-X (Smith et al.) 2023 Parnichkun et al. (2024) Despite these developments, both multi-SISO and MIMO fam- ilies typically lack post-training mechanisms to optimize leading to over-parameterrizaton and unnecessary inference cost.
Several orthogonal directions attempt reduce complexity including unstructured weight pruning or structured state pruming strategies. This work focuses on structured state pruning as translates to real computation reduction with reduced dimensionality. Parameterizing transfer functions for SISO systems can enable state-free inference, but tends to restrict the search space or guarantee stability only initialization (Parnichkun et al. 2024). Frequency-domain kemels accelerate training and

---

Published as a conference paper at ICLR 2026
inference but do not directly resolve redundancy in the learned state space (Gu et al. 2022a Gupta
et al. 2022 Zhang et al. 2023). This motivates post-training model order reduction (MOR) in deep SSMs: Identify and remove states that minimally affect the task output while preserving accuracy.
Layer-adaptive pruning provides a practical MOR route. Recent work (LAST) (Gwak et al. 2025) proposed score each learned subsystem by worst-case, frequency- then normalize scores layer-wise to enable global, cross-layer selection and pruning with bounded output distortion. Empirically, this revealed that trained SSMs are often highly compressible, achiev- ing sizable state reductions with minimal accuracy loss on long-range sequence benchmarks (e.g.. LRA and Speech Commands). While powerful, worst-case measures can be conservative for typical workloads, as they emphasize peak amplification that tasks may rarely excite.
gain (an Ho view),
This work. We introduce AIRE-Prune (Asymptotic Impulse- Response Energy for State PRUN(E)ing), a structured, post-training, layer-adaptive method for pruning states in deep state space model. The contribution of this work are as follows:
Energy-based ranking: ranks each state (mode) by its total output energy over infinite time when excited under a umit impulse input.
assigns every learned mode a closed-form total output energy score, enabling fast, principled importance ordering.
Closed-form per-mode energy score:
Cross-layer normalization global pruning score: normalizes per-layer energy scores to common scale, enabling global, cross-layer comparison and selection of the least- energetic (insignificant) states.
This extends classical modal truncation from single systems to deep stacks and aligns pruning with asymptotic impulse response energy- a typical-case criterion-rather than a worst-case gain ap- proach. AIRE-Prune is architecture-annotic (covering multi-SISO and MIMO SSMs).
Results. On Long Range Arena with S5-style MIMO SSMs, AIRE-Prune prunes on average 60.8% of states with only 0.29% average accuracy degradation without retraining. These find- ings indicate that state spaces in trained SSMs contain substantial removable redundancy and that aligning pruning with total output energy yields strong compression negligible cost to task per- formance, surpassing worst-case frequency gain-based approaches.
2
RELATED WORK
2.1 MODEL ORDER REDUCTION (MOR)
MOR approximates high- digendimensi linear systems by lower-order models with controlled error, with applications in VLSI (Antoulas & Sorensen (001). power systems (Li & White 1999). and PDE discretizations (Jones Kerrigan| 2010; Curtain & Zwart 2012). Modal truncation removes states from diagonal realization using aiquenccydooaall criteria that limit Ho distortion (Green & Limebeer 2012). Balanced truncation constructs a realization in which states are simultancously controllable and observable, then discards the least energetic ones, offering strong
guarantees and many variants (Jones & Kerrigan 2010) Curtain & Zwart 2012). However, the similarity trans- forms central balanced truncation eliminate the diagonal structure exploited by modern SSMs for efficiency and stability. Our stance: we preserve diagonal parameterization and per-state granular- ity (as in modal truncation) and extend the setting from single linear system deep stacks with nonlinearities encountered in contemporary sequence models.
2.2
LAYER-ADAPTIVS PRUNING IN DEEP NETWORKS
Allowing different layers to prune by different amounts (layer-adaptive pruning) (Morcos et al., 2019 Han et al. 2015 Mocanu et al.) 2018 8 Lee al. 2021 Xu et al. (003) generally outperforms uniform ratios (Zhu Gupta 2017] Gale et al. 2019). Approaches include layerwise magnitude thresholds/targets and global criteria that compare scores across layers under single budget. For example, provides a Frobenius norm bound on worst-case l, distortion when pruning one layer (oth-
frozen), and extends this to joint optimization over layerwise ratios, Because SSMs are governed
2

---

Published as a conference paper at ICLR 2026
by transfer functions and dynamical couplings not only static weights ron-magnitide impor- tance criterion is essential. Our stance: we adopt the layer-adaptive paradigm but base importance on total output energy, yielding closed-form, mede-separable scores that respect controllability, ob- servability, and damping.
2.3 STATE PRUNING FOR DEEP SSMs: LAST VS. THIS WORK
LAST (Gwak et al. 2025) ranks states via a worst-case frequency- domain measure: per-state Ho normalized to produce global cross-layer ranking under model-level budget, giving a peak-gain bound on output distortion and showing trained SSMs are highly compressible. This work (AIRE-Prune): we retain per-state, layer-adaptive pruning but use infinite-horizon output energy (total impulse-response energy) as a typical-case criterion that emphasizes long-run expen- diture through wtw.://liiiiitoooeeaaaiiillaaaaa
This energy view yields closed-form, layer- normalized, globally comparable scores and, in practice, enables aggressive compression with neg- ligible degradation in accuracy. Refer appendix[B.c for detailed mathematical difference.
BACKGROUND
3.1
STABILITY OF STATE SPACE MODELS
DT LTI SSM asymptotically stable iff its poles lie strictly inside the unit circle, i.c., p(A) < 1. Directly enforcing this during training is nontrivial: constraining parameters to a fixed stable set (e.g., disk/polygon) guarantees stability but can reduce expressivity, while "center-of-stability" ini- tialization helps early training yet offers guarantee thereafter. Diagonal (or diagonalizable) SSMs mitigate this by parameterizing poles directly: we use a diagonal DT realization with conjugate pairs for real I/O. When derived from CT diagonal poles with R(AŞ*) O, codor hold at step A, gives X = exp(X,* A;), hence | < 1. This diagonal parameterization preserves stability by construction and enables classical analyses (Lyapunov, modal formulas) within deep SSM layers.
3.2 DIAGONAL STATE SPACE MODELS (S5-STYLE)
High-level architecture. An S5 layer processes length-T sequence uo:T-1 R² through di- agonal, linear üime-invrrian (LTI) core followed by a pointwise nonlinearity and channel mixing. As in prior diagonal SSMs, the stack comprises: quence of ..1+nnnnlyyyrrtt blocks, and decoder tailored to the downstream task. Each block can be implemented either as multi-SISO assembly (one diagonal SISO system per channel) or single MIMO system acting jointly on the h channels. We adopt a unified MIMO description whose effective state dimension per layer is
encoder that lifts inputs to h channels,
n = ng h for multi-SISO (n,, states per channel)
1mm
for MIMO (nm shared states)
so that multi-SISO is structured special case of MIMO with block-diagone couplings.
Parameterization. We model each diagonal SSM layer (state size n; channels h) in continuous time (CT) by a=t) Ax(t) + Bu(t) and yt) Cxx() (t) + Du(t), where A diag(A)..... A,))
Caxe is diagonal, B E Caxh C' Chxn and D E haxxh"" To preserve real inputs/outputs, complex parameters appear conjugate pairs (Gu et al. 20226). Following diagonal SSM practice, each mode i carries its step size A, > 0 (collected into A IR",, so zero-order hold (ZOH)
discretization yields a diagonal DT system
0i++ Ar + Buss
yk = Cas + Dui,
(2)
where,
= diag((>):))))
B I)B,
(3)
where the Ba expression holds elementwise and extends continuously at X, 0. Stability is en- forced at the CT level (Hurwitz), i.., R(A,) < 0, which implies |ex:Ai| < 1 and thus DT con- tractivity. A pointwise nonlinearity wraps the linear core form the layer output, fo (,;;;; Σ)
3

---

Published as a conference paper at ICLR 2026
State Space Equations
Impuise Response =
IBTIC
xit+1] Ax(t] + Bult]
E 2
(1- 12,1)
ytt
Cxtt]
Sorted by
A, C Parameters
h
7. 0.6
Energy Score
uttl/((t] Input/Output x(t]
SSM Layer
Hidden State
C'
ult
hy(tl
2 0.4
Impulse Signal
Asymptotic Impulse Response
Figure 1: AIRE-Prune: Asymptotic Impulse Response Local Energy Score
of (Cox + Dug). optionally combined with residual and normalization layers. This diagonal pa- rameterization admits closed-form per-mode discretization, simple stability control, and efficient kernel/scan implementations, while remaining applicable to both multi-SISO (n = nggh) and MIMO (n nmm) instantiations under unified notation.
AIRE-PRUNE: ASYMPTOTIC IMPULSE- IMPULSE-RESPONSE ENERGY FOR STATE PRUNING IN STATE SPACE MODELS
State space model. We begin with a discrete-time (DT) diagonal state space layer
3k+1 = Are + Bus:
yk = CCak
(4)
where diag(Ar ....., An) with Ag1 < (all poles strictly inside the unit circle). This spectral condition is the DT counterpart of asymptotic stability in control: it guarantees that perturbations decay and that linear system quantities defined as infinite sums over time are finite. Diagonal (or diagonalizable) parameterizations make stability explicit through pole variables {}:} and enable mode-wise analysis.
Why energy as importance metric? In linear time-invariant (LTI) systems, "importance" has canonical meaning: how much output energy direction (state) can transmit from inputs to outputs. For stable discrete-time systems this energy is literally an area under a curve, over time (sum of squared impulse response) or over frequency (integral of squared transfer magnitude). This section develops that basic connection and, from it, derives simple state-importance score and a practical pruning rule. An impulse (or white noise) excites all frequencies equally, so the measured energy reflects aggregate amplification (realizable to real input variations), not single resonant tone.
For total (steady-state) output energy corresponding to state, consider stable discrete-time SSM layer Σ (, C') with zero initial state and an impulse input un = ô&o- The impulse response is
H C.
and the output is yı: Hg. The (squared) steady-state output energy is the discrete-time Hz energy
1
| energy =
|[G(
=
HH = Σ
ICAB
(5)
2m Jo
k=0
k=
frequency domain
time domain
where the transfer function is G(z) C(zI - A) and Parseval's identity yields the equality of time- and frequency-domain energies. This admits two equivalent interpretations that connect directly to control and signal processing: (i) impulse energy: the total squared output produced by unit impulse, accumulated from & = 0 to oo; and (ii) long-run output power under unit white noise: the steady-stat variance of yı when uy I) (by Parseval's 's identity). Both views rely |&l < to ensure that the series equation equation converges which is the case for the diagonal SSMs considered here.
4

---

Published as a conference paper at ICLR 2026
Before Global
Normalization
Post Global
A
B
C
Layer 1
E, (L)
State Pruning
AIRE-Prune
Σ
E
(L)
Layer
High Magnitude Scores
Figure 2: AIRE-Prune: Prefix- normallzedl global scoring across all the layers. (Yellow shade defines pruned states as they correspond to low magnitude score)
Per-mode energy via geometric progression (finite horizon). For single diagonal mode E,
(),, Bi.., C..)) with | < 1, the rank-1 impulse slice is H ()) C. X Bi..- Using |/w I [ ||2122 for rank-1 outer products,
H I = |Ai I B
(6)
Therefore, the (truncated) output energy contributed by mode over horizon T is the geometric sum
T-1
1 - A 22'
EA(T) = H = ||
B
Σ t=
A a
C. 12 B
1 1 2
1=
Stability of the state space model implies |A + 0 as T + oc. Taking the limit in equation gives the per-mode steady-state output energy which is equivalent to the impulse response energy.
C B 12
F3
lim (T( T
and the total layer energy can be approximated as an additive sum across modes (Appendix A. 1.2.
Σ
E.
emergy
Pruning set P of states removes their modal responses. approximately the sum of their individual contributions:
The resulting change in layer energy is
Teeneryy
||_P
energy
ΣE E, PEP
(9)
Thus, to minimize steady-state distortion at the layer level, one should remove the smallest-energy modes. We therefore define the energy (AIRE-Prune local) score for state z; as
||C,,
B 11 A1
EnergyScore E,
(10)
Special cases. If B is fixed and row-normalized (common in practice), then || = and E |C,,,,//(1111))) For bidirectional layers, |[C.,,ll is replaced by the average the forward and backward contributions.

---

Published as a conference paper at ICLR 2026
Method
State importance
Accuracy Loss vs Pruning Ratio
Random
Random
Uniform magnitude
A
B
|
Average Accuracy Loss (%)
Uniform megnitude
Global magnitude
Global magnitude
B
|
10*
Uniform Hinf
LAMP
1
[
|
Hinf
LAMP
AST
Σ
12 B B 1 111111
AIRE-Prune
|
|
Uniform Ho.
(1 x2
10a
|| |B,
Global Hoc
(1 X
CC.1 B.
35
50
60
LAST
(1-1A.
Average Pruning Ratio (%)
Σ Egsi
|C ||B,11
(1 (1-1A,1)²
(b) Accuracy loss vs. pruning ratio.
CC, B.
1- 11
AIRE-Prune (ours)
Σ Ji
|C ||
1-1
(a) State-importance definitions
Figure 3: Comparison across baselines
Intuition: Why this score is good importance measure. The score in Eq. equation 10 (a) mode-separable (no large matrix solves), (b) scale-aware through the B/C couplings, and (c) dynamics-awant through the pole damping. A state with larger |B:ell (more controllable, i.e., eas-
ier to excite), larger |C.,illz ll (more observable. i.e., easier to measure), or larger [A,l (longer mem- ory) has larger E; and therefore induces a larger accuracy or variance drop if removed. This aligns pruning with steady-state distortion and empirically enables aggressive compression with negligible accuracy loss.
From local to global: layer normalization. Per-mode asymptotic energies E can differ in scale across layers (e.g., due to encoder/decoder gains), therefore we need to normalize individual layers in such a way to have common global threshold for state pruning. The per-layer scales can differ substantially,
first sort the states of each layer by their Per-mode Asymptotic Energy in
descending order. Let E.(E) denote the i-th largest per-mode asymptotic energy in layer e, and define the prefix sum
5/0 =
Σ
( S
We then use the prefix-normalized score (AIRE-Prune)
(
(())
AIREPrund
(
(6)
(11)
s(1)
(())
+
with small è > 0 for numerical stability. This "hazard-rate" ratio is monotonically non-increasing in i, enabling an elbow-style, layer-adaptive rule: given a single global threshold т, keep the longest prefix in each layer for which Scorepreńx and prune the remaining contiguous tail.
Summary. Stability ensures that infinite-horizon energy is well-posed; energy decomposes addi- tively by modes in diagonal SSMs; the closed-form per-mode energy equation I0 gives an inter- pretable, dynamics- and scale-aware local importance; and simple layer normalization equation []] turss these into globally comparable scores for cross-layer sensitive pruning. Refer Appendix for mathematical analysis of the above.

---

Published as a conference paper at ICLR 2026
Algorithm 1: AIRE-Prune: Asymptotic Impulse-Responet Energy for State Pruning
Input: SSM Layers {e = 11..... L with parameters (A(9) (B(), C(O)); prune ratio p.
Output:
Kept-index sets { and pruned sets (P(O).
Step 1: Per-mode asymptotic energy F.
for + to do
Extract modal triplets [(1)) B 1 with 12,91 < ( 12 B (8) 12
Step 2: Per-layer sorting and prefix sums.
for L +- to L do
Sort modes by E in descending order to get E(()
E (mz)
(1) 2 ... 2...
and their indices.
Compute prefix sums: ( t -Ejs
(t) for = 1 ....
E (i)
Compute AIRE-Prumm scores: (()
where - 0
+
(6)
Step 3: Global selection
11 prune ratio p given then
12 B + Etm 1 Th -1 (1-p)
13 Sort AIRE-Prune in decending order to get global list
L {(5())
(1) l, 1) 1 S I., S i S nz).
14 Find the value equal to the B-th largest score in & Define the kept count per layer by
k + max i: (0) hT}
1s Step 4: Materialize keep/prune indices.
for f +- to L do
L K(() original indices of the top ke sorted modes; p(() +- Pruned indices.
18 return (X((),, {P(0)
EXPERIMENTS
5.1
SETUP
Models and tasks. We evaluate pruning the S5 (MIMO) SSM (Smith et al. 2023 across the six Long Range Arena (I.RA) tasks (Tay et al. 2021 and the Speech Commands 35-way keyword recognition benchmark (Warden 2018) (sequence length 16,000). All runs single NVIDIA H100 (40GB/80GB) GPU with the released S5 training and inference contiguration. Unless noted otherwise, report one-shot pruning (no retraining) and evaluate accuracy by freezing all parame- ters.
Baselines. We compare against six pruning baselines which includes random, uniform magnitude, global magnitude, LAMP (Lee al. 2021) Uniform Hoc (Gwak et al. 2025). Global Ho Gwakk
et al. 2025), LAST (Gwak et al. 2025). Figure 3a shows the state importance metric (scoring method) for the baselines.
Additional architectures. Beyond S5, we also evaluate AIRE-PRUMM on S4D and Mamba (S6) backbones test generalization across SSM families. Task-wise pruning/accuracy numbers for these models on LRA are summarized in Table
Pruning ratios. For methods that allocate layerwise budgets (Global Ho., AIRE-Prunc), we report average pruning ratios across layers. We sweep pruning ratios from (0% to 100%) with step size

---

Published as a conference paper at ICLR 2026
ListOps: Accuracy vs, Pruning Ratio
Accuracy Pruning
Retrieval: Accuracy Pruning
50
Accuracy (N)
80
Accuracy 180)
Aucuracy
20
60
22901
50
50
Bardon
40
Pruning ratio
ratio
Image: Accuracy vs. Pruning Ratio
Pothfinder: Accurocy vs. Pruning Rotio
PathX: Accuracy vs. Pruning Ratio
100
Accuracy (%)
Accuracc
Recureg
80
80
8
60
60
R00000 EURSS
20
40
50
80
100
20
20
60
80
100
20
40
00
00
100
Pruning ratlo 196)
Pruning ratio (%)
Pruiigg natio (56)
Random
Uniform
Global
LAST
AIRE-Prune (ours)
Figure 4: Trade-off curves between pruning ratio and accuracy for pruned S5 models across tasks in the LRA benchmark. Baselines LAST, Uniform Ho.. Global Ho are refered from (Gwak et al. 2025
Table 1: Task-wise accuracy LRA and Speech Commands with S5. Pruning ratio (Prun.) and post-pruning accuracy (Ace.). Numbers for Full/Uniform /Goball HLALTT (Gwak et al. 2025) are compared with AIRE-Prune (ours).
ListOps (2,048)
Text (4,096)
Retrieval (4,000)
Image
Pathfinder (1,024)
Path-X (16,384)
Speech (16,000)
(1,024)
Model
Prun.
Acc.
Prun.
Acc.
Prun. Acc.
Prun. Acc.
Prun.
Acc.
Prun.
Acc.
Prun.
Acc.
$5 Full
0%
61.48
0%
88.88
0%
91.20
0%
87.30
0%
95.15
0%
98.41
0%
96.43
Uniform Ho
0%
61.48
60%
82.49
50%
90.29
30%
86.45
30%
71.38
30%
90.90
20%
96.20
Global Hoc
0%
61.48
60%
88.56
50%
90.93
30%
87.04
30%
57.20
30%
69.21
20%
96.21
LAST
0%
61.48
60%
88.52
50%
90.42
30%
86.34
30%
94.45
30%
97.95
20%
96.31
AIRE-Prune
(ours)
20%
61.48
80%
88.24
50%
90.11 11 65 %
87.30
80%
95.15
70%
98.41
45 %
96.40
of 10% for Uniform Hoc Gwak et al. 2025), Global Ho (Gwak et al. 2025 LAST (Gwak al. 2025) and AIRE-Prune (our work) as in Figure conjugate pair per layer.
A ratio of 100% would leave a single complex-
5.2 ANALYSIS
Long Range Arena (S5). We evaluate AIRE-Prune on the LRA suite (Tay et al. 2021). which probes long-range dependencies with sequence lengths from 1,024 to 16,384. In all runs apply a single global threshold to the reefi-nnormalied AIRE scores and prune without any retraining. Table reports the post-pru post-pruning accuracies.
Speech Commands (S5). On Speech Commands, AIRE-PRUMS achieves 45% state pruning while essentially preserving accuracy (96.40% vs. 96.43% for the full S5 model; Table [ This mirrors the high-compressibility regime observed on several LRA tasks and indicates that AIRE's energy-based ranking transfers cleanly from synthetic range benchmarks to real audio classifi- cation.
Aggregate outcome. Averaged across tasks, AIRE-Prune removes 60.8% of states with only 0.29 pp accuracy drop (Figure 3b) last row). In the common "< 1 pp loss" regime, compressibility is task dependenn::::: Text, Pathfinder, and Path-X sustain 80%, 80 and 70% pruning, respectively, while matching or nearly matching full accuracy (Text: 88.24 vs. 88.88; Pathfinder: 95.15 vs. 95. 95,15; Path-X: 98.41 vs. 98.41; Table Tp. Retrieval and Image tolerate 50-65% % pruning at < pp loss (Re-
8

---

Published as a conference paper at ICLR 2026
Task-wise accuracy on LRA with S4D and Mamba. Pruning ratio (Prun.) and post- pruning accuracy (Acc.). Numbers for Full/Uniform /Globll LLST (Gwak et al. 2025) (S4D) are compared with AIRE-Prune (ours).
Table 2:
ListOps (2.048)
Text (4,096)
Retrieval (4,000)
Image (1,024)
Model
Prun.
Acc.
Prun.
Acc.
Prun.
Acc.
Prun.
Acc.
S4D Full
56.42
86,40
90.46
0%
77.02
Uniform ooc
10%
55.82
80%
86.02
60%
89.87
0%
77.02
Global Hoo
10%
49.95
80%
86.20
60%
89.84
0%
77.02
LAST
10%
56.27
80%
85.95
60%
89.46
0%
77.02
AIRE-Prunc (ours)
35%
56.07
%
85.83
76.32
Mamba Full
38.02
82.98
0%
72.14 14
0%
69.82
LAST
-%
20%
81.32
45%
71.37
45%
69.13
AIRE-Prune (ours)
45%
81.19
60 %
71.45
80%
68.97
Table Performance (S5): pruning ratios and resulting inference speedups and parameter-memory reductions for S5
Metric / Dataset +
ListOps
Text
Retrieval
Image
Pathfinder
Path-X
Pruning ratio
20%
80%
50%
65%
80%
70%
Inference speedup (x)
1.23 x
2.49
1.82
x.93
2.82
2.86x
Model reduction (%)
19.3%
58.7%
38.3%
45.3%
64.7%
64.2%
trieval: 90.11 vs. 90.93; Image: 87.30 vs. 87.30). ListOps is sensitive, admitting only 20% pruning before degradation; this operating point, accuracy is preserved (61.48 61.48).
Comparison to baselines. At high compression (right half of each panel in Fig. 4. AIRE-PRUNE dominates Uniform/Global Hos magnitude rules and outperforms LAST under the <1 pp loss crite- rion. This is mirrored in the average accuracy loss of Figure 3b) AIRE-PRUNE (0.29 pp) improves over Uniform Hox (4.32 pp), Global Hoe (7.51 pp), and LAST (0.52 pP), while achieving much higher average pruning ratio (60.8% % vs. 33% for those methods; ~= 1.8x higher).
Generalization to S4D and Mamba. Table 2 shows that AIRE-PRUNE also yields competitive pruning/accuracy trade-off's on S4D and Mamba. On S4D, AIRE attains up 90 % pruning on Text and 40% on Image while closely tracking full-model accuracy. On Mamba, AIRE safely prunes 45% of states on Text, 60% on Retrieval, and 80% on Image with negligible degradation. These results suggest that the asymptotic energy ranking is not tied to specific SSM parameterization and can serve as drop-in pruning rule across modern state-space architectures.
Superiority over state of the art. Against the strongest baselines (Uniform/Global Hoc and LAST), AIRE-PRUNE delivers strictly larger safe no-retrain budgets (defined at < 1 accuracy drop) on five of six LRA tasks while matching or exceeding accuracy. From Table [ AIRE's safe pruning exceeds the best baseline by +20 pp on ListOps, +20 pp on Text, +35 pp on Image, +50 pp on Pathfinder, and +40 pp on Path-X; x; on Retrieval, the listed AIRE point 50% misses the < pp rule by only 0.09 pp. yielding the best baseline. We attribute this margin to AIRE's time-accumulated energy lens and prefix nor- malization, which together create strong head-tail separation and enable single global threshold to keep essential modes intact and collapse low-energy layers, a behavior not observed magnitude- or peak-gain-basee based scoring.
advantage at that specific sample yet remaining competitive with
Inference efficiency and model size. We next quantify the system-level impact of pruning-by- removal. For each LRA task, we transfer the kept states into a ceamension-reducee S5 model and measure inference throughput and parameter reduction tasks, AIRE's s pruning schedules yield x22x229 x speedups and 19% 90%6555% % reductions in parameter count, with larger gains high-pruning regimes such Pathfinder and Path-X. This confirms that
NVIDIA H100 GPU (Table Across

---

Published as a conference paper at ICLR 2026
AIRE Prun-Prreeeee (Ust
AiRte-Prune (MDB)
AIRE Prune (Retrievall
Kegt maIn
Kegt Sbeeee
Host
AIRE-Prune (CIFAR-10)
AIRE-Prune (Pathfinder)
AIRE-Prune (PathX)
Hagt
0%
20%
40%
60%
80%
Figure 5: Layer-wise pruning ratio as we increase the global pruning threshold for S5 models across tasks in the LRA benchmark.
aggressive state pruning translates into meaningful computational and memory savings, not just abstract sparsity. Similar analysis for Mamba has been mentioned in Appendix C.3
Accuracy drop vs. pruning ratio. Figure 4 exhibits a pronounced elbow/step for AIRE-Prune accuracy remains near the full model up to high pruning threshold, then drops sharply; in contrast, baselines degrade smoothly as pruning increases. We interpret the step as evidence that AIRE's s rank- ing separates important from unimportant states with a wide margin, while smooth curves indicate that pruning mixes both groups throughout the sweep. Takeaway. The step-like profile is robust only when using AIRE energies and prefix normalization, which together keep the high-margin head in- tact and prune contiguous tails. Baselines show smooth decay because their scores offer weaker head-tail separation, so valuable states are progressively removed throughout the sweep.
Layer-wise pruning profile. Figure shows kept-state percentages per layer as the global ratio increases, revealing task-specific structure. On Path-X and Pathfinder, AIRE progressively thins later layers and, at high ratios, can remove entire layers whose states fall below the global threshold, turning fine-grained sparsity into block-level deletions. This indicates that our method can help achieve low-latency SSMs, as we are able to prune layers without losing accuracy. For Text and Retrieval, earlier layers retain larger fractions while mid/late layers contribute most of the budget, consistent with encoder-style front-end processing. Image exhibits more even shrinkage until its elbow, after which few layers collapse. ListOps shows non-trivial contribution from all layers even at small ratios, explaining its early elbow and our conservative cap around 20%. This heterogeneous, layer-seled seleetveeeetive behavior contrasts with reported layer pruning profile of LAST (Gwak et al. 2025) which tends to have more uniform pruning among layers at similar global budgets.
6 CONCLUSION
We introduced AIRE-PRUNE, a training-free pruning criterion for diagonal (or diagonalized) state- space models that ranks modes by asymptotic output energy, combined with prefix normalization that puts scores from different layers on common scale. With single global threshold and no retraining, AIRE-PRUNE achieves strong compression/accuracy trade-offs across the LRA bench- mark.
Summary of findings. Averaged over tasks for S5, AIRE-PRUNE removes 60.8 % of states while incurring only 0.29 percentage points of accuracy loss. Under a <1 pp loss budget, safely prunes 80% on Text and Pathfinder, 70% on Path-X, 65% on Image, 50% on Retrieval, and even 20% on the often-labeled "incompressible" ListOps. Across the board, AIRE-Prum dominates across all the baselines.
Practical implications. Because scores are computed once and applied globally, practitioners can: (i) set a target budget without per-layer tuning; (ii) sweep to the elbow and stop just before it to maximize savings at fixed accuracy; and (iii) realize system-level gains by deleting entire layers whenever their states fall below threshold, converting fine-grained sparsity into stoctestr reductions that are easier deploy hardware.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Athanasios C. Antoulas and Danny C. Sorensen. Approximation of large-scale dynamical systems: An overview. International Journal of Applied Mathematics and Computer Science, 1195):1093-
2001.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milng et al. Jax: composable transformations of python+ numpy programs. 2018.
Ruth F. Curtain and Hans Zwart. An Introduction to Infinite-Dinnn Siinitt-DDimesioonall Linear Systems Theory, volume 21 of Texts in Applied Mathematics. Springer Science Business Media, 2012.
Trevor Gale, Erich Elsen, and Sara Hooker. The state of sparsity in deep neural networks. arXiv preprint arXiv: 1902.09527444 2019.
M. Green and D. J. N. Limebeer. Linear Robust Control. Courier Corporation, 2012.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv 00752
preprint arXiv:23/2.00752 2023. URL. U
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. In Advances Neural Information Processing Systems, volume 33, pp. 1474-1487, 2020.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Com- bining recurrent, convolutional, and continuous-timg models with linear state space layers. In Advances in Neural Information Processing Systems, volume 34, pp. 572-585, 2021.
Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initial- ization of diagonal state space models. In Advances in Neural Information Processing Systems, volume 35, pp. 35971-35983, 2022a.
Albert Gu, Karan Goel, and Christopher Ré. Selective state spaces for long sequences. In Interna- tional Conference on Learning Representations (ICLR), 2022b.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2022c. URL https: //openreview. net./ wttp://wwwww...eeee.teeeeeee.eee www.rrrmm.......olo.oooo..
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces. In Advances in Neural Information Processing Systems, volume 35, pp. 22982- 22994, 2022.
Minscon Gwak, Scongrok Moon, Joohwan Ko, and PooGyeon Park. Layer-adaptive state pruning for deep state space models. In Proceedings of the 38th International Conference on Neural Information Processing Systems, NIPS '24, Red Hook, NY, USA, 2025. Curran Associates Inc. ISBN 9798331314385.
Song Han, Jeff Pool, John Tran, and William Dally. Learning both weights and connections for efficient neural networks. In Advances Neural Information Processing Systems, volume 28, 2015.
B. I. Jones and C. Kerrigan. When is the discretization of spatially distributed system good enough for control? Automatica, 46(9): 14e2-1462614.... 2010.
Thomas Kailath. Linear Systems. 01 1980.
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Tech- nical report, University of Toronto, 2009.
Lee, S. Park, S. Mo, Ahn, and J. Shin. adyeriad sparsity for the magnitude- -bssdd pruning. In International Conference Learning Representations (ICLR), 2021.
11

---

Published as a conference paper at ICLR 2026
J.-R. Li and J. White. Efficient model reduction of interconnect via approximate system grami- ans. In Proceedings the the 1999 IEEE/ACM Interaational Conference on Computer-Aided Design (ICCAD). pp. 380-383. IEEE, 1999.
Drew Linsley, Junkyung Kim, Vijay Veerabadran, Charles Windolf, and Thomas Serre. Learning long-range spatial dependencies with horizontal gated recurrent units. In Advances in Neural Information Processing Systems, volume 31, 2018.
Andrew Maas, Raymond Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng. and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pp. 142-150, 2011.
Decebal Constantin Mocanu, Elena Mocanu, Peter Stone, Phuong H. Nguyen, Madeleine Gibescu, and Antonio Liotta. Scalable training of artificial neural networks with adaptive sparse connec- tivity inspired by network science. Nature Communications, 9(1):2383, 2018.
Ari S. Morcos, Haonan Yu, Michela Paganini, and Yuandong Tian. One ticket to win them all: Generalizing lottery ticket initializations across datasets and optimizers. In Advances Neural Information Processing Systems, volume 32, 2019.
Nikita Nangia and Samuel R. Bowman. Listops: A diagnostic dataset for latent tree learning. arXiv preprint 1804 06002,, 2018. URL httpt:tt//// /ttp:s////avvv.eea
1804. 06028
R. N. Parnichkun, Stefano Massaroli, Andrea Moro, J. T. Smith, Ramin Hasani, Mathias Lechner, Qian An, Christopher Ré, Hajime Asama, Stefano Ermon, al. State-free inference of state- space models: The transfer function approach. arXiv preprint arXiv:2405.06147 2024. URL https: //tp://aaivvvvvi.a wwtp/////////.....o/
Dragomir R. Radev, Paramveer Muthukrishnan, and Vahed Qazvinian. The acl anthology net- work corpus. In Min-Yen Kan and Simone Teufel (eds.), Proceedings of the 2009 Workshop on Text and Citation Analysis for Scholarly Digital Libraries, pp. 54-61, Suntec City, Singapore, August 2009, Association for Computational Linguistics. URL attpss:/aclantholl. orq/w09-36077
J. T. Smith, Adam Warrington, and Scott W. Linderman. Simplified state space layers for se- quence modeling. In International Conference on Learning Representations (ICLR), 2023. arXiv:2208.04933.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sehastian Ruder, and Donald Metzler. Long range arena: benchmark for efficient transformers. In International Conference Learning Representations (ICLR), 2021.
Pete Warden. Speech commands: A dataset for limited-vocabulary speech recognition, 2018. URL https: org/ 188800400 .33209
Xu, Z. Wang, X. Geng, M. Wu, X. Li, and W. Lin. Efficient joint optimization of layer-adaptive weight pruning deep neural networks. Proceedings of the IEEE/CVF International Confer- ence on Computer Vision (ICCV), pp. 17447-17457. 2023.
M. Zhang, K. K. Saab, M. Poli, T. Dao, K. Goel, and C. Ré. Effectively modeling time series with simple discrete state spaces. The International Conference on Learning Representations (ICLR), 2023. /tLntttss://ppeeeeeeee http///epeeiieeeeeee/emeeeeeeeee
:
n forum? dimmmiid-3313 d=3mR 3
Michael Zhu and Suyog Gupta. To prune, or not to prune: Exploring the efficacy of pruning for model compression. arXiv preprint 17..0:1710.01878, 2017.
12

---

Published as a conference paper at ICLR 2026
A APPENDIX
A.1 MATHEMATICAL PROOFS
A.1.1
VECTOR/MATRIX INNER PRODUCTS AND NORMS
Euclidean space. For e C the standard inner product and induced 2-norm are
y) == x: y, ||ح||| = ax..
The Cauchy-Schwaad inequality gives (ax, y) S ||0:|2 1lyll2-
Frobenius (Hilbert-Schmidd) structure. For A, B E Cm the Frobenius inner product and norm are
== | == (T)
Σ |as (ass")
This norm is: (i) unitarily invariant ( AV = |||e for unitary U, V), (ii) compatible with (
vectorization ( (||A|| F [[voc( 222 and (iii) submultiplicative with respect the spectral norm:
||AB|| I ||A l B|| and || | I |||| B
Rank-1 outer products. (wv") )ig
Given C and C", the outer product uv" C has entries
0e U.
A.1.2 IDENTITY FOR RANK-1 TERMS
Lemma (Frobenius norm of a rank-l outer product), For u E " and v e C",
|
||
|| 10
Proof.
Uuv 11 = Σ lus Uzf uy12
Σ |u
Σ (151²) = 1|u112 |10
i.j
A.1.3 "ENERGY" INTERPRETATION OF 1
For an LTI system with y, = Σx20 Hy 4 and zero-mean white input u, ~ CO 4x) inde- pendent across I, the instantaneous output power contributed by the k-lag kernel
E H 1 = E Err(", H. x = tn H H E E x1) tr H = H
Thus, || Hy 1h- equals the expected output energy (power for unit-variance inputs) contributed by the k-th impulse slice. The same holds componentwise for H (i)
|| (
( quantifies the energy ( |P when the rank-1 terms |HHl = tr(H; Hy)
carried by the i-th mode at lag l. Moreover, since || HH = Σ, H
are mutually orthogonal in the Frobenius inner product, or more generally
always, this norm provides natural, additive energy accounting across modes and lags.
A.1.4 CONNECTION TO THE H2 NORM (TOTAL ENERGY)
For stable MIMO L.T system with transfer matrix G(z) and impulse sequence (H.).
| 2
tr G(G du Σ
2n
by Parseval/Plancherel (discrete-timi) Hence the Frobenius squarred of impulse slices sums to the total output energy for white inputs. This makes H, natural energy density over lags.
13

---

Published as a conference paper at ICLR 2026
A.1.5 REAL VS. COMPLEX DATA; BIDIRECTIONALITY
If the model uses complex pairs to represent real-valued dynamics, terms appear conjugate pairs whose sum is real. Use v (not in complex algebra; in purely real settings, o* = v For bidi- rectional SSMs, one can form an augmented LTI with block-diagonal forward/backward dynamies; the energy accounting above applies componentwise and adds.
B FROM ENERGY-BAR -BASED PRUNING TO WORST-CASE CERTIFICATES
B.1 SCOPE AND PROMISE (WHAT THIS APPENDIX DELIVERS)
What AIRE-Prune proposes. An energy-based post-training, layer-adaptive pruning rule for diagonal (or diagonalizable) state-space layers that uses a single global threshold.
What add here. A precise, worst-case (Hx) error certificale that complements the original rypical-case (energy/H2) rationale, and an end-to-end distortion bound for resid- ual stack with Lipschitz wrappers.
Why it matters. You keep the same AIRE ranking/threshoiding. but now with a com- putable, auditable guarantee; we also position the result against LAST (Gwak et al. 2025) and highlight the mathematical differences.
B.2 BACKGROUND: DIAGONAL SSM LAYERS AND ENERGY (NO PRIOR KNOWLEDGE ASSUMED)
Diagonal/diagonaaiized SSM layer. Let #x+1 Ar Bug and yn Car, with A diag( ; An). [A,] < for stability. The impulse response and frequency response are
Σ C. 1 AGe-sw
H, = CA'B (L = 0-1,2,....)
G( = C(I Ae Ae
Each mode i contributes 1 impulse slice H Ci..
Energy (H2) of layer and a mode.
Parseval implies
H. 2n /
B
E
[G
Using [uo"11 = ||
2 and the geometric series,
||C,,.12 B..
E
|
E.
1 1
energy
0=0
i=1
Intuition. E; is the mode's long-run output power (e.g., under unit white-noise inputs). If B rows are normalized, then E; C. 1l3 12,12).
B.3 AIRE-PRUNE: THE ALGORITHM
1. Per-mode energy. Compute E, for each state in each layer.
2. Sort & prefix sums (per layer). Sort E(p 2 ... 2 E(n) set S(i) = Ejsi EGj
3. normmaiedr score.
E
0 0<ELL.
5(i) +
This score is monotone non-increasing in i; single global threshold causes each layer to keep its longest prefix and prume contiguous tail.
4. Global threshold. Concatenate all s(i) across layers, choose for the desired global budget, and materialize kept/pruned indices per layer.
Practical effect. The prefix rule yields clear head-tail separation and, at higher thresholds, may delete entire low-energy layers, which helps wall-clock latency.
14

---

Published as a conference paper at ICLR 2026
B.4 WHY DISCUSS WORST-CASE BOUNDS IF ENERGY ALREADY WORKS?
The energy lens (H2) explains typical/average behavior and aligns with many workloads. We there- fore derive oo certificate from the same energy quantities AIRE already computes, plus single stability factor. The method remains unchanged; the mathematics provides assurance why AIRE- Prune works.
WHY THE BOUND SHOULD HOLD
Sort layer's states by energy. Cumulative energy vs. kept states shows an elbow: most energy sits short head. Pruning the tiny tail barely changes the typical output. Our result shows the worst- case change is also small unless pruned poles are extremely close to the unit circle. The stability margin enters through a single factor.
B.5 PER-LAYER WORST-CASE (Ho)) CERTIFICATE FROM ENERGY TAILS
Goal. We bound the worst-case (induced & - E2) error of layer after pruning set of modes T:
|EGi(eny ||:
Ci 1 Aie-Jw
G G||
sup
Gi(e2))
wE(0,2n)
iET
Step 1: Triangle inequality under the Ho norm. The induced (operator) 2-norm satisfies || +
BI S ||A||2 + BBl2. Taking the supremum over w preserves the inequality:
(triangle)
sup Σ Gi(e"))
(G: 112
sup
iET
(sup subaddinivity)
I
sup
Σ
Gi
1ET
1ET
Hence,
S
Σ G 2o.
(12)
IET
C.,
Step 2: Per-mode Ho envelope.
Write G:(e)w)
= 1 Açe-jw For rank-1 matrix we,
| ||2 ||/||||0||2 define
a; Ic,,lla Bi,,ll = -IC.,,
For any w, the reverse triangle inequality gives
Nfw- 2 - |Arl. Therefore,
|G, = 1 1 1-- ju
Oi
| < 1- PA
(13)
A
Step 3: Express the peak envelope via the energy E,, By definition of the per-mode energy (asymptotic impulse/white-non noise energy),
|
1 B
a
E
12
X
Solving for a; and substituting into equation
13
0; V AA
E,
V A 2
E.
1 +
>
Av
Ax
Thus, for every pruned mode
1 1 A
G.
I VEA
(14)
A
15

---

Published as a conference paper at ICLR 2026
Step 4: Uniform "stability factor" over the pruned set. Let
1+ p
max ET
A
and
k(p) -
V
1-p p
The scalar function x +
is increasing on [0, 1), so for all i E T,
Ad A4
S x(p).
Combining with equation 14 yields the uniform bound
||; I x(()) VE,
VIET.
(15)
Step First aggregate bound (linear in V.)). Substitute equation 15 into equation 12
Σ k(()) E <(p) Σ
Ex.
(16)
Step 6: Second aggregate bound via Cauchy-Schwarz (root-of-sum). For nonnegative {a:}
and m = (7), Cauchy-Schweere Schwarz gives Ei=1 a S mm E o.. Apply this with a, = VEi:
EVE, S VIT Σ
Ex.
RET
iET
Combining with equation
16 yields
ES k(p) VITI
Σ XET
E,
(17)
Step 7: Final certificate (take the minimum). Both equation 16 and equation 17 are valid; each be tighter depending whether the pruned tail is concentrated or diffuse. We therefore report the minimum:
|| G pox s(p) min EVE. T
Σ Fa SET
SET
Why the bound is finite (stability). Since the layer is stable, |Ayl 1 for all i, each E, of/(L 12,13) finite and p which implies k(p) /(1 p)/(1 (1 p) co.
Special cases. If B rows normalized, then E. 11 (the derivation above ap- plies verbatim). For bidirectional layers, sum (or average) the forward/backwwrd ||C||2 contributions when forming Eg; the rank-1 and norm inequalities used here are unchanged.
B.6 MATUEMATICAL COMPARISON TO LAST (AND THEORY POSITIONING)
af
Per-mode quantities (shorthand). a;
||C,,sla |B, 12.
E,
9; (20)
1
|
2
1+p
p = maxiet [del. <(p) V 1-p
1 A
Layer-level envelopes (apples-
LAST (Hoo-frst):
|G - - S Σ
a₁
(18)
1 Ax
ST
AIRE (energy-first + Hoo):
G م|| $ k(p) min
Σ
a:
VT
Σ
0²
A
ET - 1
ET
(19)
16

---

Published as a conference paper at ICLR 2026
Key mathematical differences.
1. Aggregation. LAST sums peak gains linearly, Σ a1 >xl). AIRE aggregates energies either linearly in v E, or sublinearly via V Σ E,, which is strictly tighter for diffuse tails (Cauchy-Schwarz gap).
2. Pole-radius dependence. As 1- 1A|2 2(1 | Both equation 18 and the first term in equation scale like (1 (1 uniform across the pruned set.
AIRE's k(p) A makes this dependence explicit and
3. Quantities needed. LAST needs a, and (1 (A,)) per pruned mode. AIRE needs E (already computed for ranking) and a single p. No new per-mode statistics are required beyond AIRE's pipeline.
4. Cross-layer coupling. LAST"s guarantee is tied to per-mode peak surrogates. AIRE's certificate tail-centric (energy tail one p), aligning with prefix-normalized, contiguous- tail pruning and explaining whole-layer drops.
When AIRE can be tighter.
Diffuse tails: if [T] large but Σ E, is tiny. then
x(p)
IT
Σ E
Ar
making AIRE's root-of-sum form substantially less conservative.
Conceptual positioning. LAST wosst-case-first and can be conservative when worst frequen- cies are rarely excited. AIRE is typical-case-fing (energy/H2). empirically enabling larger safe pruning.
C EXPERIMENTAL DETAILS
All S5 experiments were implemented in JAX (Bradbury et al. 2018) and run on a single NVIDIA A100 / H100 accelerator (40 GB or 80 GB VRAM). Our models and training code are based on public S5 (Smith et al. 2023 1 S4D (Gu et al. 2022a and Mamba(S6) Gu & Dao 2023 implementation. Unless otherwise stated, we employ bidirectional SSM layers for all LRA tasks, and use the parallel-scan inference kernel provided by S5.
C.I BENCHMARKS / TASKS
We evaluate the Long Range Arena (LRA), suite of long-context sequence problems span- ning symbolic reasoning, byte-level text, document retrieval, and flattened vision, designed to probe modeling of dependencies across sequences up to 16k tokens.
ListOps
10-way classification over extended ListOps expressions (Nangia & Bowman, 2018). In- puts are single-channel sequences (max length 2,048) encoding digits, operators, and bracket markers one-hot vectors over 17 tokens. Splits: 96k train, 2k validation, 2k test.
Text Binary sentiment classification on IMDB reviews at the byte level MMass al. 201 Each example is single-channel sequence up to 4,096 tokens, using 129-symbol one- hot alphabet. Splits: 25k train, 25k test.
Retrieval Binary document-paie pair classification on the ACL Anthology Network (Radev al. 2009). The goal is to predict whether two documents share equivalent citation links. Each doc- ument is byte-tokenized with a 97-symbol one-hot encoding and capped at 4, 4,000 tokens. Splits: 147,086 train pairs, 18,090 validation pairs, 17,437 test pairs.
Image 10-way classification on flattened CIFAR-10 (Krizhevsky & Hinton 2009 represented as single-channel sequences of length 1,024.
hateer//gthhub.com/hnnemmaaaa//
17

---

Published as a conference paper at ICLR 2026
Table Training configurations for S5 models on the six LRA tasks. All runs use batch normaliza- tion, mre-normalization and Amax 0.1. Rmm SSM state dimension; J: blocks for diagonal/block
init of A; D: dropout; LR: global learning rate; SSM LR: learning rate for SSM-only params; B:
batch size; E: epochs; WD: weight decay; min: minimum step.
Task
L
Dm
LR
SSM LR
WD
min
ListOps
8
128
16
8
0
0.003
0.001
50
40
0.07
0.001
Text
256
192
12
0.1
0.004
0.001
50
35
0.07
0.001
Retrieval
128
256
16
0
0.002
0.001
32
20
0.05
0.001
Image
6
512
384
3
0.1
0.005
0.001
50
250
0.07
0.001
Pathfinder
192
256
8
0.05
0.005
0.0009
64
200
0.07
0.001
Path-X
128
256
16
0.002
0.0006
75
0.05
0.001
Table 5: AIRE-Prunc pruning ratios and resulting inference speedups and parameter-memont re- ductions for Mamba (S6)
Metric / Dataset
ListOps
Text
Retrieval
Image
Pathfinder
Path-X
Pruning ratio
45%
60%
80%
Inference speedup (x)
1.61 x
1.57
2,03x
Parameter reduction
19.1%
8.1%
12.1%
Pathfinder Binary classification on flattened Pathfinder stimuli (Linsley et al. 2018 determining whether two points are connected by target path among distractors. Sequences are single- channel with length 1,024. Splits: 160k train, 20k validation, 20k test.
Path-X A scaled, more challenging Pathfinder variant (Linsley et al. 2018) with single-channel sequences of length 16,384, again testing path connectivity under heavy clutter.
Speech Speech command is 35-way classification task 1-second spoken utterances (Warden 2018). Each audio clip is a single-channel waveform length 16,000 samples (16 kHz). For experiments with different sampling rates, additionally create downsampled ver- sion at 8 kHz. The resulting split contains 24,482 training examples, 5,246 validation examples, and 5,247 test examples.
C.2 HYPERPARAMETER
We follow the LRA protocol across six tasks-ListOn Text, Retrieval, Image, Pathfinder, and Path-X-tuning depth (L), hidden width (h), and SSM state size (nm) per task. Byte-level Text and long-context Retrieval favor larger J (more blocks) and higher 7.m to capture long-range depen- dencies; flattened Image/Pathfinde variants use wider channels (h) with moderate dropout D for regularization. Learning rates decoupled, with a smaller SSM LRR than the global LR stabilize cigen-parameter updates, and modest weight decay (WD) throughout. Batch sizes and epochs E reflect dataset scale (e.g., longer training for Image/Pathfindrr, while Amio fixes the minimum step size used in state discretization. This setup (as in Table provides consistent, comparable training recipe across heterogeneous long-context workloads.
C.3 EFFICIENCY IMPACT OF DIMENSION DIMENSION-REDUCED S5 MODELS
To assess the practical efficiency impact of AIRE-Prune for mamba (S6), we implemented pruning by removal in addition to the standard masking-based implementation, transferring the retained high-importance states into dower-dimensional S5 model. step throughput and the corresponding parameter reduction achieved by the pruned S5 models on an NVIDIA H100 GPU. Shrinking the state dimension consistently improves both compute cost and model size, with the magnitude of savings varying across tasks as a function of channel width.
the average evaluation-
18


Reviews:
[{"review_id": "49wZbqQ2Wg", "reviewer": "Reviewer_4JBM", "summary": "The paper proposes a post-training pruning method for state space models (SSMs) that reduces state dimensions by minimizing long-run output energy distortion. Unlike previous worst-case gain approaches, their method assigns each state a closed-form asymptotic impulse-response energy score. These scores are normalized layer-wise for global comparison. The method achieves good results on Long Range Arena benchmarks, pruning 60.8% of states on average with only 0.29% accuracy drop without retraining.", "questions": "Are entire layers ever actually pruned? The paper states: \"our method can help achieve low latency SSM model as we are able to prune layers\"", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Good empirical results: 60.8% average pruning with only 0.29% accuracy drop without retraining on LRA\n- Closed-form solution for importance scores seems efficient\n- Energy-based metric is well-motivated from control theory perspective and has clear mathematical grounding\n- Works for both SISO and MIMO SSMs", "weaknesses": "- Minor: missing reference details in line 117\n- Limited to diagonal/diagonalizable SSMs and doesn't extend to input-selective models like Mamba\n- Only evaluated on Long Range Arena without speech and language benchmarks\n- No retraining experiments to show potential further improvements\n- Comparison mainly against only one recent baseline (LAST)\n- No discussion of computational overhead of computing energy scores\n- Limited analysis of why certain tasks (ListOps) are more sensitive to pruning", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mR0kJK5NHF", "reviewer": "Reviewer_p2pk", "summary": "This paper introduces a novel state pruning technique for SSMs. The core idea is to use the total impulse-response energy as a criterion to identify and remove less important states. The authors demonstrate that this method enables more aggressive pruning ratios compared to existing baselines, achieving new state-of-the-art results on the benchmark task.", "questions": "1.  **Motivation and Practical Impact:**\n    * Could the authors provide a clearer motivation for state pruning, especially in the context of models like Mamba where the state itself is a small fraction of the total parameters?\n    * When the paper claims to prune \"60% of states,\" what is the corresponding impact on the total parameter count of the model?\n    * Does this state pruning translate into tangible performance gains, such as reduced latency (in inference or training) or lower memory usage? We request that the authors provide empirical measurements for these metrics.\n\n* **Generalizability:**\n    * How well does the proposed impulse-response energy criterion generalize beyond the S5 model? Have the authors tested this approach on other significant SSM architectures, such as Mamba or S4?\n    * To what extent are these results dependent on the specific benchmark used? How does the method perform on other tasks or datasets?\n\n* Could the authors revise Sections 3.2 and 4 to make them more readable", "limitations": "", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1.  Novel Pruning Criterion: The paper proposes a new and theoretically motivated pruning technique for SSMs. \n2.  Strong Empirical Results: The method achieves state-of-the-art performance on the tested benchmark (S5 model), effectively demonstrating its capability to outperform prior pruning approaches in that specific setting.", "weaknesses": "1.  **Clarity of Methodology (Sections 3.2 & 4):** The paper's core methodology is difficult to understand as written. Section 3.2 lacks clarity and citations to fully grasp. Furthermore, the theoretical part (Section 4) is hard to understand\n\n2.  **Limited Evaluation and Generalizability:** The empirical validation is a significant weakness. The approach is only tested on a single model (S5) and a single benchmark. This narrow scope makes it impossible to assess the generalizability of the technique. It remains unclear whether these performance gains would translate to other important SSM architectures (e.g., Mamba) or to different tasks and datasets. The authors should expand their evaluation to include more models and benchmarks.\n\n3.  **Unclear Motivation and Practical Impact:** The paper's motivation for state pruning is not well-articulated, and its practical impact is obscure. The authors claim, for example, to prune \"60% of states,\" but it is unclear what this means in practice. In many modern SSMs (like Mamba), the state itself represents a small portion of the total parameters. The paper fails to connect state pruning to crucial downstream metrics. The authors must clarify:\n    * What is the impact of state pruning on the total parameter count?\n    * Does this pruning translate to tangible latency improvements (e.g., in inference or training) or memory reduction?\n\n    Without this context, the practical benefits of the proposed method are unclear.\n\nMinor:\n\n* Missing Citations: Several claims and components lack proper attribution.\n    * Section 3.1 makes assertions that require supporting citations.\n    * The S5 model is mentioned in Section 3.2 but is not cited.\n    * There appears to be a missing reference at L151.\n* Undefined Acronyms: The terms \"CT\" (Continuous-Time) and \"DT\" (Discrete-Time) are used before they are formally defined, which could confuse readers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dXP11bFK8Z", "reviewer": "Reviewer_HyDj", "summary": "This paper introduces AIRE-Prune, a structured post-training pruning method for state space models (SSMs) that uses asymptotic impulse-response energy to identify and remove redundant states. The method assigns each state a closed-form energy score representing its total output contribution over infinite time, then uses layer-wise prefix normalization to enable global cross-layer pruning decisions. Evaluated on S5 MIMO models across Long Range Arena benchmarks, AIRE-Prune achieves 60.8% average state pruning with only 0.29% accuracy degradation without retraining, substantially outperforming the prior LAST method (33% pruning at 0.52% loss). The approach extends classical modal truncation from single systems to deep SSM stacks and provides both typical-case (H₂ energy) and worst-case (H∞) theoretical justifications.", "questions": "- Mamba Extension (Critical): Can you extend AIRE to input-dependent SSMs? For Mamba, where Bₜ = sB(xₜ), Cₜ = sC(xₜ), how would you estimate Eᵢ? Monte Carlo over calibration data? Expected energy under input distribution?\n\n- Direct LAST Comparison: Can you provide results at matched pruning ratios (e.g., both methods at 50%, 60%) to help me get an equivalent comparison? What is the compute cost ratio for scoring?\nRetraining: What accuracy can you achieve with 1-5 epochs of fine-tuning after AIRE pruning? This would strengthen claims about practical deployment.\n-Layer Collapse: When entire layers are removed, does this create training instabilities if later fine-tuning? Have you tested this?\n\n-Failure Modes: ListOps only tolerates 20% pruning. Is this because: (a) the task truly needs full capacity, (b) S5 architecture is suboptimal for this task, or (c) AIRE scoring doesn't capture syntactic structure dependencies? Can you investigate?\n\n-Computational Complexity: What is the wall-clock time for computing AIRE scores vs. LAST scores for a typical S5 model? Is O(n) per state vs. O(n) amortized?\n\n\nI appreciate the authors for their wonderful work. I learned a lot while reading the manuscript.  I do understand that I asked a lot of questions. However, as long as the critical and decent number of other questions/concerns are answered, I am happy to increase my score.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Strong Theoretical Foundation: The paper elegantly bridges classical control theory (modal truncation) with modern deep learning (layer-adaptive pruning). The energy-based criterion has clear physical interpretation: states with low Eᵢ = ‖C:,i‖²₂‖Bi,:‖²₂/(1-|λᵢ|²) contribute minimal long-run output energy.\n\nPractical Algorithm with Closed-Form Solutions: Unlike iterative methods, AIRE-Prune requires only: (1) compute Eᵢ per mode, (2) sort and compute prefix sums, (3) apply global threshold. No matrix inversions or optimization loops needed.\n\nImpressive Empirical Results: Achieving 60.8% pruning at 0.29% accuracy loss substantially improves over LAST (33% at 0.52% loss). The step-function accuracy profiles (Figure 4) suggest the method effectively separates critical from redundant states.\n\nComprehensive Mathematical Analysis: Appendix B derives H∞ certificates showing ε ≤ κ(ρ)min{∑√Eᵢ, √|T|∑Eᵢ}, proving the energy-based method also bounds worst-case distortion. This dual justification is valuable.\n\nActionable Architectural Insights: Layer-wise profiles show task-dependent patterns—some layers can be entirely removed (enabling block-structured speedups), which is more deployment-friendly than fine-grained sparsity.", "weaknesses": "Severely Limited Experimental Scope:\n- Only S5 models evaluated: No experiments on Mamba, Mamba2, or hybrid architectures which dominate current practice\n- Only LRA benchmark: Missing speech (Speech Commands), language modeling (WikiText), or modern long-context tasks\n- Input-selective SSMs (Mamba) have input-dependent B, C—the energy formulation assumes these are fixed. How does AIRE extend to this case?\n\n\nIncomplete Comparison with LAST:\n\n- Table 1 shows different pruning ratios per task, making direct comparison difficult. Need controlled experiments at {30%, 40%, 50%, 60%, 70%} pruning with both methods. Computational cost comparison missing: Is AIRE faster than LAST to compute scores?\n\n\nMissing Key Experiments:\n-No retraining: Even 1-2 epochs could show full potential. Other SSM pruning work achieves near-zero loss with minimal fine-tuning\n- No ablations: What if using median normalization instead of prefix-sum? How sensitive to ε?\n- No failure analysis: Why does ListOps only tolerate 20% pruning? Is this fundamental or fixable?\n\n\nLimited Technical Analysis:\n\n- Equation 8 assumes |λᵢ|²ᵀ → 0, but convergence rate matters for numerical stability near the unit circle. Claim that entire layers can be removed needs gradient flow analysis—does this harm subsequent fine-tuning?\n-No discussion of how to extend to non-diagonal parameterizations (e.g., low-rank structured matrices)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rcVBEqCnMt", "reviewer": "Reviewer_EJJM", "summary": "This paper proposes AIRE-Prune, a post-training compression method to reduce the dimension of latent state space of a trained state-space model. The method is based on the idea of maximally preserving the energy of LTI systems while truncation. Results show that on Long-Range Arena, roughly $60\\\\%$ of states can be pruned without defecting the model's performance.", "questions": "1. I am confused about eq. (9) and the statement above that \"the total layer energy is additive across modes.\" According to the definition in  eq. (5), the energy is defined as the $L^2$ norm of the transfer function. For a diagonal LTI system, the transfer function is the sum of partial fractions that correspond to all states. However, there is no guarantee in general that these partial fractions are mutually orthogonal in $L^2$. If so, how can it be that the total layer energy is additive across modes?\n2. Can you show the energies of the states being pruned in your experiments also as a function of \"pruning ratio\"? This could potentially explain the \"elbow\" that you observe in Image, Pathfinder, and PathX.\n3. Have you tried the same method on S4; if so, how does the result compare?\n4. How does the model benefit from post-training compression in terms of time and space usage?\n5. While I understand intuitively why the $\\mathcal{H}_2$-based method works, as explained in the paper, I don't have intuition in why it performs better than an $\\mathcal{H}\\_\\infty$-based one. Can the author(s) provide more justification?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* The results of compressing S5 on LRA look compelling. The proposed method seems to perform significantly better than existing methods.\n* The energy score of every state can be easily computed without resorting to large matrix operations or simulations, making the approach efficient for large models.", "weaknesses": "* The discussion of energy, while intuitively, is not rigorous. More theory of why the proposed method tends to work better than the existing ($\\mathcal{H}\\_\\infty$-based) ones would strengthen the paper.\n* The experiment section could involve more diagnostic experiments. For example, it would be interesting to see which states are pruned by the proposed method and more diagnosis around the kink would be useful.\n* While I believe the paper contains nice ideas, its presentation has significant issues:\n   * \\\\citep and \\\\citet are completely misused.\n   * Similarly, en dashes (- in LaTeX) and em dashes (-- in LaTeX) are also misused all over the place.\n   * ICLR papers use boldface letters for matrices and vectors.\n   * I would avoid the subsection-style, which could disrupt the overall narratives of the paper.\n   * The paper is not proofread as it contains many typos. For instance, on line 151, the equation is not properly cross-referenced; on line 150, you wrote \"zero-order-hold\" but you use \"zero-order hold\" elsewhere, which is more standard; the caption of Figure 2 misses the right parenthesis, etc.\n\nThe paper, in its current form, is not acceptable. I would raise my score to 4 if the author(s) could fix the presentation issues outlined above. I will be happy to further increase the score if my questions below are properly addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)