Hankel Singular Value Regularization for Highly Compressible State Space Models

AI summary

Overview

Research area: Machine learning for long-sequence modeling, specifically compression of deep neural networks that use state space models (SSMs) as layers, drawing on classical system and control theory.

Technical level: Advanced. The paper assumes familiarity with linear time-invariant systems, Lyapunov equations, gramians, Hankel operators, and balanced truncation, although the core idea can be stated plainly.

Scope: The paper proposes a training-time regularizer on the Hankel singular values of SSM layers so that the resulting networks can be compressed after training with standard model-reduction tools, and it develops an algorithm that makes this regularization computationally tractable.

What This Paper Is About

Deep networks built from state space model (SSM) layers perform well on long-range sequence tasks, but their layers are hard to shrink after training, so inference remains expensive. The cost of an SSM layer is governed by its order n, the dimension of the internal state, and ordinary training leaves the Hankel singular values of these layers decaying slowly, which means no much smaller system can approximate them well. The goal of this paper is to train SSM layers so that their Hankel singular values decay fast, making the layers highly compressible while keeping accuracy high.

Key Contributions

1. A system-theoretic account of SSM compressibility. The authors connect SSM layers to the Hankel operator and its singular values, deriving the conditions under which a layer of order n can be replaced by a system of much smaller order r with bounded error.

2. A differentiable Hankel singular value regularizer. They propose penalizing the sum of the Hankel singular values (the Hankel nuclear norm) of all SSM layers during training, and they prove (Proposition 2) that although the individual Hankel singular values are not differentiable with respect to the system matrices, their sum depends smoothly on the matrices A, B, and C.

3. An algorithm that lowers the cost of computing gramians from 𝒪(n³) to 𝒪(n²). By parametrizing the system matrix as a block-diagonal collection of scaled 2-by-2 rotation matrices, the Lyapunov equations decompose into q = n/2 small Lyapunov equations plus q(q−1)/2 Sylvester equations on 2-by-2 blocks that can be solved independently and in parallel. A matching associative parallel scan is also introduced for this parametrization, avoiding explicit products of the blocks.

4. Demonstrated compressibility on Long Range Arena benchmarks. The contributions list claims up to a 10× improvement in accuracy for strongly compressed models on Long Range Arena benchmarks, relative to a baseline referred to as "Pruning LAST."

Main Findings

• Standard training produces poorly compressible layers. For the LRA image benchmark, the Hankel singular values of learned SSMs decay slowly, so no much smaller system with r ≪ n can approximate the sequence-to-sequence map well. This agrees with earlier observations about SSM compression on standard LRA benchmarks.

• Regularization produces fast Hankel singular value decay. Adding the Hankel nuclear norm of all layers to the training loss nudges the optimizer toward SSMs with many low-energy states that contribute little to the layer output, so truncating only those states changes the sequence-to-sequence map only marginally.

• Compressibility gains up to 10×. The abstract states that regularized state space layers are up to 10× more compressible than standard state space layers while maintaining high accuracy; the contributions section phrases this as up to a 10× improvement in accuracy for strongly compressed models on Long Range Arena benchmarks.

• Cost reduction in the regularizer. Solving the Lyapunov equations with the standard Bartels-Stewart algorithm costs 𝒪(n³); vectorizing the Lyapunov equations directly costs 𝒪(n⁶); the block-structured approach in this paper costs 𝒪(n²). The computational cost of the regularizer is independent of the sequence length n_s.

• Scan cost. Generating an output sequence of length n_s with the proposed parallel scan scales as 𝒪(log(n_s)·n), assuming n_s processors, the same scaling as with diagonal system matrices in S5.

• Universality of the rotation parametrization. Proposition 1 states that for any linear time-invariant system of order n, an infinitesimal perturbation exists such that the perturbed model's sequence-to-sequence map can be described by an SSM with the proposed rotation-block matrices, with α and ρ in ℝ^(ñ) where ñ ≤ n.

• Experimental setup shown. Figure 2 plots test accuracy percentage (vertical axes) against compression level from 0 to 100 or 40 to 100 (horizontal axes) for sCIFAR10 (grayscale), sMNIST, IMDB, PATH, and PATH-X, comparing HSVR (ours), no regularization, and Pruning LAST. The numerical accuracy values underlying these curves are not reported in the available text.

• Parameter count grows linearly in the order. The number of SSM parameters per layer, the dimension of θ^(ℓ), is 2q + n(p−1) + pn + p, which grows linearly in n, matching diagonal system matrices as used in S5 layers.

Methodology in Plain English

The authors treat each SSM layer as a linear time-invariant dynamical system with state x, input u, and output y, driven by four matrices: A (how the state evolves), B (how input enters), C (how state is read out), and D (a direct feedthrough). The "size" of such a layer is the state dimension n, and the question is how small a replacement system can be without changing the layer's input-to-output behavior much.

The classical answer comes from the Hankel singular values: if those values decay quickly, a small system captures almost all of the behavior, and the error of the reduced system is bounded by twice the norm of the input times the sum of the discarded singular values. The Hankel singular values themselves are the square roots of the eigenvalues of the product PQ, where P and Q are the controllability and observability gramians obtained by solving two discrete Lyapunov equations.

The training procedure therefore adds the sum of all layers' Hankel singular values to the loss, which pushes the optimizer toward layers whose singular values drop off quickly. Making this practical requires two ingredients. First, the system matrix is written as a block-diagonal stack of scaled 2-by-2 rotation blocks, each described by a scale ρ and an angle α; stability is enforced by pushing ρ through a tanh so it stays below one, and α through a tanh so it lies in [0, π]. This yields a dense-in-the-space-of-systems parametrization (Proposition 1) with a differentiable, stability-guaranteed form. Second, because the system matrix is block diagonal, the Lyapunov equations split into many small, independently solvable pieces, reducing the cost from 𝒪(n³) to 𝒪(n²).

After training, compression is post-processing with no retraining: each layer's gramians are computed, balanced truncation is applied via the standard square-root method to get a reduced system of order r, and the reduced order is chosen per layer by an energy criterion such as retaining 99% of the layer's Hankel singular value energy, or by distributing a total state budget across layers with a bisection algorithm. Because balanced truncation gives balanced but not diagonal systems, the reduced systems are then diagonalized by an eigenvalue decomposition so the fast associative scan can still be used.

Why This Matters

The work connects ideas from classical system and control theory to deep learning practice, showing that compressibility can be designed into a network during training rather than repaired afterward. It also extends prior Hankel singular value regularization work — which used the same rank for every layer, diagonal system matrices, and physics problems — to layer-dependent rank adaptation with block-structured systems, demonstrated at the scale of Long Range Arena language and vision benchmarks.

Real-world applications suggested by the paper's framing:

• Long-sequence inference on constrained hardware. Lower-order SSM layers directly reduce the state dimension that drives inference cost, which matters for deploying sequence models where compute or memory is limited.
• Audio generation, listed as an application area for structured SSMs.
• Vision tasks, also listed as an application area, and reflected in the sCIFAR10 and sMNIST experiments.
• Long-range sequence modeling generally, including the language-modeling-style benchmarks emphasized by the authors' Long Range Arena experiments.

Industry relevance: Neural network compression via quantization, pruning, and distillation is described as increasingly important as models grow. This method adds a fourth, training-time route that produces models ready for off-the-shelf reduction tools without the retraining step that earlier direct model-reduction approaches to SSMs require.

Future Directions

• Extending beyond the demonstrated benchmarks. The authors position their Long Range Arena results as evidence of relevance to language modeling tasks beyond the physics problems of prior work; a broader evaluation across long-sequence domains is a natural next step.
• Tighter coupling between regularization strength and compression order. The regularizer strength and the per-layer retained-energy threshold (or total state budget) are treated as separate choices; the relationship between the two is left open.
• Further efficiency gains in the gramian computation. The block-diagonal rotation parametrization was chosen precisely because it admits 𝒪(n²) gramian solves; whether other parametrizations admit comparable structure is not addressed.
• Numerical details of the compression pipeline at scale. The truncated content does not report the full set of accuracy values behind Figure 2 or the exact compressed orders used, leaving detailed characterization of the achievable compression-accuracy trade-off for future reporting.

Target Audience

This paper is best suited to researchers and engineers working on efficient deep learning, sequence modeling architectures, and model compression, particularly those already familiar with state space models such as S4 and S5. It will also interest applied mathematicians and control theorists working on model reduction and balanced truncation who are curious about applications to neural networks. Readers without a background in linear systems and Lyapunov equations will find the derivations demanding, though the central idea — regularize during training so that post-training compression works — is accessible without them.

Authors’ abstract

Deep neural networks using state space models as layers are well suited for long-range sequence tasks but can be challenging to compress after training. We use that regularizing the sum of Hankel singular values of state space models leads to a fast decay of these singular values and thus to compressible models. To make the proposed Hankel singular value regularization scalable, we develop an algorithm to efficiently compute the Hankel singular values during training iterations by exploiting the specific block-diagonal structure of the system matrices that we use in our state space model parametrization. Experiments on Long Range Arena benchmarks demonstrate that the regularized state space layers are up to 10$\times$ more compressible than standard state space layers while maintaining high accuracy.