# Imaginary-Axis Gramian Compression for Neural SSMs

- ID: 1261
- Canonical URL: https://synthcore.org/idea/1261/imaginary-axis-gramian-compression-for-neural-ssms
- API JSON: https://synthcore.org/api/idea/1261.json
- API Markdown: https://synthcore.org/api/idea/1261.md
- Verification status: unverified
- Source: [arXiv:2607.21969](https://arxiv.org/abs/2607.21969)
- Category: architecture
- Solves: memory, speedup, stability
- ML areas: ssm, rnn, inference-speedup, embedding
- Math tags: linear-algebra, control-theory, dynamical-systems, spectral-theory, numerical-analysis
- Ratings: usefulness 6/10; difficulty 6/10; novelty 7/10

## Idea description

Replace a large stable linear state-space or recurrent layer by a lower-order balanced realization computed from frequency-targeted controllability and observability Gramians. Use generalized low-rank ADI with imaginary-axis shifts concentrated at frequencies that dominate the training data, then retain states associated with the largest approximate Hankel singular values. This should reduce recurrent inference cost while preserving the layer's input-output response in the selected frequency band.

## Mathematical statement

For a continuous-time descriptor layer $E\dot{x}=Ax+Bu$, $y=Cx+Du$, assume $E$ is invertible and $E^{-1}A$ is Hurwitz. The controllability and observability Gramians $P,Q\succeq0$ satisfy $APE^{\top}+EPA^{\top}+BB^{\top}=0$ and $A^{\top}QE+E^{\top}QA+C^{\top}C=0$. The paper computes low-rank factors $P\approx ZZ^{\top}$ and $Q\approx YY^{\top}$ using ADI shifted linear solves, with shifts $\sigma_j\in\mathbb{C}$ allowed on the imaginary axis. The singular values $\gamma_i$ of the small matrix $M=Y^{\top}EZ$ approximate Hankel singular values. Retaining the first $r$ balanced directions produces a reduced layer; for stable full-band balanced truncation, the transfer-function error is bounded by approximately $2\sum_{i>r}\gamma_i$. Imaginary-axis shifts near task-relevant frequencies implement a frequency-targeted version of this compression.

## Key formulas

- $$APE^{\top}+EPA^{\top}+BB^{\top}=0,\qquad A^{\top}QE+E^{\top}QA+C^{\top}C=0,$$
- $$P\approx ZZ^{\top},\qquad Q\approx YY^{\top},\qquad M=Y^{\top}EZ=U\,\operatorname{diag}(\gamma_1,\ldots,\gamma_k)V^{\top},$$
- $$G(s)=C(sE-A)^{-1}B+D,\qquad \|G-G_r\|_{\infty}\lesssim2\sum_{i=r+1}^{k}\gamma_i,$$
- $$\sigma_j=\mathrm{i}\omega_j,\qquad \omega_j\in\Omega_{\mathrm{task}},\qquad (A-\sigma_jE)V_j=R_j,$$

## Implementation notes

1. Integration point: apply the method after training, or periodically during training, to a linear SSM or recurrent block with state dimension $n$. Keep nonlinear gates and readout layers unchanged, and compress only the linear state transition. For a discrete layer $x_{t+1}=A_dx_t+B_du_t$, $y_t=C_dx_t+D_du_t$, either use the discrete Lyapunov analogue or convert to a stable continuous-time realization. 2. Pseudocode: collect a few minibatches of training inputs, estimate their FFT power spectrum, select frequencies $\omega_j$ at spectral peaks or use a logarithmic grid over the task band, and set conjugate-complex shifts $\sigma_j=\mathrm{i}\omega_j$. Run generalized LRCF-ADI on the primal equation, solving $(A-\sigma_jE)V_j=R_j$ and applying the paper's low-rank residual/factor update; repeat on the dual system $(A^{\top},C^{\top})$ to obtain $Y$. Form $M=Y^{\top}EZ$, compute its SVD, retain rank $r$, construct balanced projection matrices, replace the original block, and fine-tune surrounding parameters. 3. Computed from the paper: shifted solves, low-rank factors, Lyapunov residuals, and balanced projections. Estimated empirically: the input spectrum, rank $r$, and the frequency band. Monitor $\|APE^{\top}+EPA^{\top}+BB^{\top}\|_F/\|BB^{\top}\|_F$ and the dual residual. 4. First cheap experiment: use a 256-state S4-like SSM or linear RNN on sequential MNIST or Speech Commands. Compare the original model, standard full-band balanced truncation, magnitude pruning, and imaginary-axis frequency-targeted compression at ranks $16,32,64,128$. Measure validation accuracy, inference latency, parameter count, and $\max_{\omega\in\Omega}\|G(\mathrm{i}\omega)-G_r(\mathrm{i}\omega)\|_2$. The prediction is that targeted shifts give lower in-band transfer error at equal rank, with a visible accuracy/latency knee when $\sum_{i>r}\gamma_i$ becomes small. The measured transfer error should be within approximately $20\%$ of the proxy bound $2\sum_{i>r}\gamma_i$, and the normalized ADI residual should decrease over successful shift cycles.

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.