Beyond Unit-Circle Eigenvalues: Invariant Bases for Stable State Space Dynamics

AI summary

Overview

Research area: Machine learning for scientific simulation — specifically stable long-horizon sequence modelling with diagonal state space models (SSMs) and structure-preserving operator design. Technical level: Advanced. The paper is built on linear algebra, operator theory (Cayley transforms, Kronecker products, graph Laplacians) and conservation laws, though the core argument is stated in plain terms. Scope: The paper argues that conserved quantities in a diagonal recurrence are layered — each layer requiring a strictly stronger algebraic condition on the transition operator — and proposes the Kronecker–Cayley (KC) transition, which assembles all the required conditions by construction rather than learning or post-correcting them.

What This Paper Is About

Diagonal state space models advance a latent state by applying the same diagonal transition at every step of a rollout, at linear cost in sequence length. Because nothing in an unconstrained recurrence is held fixed, the state drifts: it decays toward zero or grows without bound, and quantities like total mass or momentum are lost. The paper's goal is to determine exactly which quantities a diagonal recurrence can preserve, what condition on the operator each one requires, and why adding a conservation penalty to the training loss does not supply those conditions — then to build a transition that satisfies them all at once.

Key Contributions

1. A layered-conservation theorem. The authors prove that the conserved quantities of a diagonal recurrence are nested: unit-modulus eigenvalues (|d_k| = 1) preserve only the norm; vanishing row sums (L𝟏 = 0) additionally preserve the summed quantity; and the Kronecker form H = L ⊗ I_m additionally preserves an angular pairing. Each condition implies the one before it, and no converse holds.
2. The Kronecker–Cayley (KC) transition. The generator is a graph Laplacian over N object subspaces tensored with the identity, advanced by a Cayley transform. Because Cay(L ⊗ I_m) = Cay(L) ⊗ I_m, the step stays diagonal in the modal basis at cost O(N³ + N²m) rather than O(N³m³). The authors prove the conserved set is exactly a 2m²-dimensional family of quadratic forms, held at every parameter value.
3. An applicability criterion and a negative result for learned conservation. Corollary 5 characterises which physical quantities (quadratic in a linearly decoded output) are held exactly, and Proposition 6 gives a lower bound showing any unitary step incurs per-step error on contracting or expanding dynamics. Proposition 7 shows a trained modal basis can only conserve the summed quantity by freezing every mode, confirmed empirically.
4. Empirical validation across physics, chemistry and engineering. Four conservation benchmarks plus three N-body systems and operator-level tests in molecular dynamics, chemical kinetics, pharmacokinetics and power-grid swing dynamics.

Main Findings

• Spectrum constraints are not conservation constraints. A transition with every eigenvalue on the unit circle holds the norm to machine precision, yet momentum drifts by 141 percent and the angular pairing by 286, because the norm is a property of the spectrum whereas a sum is a property of the coordinates.
• KC holds mass where every diagonal baseline loses it. On four benchmarks (dam-break and advection from PDEBench, spherical shallow water from PDEArena, and a conservative Allen–Cahn flow) at seeds 42–44, KC's drift is 2.0 × 10⁻⁵ on Allen–Cahn, 2.6 × 10⁻³ on advection, 8.3 × 10⁻³ on the dam break and 483 on the sphere (a discretisation residual, since the reference solver's own value is 30.8). Every other diagonal transition loses between 3 and 2 × 10¹¹.
• Correction at the output does not replace the condition. CPFNO's learned corrector holds mass near 10⁻³ for fifty steps while the prediction before correction drifts by 10⁴ from the first step — seven orders of magnitude closed. But four of five seeds diverge between steps 48 and 83 for the learned corrector, and the fixed uniform corrector diverges on every seed between steps 58 and 97. KC is the only model reaching the final horizon with bounded error on every seed.
• Conservation and accuracy are separate. On advection, mass correction beats KC on drift (2.8 × 10⁻⁴ vs 2.6 × 10⁻³) on a field whose nRMSE of 1.85 is worse than the constant field's 0.702. On the one dam-break seed where FNO reaches the horizon, it is twice as accurate as KC while drifting its mass by 230.
• Each condition holds exactly one quantity. Removing one condition at a time confirms Corollary 3: a matched perturbation of the trained Laplacian with row sums kept costs nothing (3.5 × 10⁻⁵ momentum drift), while breaking row sums costs 110 percent of the summed quantity, with energy at 2.0 × 10⁻⁶ in both cases. At the operator level over 2000 steps, four norm-preserving operators are equally stable below 10⁻², yet Haar unitary, dense Cayley and diagonal unitary lose the summed quantity at 169.7, 170.7 and 172.5 percent respectively, against 4.0 × 10⁻³ for KC.
• Training reaches conservation only by freezing the dynamics. A free orthogonal basis trained under a conservation penalty (sixteen seeds at λ = 10, with a sweep over λ ∈ {0, 0.1, 1, 10}) conserves only on the seeds that collapse to d_k = 1 for every mode, meaning the latent no longer evolves; the basis stays near its initialisation rather than moving to the conserving configuration KC satisfies by construction.
• N-body momentum is held. KC holds linear momentum to 2 × 10⁻⁶ against 0.5 to 8 for every baseline, and at the operator level holds the conserved quantity of molecular dynamics, chemical kinetics, pharmacokinetics and power-grid swing dynamics to within 10⁻² percent.
• The construction has a stated boundary. On a dissipative system it fits less accurately than baselines that can contract, exactly as Proposition 6 predicts: the per-step error is bounded below by the difference in latent norms, and no choice of L, θ or decoder reduces it.
• The guarantee is cheap in parameters. The learned coupling network carries 3,963 parameters against roughly 15 M in the shared encoder and decoder, so conservation is not bought with transition capacity.

Methodology in Plain English

The authors start from a single algebraic fact: for a transition built as the Cayley transform of a Hermitian generator H, a quadratic quantity z†Gz is preserved for every state exactly when G commutes with H. That converts the question "what does this model conserve?" into "what is the algebraic shape of H?" — a property of the pattern, not of the numbers placed in it, so it holds at every parameter value and needs no correction after the fact.

They then design the pattern deliberately. At every rollout step, a small network reads the distance between each pair of objects and the current action and outputs one real scalar per pair. Those scalars are arranged into a graph Laplacian L — off-diagonal entries are the negative couplings, diagonal entries are set so each row sums to zero — and the generator is H = L ⊗ I_m. The Cayley transform of this generator is unitary for any Hermitian H, and the Kronecker structure passes through the transform, so the step rotates the N objects into one another while leaving every coordinate channel untouched.

Because the Laplacian depends on the state, its eigenbasis changes each step, costing one N × N eigendecomposition and two modal transforms — O(N³ + N²m). This is the price of input dependence that fixed-basis models like Mamba avoid by holding the basis constant.

To test the claims, the authors hold the encoder and decoder fixed across baselines so that "level A" (free spectrum), "level B" (unit modulus), "level C" (one mode per object fixed at unity) and the full KC assembly reproduce the layers of their corollary. They also run ablation controls that remove exactly one condition from a trained Laplacian, and operator-level tests with no encoder or decoder at all, where four norm-preserving operators are compared across 2000 steps.

Why This Matters

Conservation in learned simulators is usually treated as an accuracy problem — add a penalty, add a corrector, train harder. This paper reframes it as a structural question with a sharp answer: the identity of the conserved quantity follows from the algebra of the operator, and the basis in which the spectrum acts matters as much as the spectrum itself.

Real-world applications:

• Climate and weather emulation. The paper cites a documented case where a water-budget corrector hid a degrading climate model — precisely the failure mode observed in the CPFNO experiments here.
• Molecular dynamics, where velocity rescaling is known to drain kinetic energy into centre-of-mass motion until internal degrees of freedom freeze.
• Chemical kinetics and pharmacokinetics, where mass-conserving schemes leave the positive orthant and must be projected back onto the reaction simplex.
• Power-grid swing dynamics, where the conserved quantity is a structural invariant of the network rather than a modelling convenience.

Industry relevance: scientific machine learning platforms, simulation software vendors, and any pipeline where a surrogate model must run for thousands of autoregressive steps without a post-hoc fix. Because the guarantee is enforced at every parameter value, it survives training, retraining and hyperparameter changes without re-derivation — and it costs only a few thousand parameters in the coupling network.

Future Directions

• Non-quadratic quantities remain open. Corollary 5 characterises only quantities that are quadratic under a linear decoder; a potential energy of the form Σᵢ<ⱼ ‖rᵢ − rⱼ‖⁻¹ is explicitly outside the criterion's reach, and energy is reported as held only where the potential is quadratic.
• Dissipative dynamics need a different construction. Proposition 6 establishes that no unitary step can avoid a per-step error floor on contracting systems, so extending the guarantee to dissipative regimes requires leaving the unitary family entirely.
• Scaling the eigendecomposition. The O(N³ + N²m) per-step cost from the state-dependent basis is a practical ceiling on the number of interacting objects; cheaper ways to update the modal basis without losing the guarantee are unresolved.
• Closing the loop on learned bases. The paper shows training reaches conservation only by freezing the spectrum; whether a parameterisation exists that lets a learned basis reach the conserving configuration and keep modes rotating is left as an open design question.

Target Audience

Researchers and practitioners in scientific machine learning, structure-preserving numerics and sequence modelling — particularly those working on neural surrogates for physical simulation who need long-horizon rollouts to respect conservation laws. It will also interest readers of the state space model literature (S4, LRU, Mamba-style architectures) who want to understand the limits of spectral stability constraints. A background in linear algebra is assumed; the empirical results are readable without following the proofs.

Authors’ abstract

Diagonal state space models process long sequences at linear cost, yet the same transition is applied at every step of a rollout, so an unconstrained recurrence holds no quantity fixed and the state decays or grows without bound. Current designs constrain one quantity, and preserving one does not guarantee preservation of the others, so a model can hold its norm to machine precision while momentum drifts by hundreds of percent. We claim that the conserved quantities of a diagonal recurrence are layered, that each layer requires a strictly stronger condition on the operator, and that a unit-modulus spectrum reaches only the outermost. Acting on the claim, we enforce the conditions required for every layer by construction. The generator of the Kronecker-Cayley transition (KC) is a graph Laplacian over object subspaces, $\mathbf{H}=\mathbf{L}\otimes I_m$, advanced by a Cayley transform, and the step stays diagonal at $O(N^3+N^2m)$ cost. The conserved set is exactly a $2m^2$-dimensional family of quadratic forms, held at every parameter value, so the quantities and the dynamics the construction reaches are fixed before training. On four conservation benchmarks, KC holds total mass at or near the reference floor while every diagonal baseline loses it by orders of magnitude. We further compare against mass correction and find that two correctors fail once the underlying prediction diverges, which shows that assembly rather than correction is necessary for a guarantee that holds across the rollout. A free modal basis trained under a conservation penalty conserves only by freezing its spectrum, whereas the assembled operator conserves with every mode but one free to rotate. Beyond these benchmarks, KC holds the linear momentum of three $N$-body systems and, at the operator level, the conserved quantity in molecular dynamics, chemical kinetics, pharmacokinetics and power-grid swing dynamics.