Layer-Adaptive State Pruning (LAST)

• The paper introduces LAST, extending modal truncation to multi-layer, nonlinear SSMs by pruning redundant state dimensions based on energy contributions.
• It computes local H-scores normalized into LAST scores to enable a global, one-shot pruning threshold without necessitating retraining.
• Empirical results across benchmarks demonstrate that LAST achieves significant pruning with minimal accuracy loss, validating its practical efficiency.

Layer-Adaptive State Pruning (LAST) is a structured, control-theoretic framework for reducing the state dimensionality of deep state-space models (SSMs) while provably controlling worst-case output distortion. It extends modal truncation concepts from single-layer, linear time-invariant (LTI) systems to multi-layer, nonlinear deep SSMs by introducing globally comparable, layer-adaptive pruning criteria based on energy contributions as measured by $\mathcal{H}_\infty$ norms. LAST enables systematic, one-shot identification and removal of redundant state dimensions in existing SSM architectures with minimal impact on predictive accuracy and without requiring retraining (Gwak et al., 2024).

1. Optimization Objective and Energy-Based Pruning Criteria

LAST targets the reduction of state dimensions in deep SSMs, where an input sequence $u$ is transformed via $L$ successive linear state-space layers $\Sigma^{(1)}, \dots, \Sigma^{(L)}$, each coupled with a 1-Lipschitz activation $\sigma$, yielding output $y_\text{full} = f_\sigma(u; \Sigma^{(1:L)})$. The central objective is to prune a subset $P^{(l)}$ of states per layer so as to minimize the worst-case output energy distortion:

$$\min_{P^{(1)},\ldots,P^{(L)}} \left\| f_\sigma(u; \Sigma^{(1:L)}) - f_\sigma(u; \hat\Sigma^{(1:L)}) \right\|_2^2$$

where $\hat\Sigma^{(l)}$ denotes $\Sigma^{(l)}$ with the states in $u$0 removed, under a global pruning budget constraint.

The influence of each state is assessed using the discrete-time LTI transfer function matrix $u$1 for an SSM layer. The $u$2 norm, defined as

$u$3

(where $u$4 is the largest singular value), yields the output energy bound

$u$5

In the common diagonal parameterization, each state $u$6 forms a SISO subsystem with pole $u$7 and vectors $u$8; its transfer function is $u$9, and

$L$0

2. Pruning Score Derivation: From Modal Truncation to Layer-Adaptive Scoring

For a single diagonal SSM, the squared error induced by pruning states $L$1 is bounded by

$L$2

States are thus ranked by local “H-scores”:

$L$3

For multi-layer SSMs, pruning a state in layer $L$4 propagates through all downstream layers, and the error is similarly bounded by products of $L$5-norms across layers.

To allow cross-layer comparability, the “LAST score” for state $L$6 in layer $L$7 normalizes $L$8 by the cumulative H-energy in that layer:

$L$9

with states first sorted in descending $\Sigma^{(1)}, \dots, \Sigma^{(L)}$0. This energy normalization provides a global criterion for pruning and enables a single threshold selection across all layers.

3. LAST Algorithm: One-Shot Layer-Adaptive State Pruning

The LAST algorithm proceeds as a single-pass post-hoc procedure, requiring no retraining. The following pseudocode summarizes its workflow:

• Input: Trained SSM layers $\Sigma^{(1)}, \dots, \Sigma^{(L)}$1; global pruning ratio $\Sigma^{(1)}, \dots, \Sigma^{(L)}$2.
• For each layer $\Sigma^{(1)}, \dots, \Sigma^{(L)}$3:
  1. Compute local H-scores $\Sigma^{(1)}, \dots, \Sigma^{(L)}$4 as above.
  2. Sort states in descending $\Sigma^{(1)}, \dots, \Sigma^{(L)}$5.
  3. Compute cumulative energies $\Sigma^{(1)}, \dots, \Sigma^{(L)}$6.
  4. Compute LAST-scores $\Sigma^{(1)}, \dots, \Sigma^{(L)}$7.

• Collect all $\Sigma^{(1)}, \dots, \Sigma^{(L)}$8 from all layers, sort ascending.

• Select threshold $\Sigma^{(1)}, \dots, \Sigma^{(L)}$9 so that fraction $\sigma$0 of all states satisfy $\sigma$1.
• Prune every state with $\sigma$2.
• Output: Pruned SSM $\sigma$3.

This workflow leverages precomputed scores from the learned parameters and acts in a one-shot manner.

4. Theoretical Guarantees

Stability

The diagonal SSM parameterization assumes continuous-time poles $\sigma$4 strictly in the left half-plane (Hurwitz condition). Their zero-order-hold discretization ensures $\sigma$5, guaranteeing stability of each layer post-pruning.

Output Distortion Bound

For pruned sets $\sigma$6 in each layer, the cumulative output error satisfies

$\sigma$7

Hence, pruning states with small LAST scores provably results in a controlled upper bound on worst-case output distortion.

5. Empirical Results Across Benchmark Tasks

LAST was evaluated on standard sequence processing and classification benchmarks—Long Range Arena (LRA) tasks (ListOps, Text, Retrieval, Image, Pathfinder, Path-X), Speech Commands (raw 16 kHz/8 kHz), and pixel-level classification (sMNIST, psMNIST, sCIFAR). Two SSM variants were used: S4D (parallel SISO diagonal systems) and S5 (multi-input multi-output, MIMO, diagonal system).

Key empirical findings (one-shot, no retraining):

Model	Avg. % Pruned (All)	Avg. % Pruned (Compressible)	Avg. Accuracy Drop (All)	Avg. Accuracy Drop (Compressible)
S4D	25	33	0.32	0.42
S5	33	36.7	0.52	0.58

On selected LRA tasks, accuracy loss remained under 1% at pruning ratios up to 80% (S4D/Text), 60% (S5/Text), 60% (S4D/Retrieval), 50% (S5/Retrieval), 10%–30% (Pathfinder/Path-X, model dependent). For pixel tasks (sMNIST baseline 99.55%), S5 retained 99.01% accuracy with 50% pruning. For Speech Commands, at 20% pruning, accuracy dropped negligibly (96.43%→96.31%).

6. Significance, Limitations, and Implications

LAST establishes a globally consistent, theoretically justified framework for one-shot, layer-adaptive pruning in deep SSMs. Its control-theoretic approach quantifies state importance via energy contributions, enabling aggressive model compression with minimal accuracy loss and no need for retraining. The revealed redundancy and compressibility in SSMs suggest further research into structural optimization and the efficient design of deep sequential models. While the approach relies on diagonal parameterizations and the Hurwitz condition, the underlying methodology may inspire analogous schemes for more general architectures.