Imaginary-Axis Gramian Compression for Neural SSMs

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.