Jacobian-Frozen Stable Rollouts

Integrate this at the rollout layer of a recurrent world model, neural state-space model, or SSM that predicts latent states. At the beginning of every training sequence, compute one Jacobian-vector-product-based estimate of $A_k$ and $B_k$ at the observed initial state; construct $c_k$; then use the affine recurrence for an auxiliary horizon loss or for inference-time planning. Pseudocode: (1) evaluate $y=g_\theta(x_k,u_k)$; (2) obtain $A_k$ and $B_k$ with automatic differentiation, Hutchinson/vector-Jacobian products, or an explicit Jacobian for small latent dimensions; (3) set $c_k=y-A_kx_k-B_ku_k$; (4) repeatedly update $\hat{x}\leftarrow A_k\hat{x}+B_ku+c_k$; (5) minimize prediction loss plus $\lambda[\max(0,\rho(A_k)-r)]^2$. Compute $\rho(A_k)$ exactly for small states, or estimate it with 10--20 power iterations; use $r<1$, such as $0.98$, as the target. The paper supplies the frozen-affine construction; the contraction penalty and neural rollout use are adaptations. The first cheap experiment is a GRU or MLP state-space model trained on damped nonlinear oscillator or Lorenz-63 sequences, comparing ordinary nonlinear rollout, frozen-affine rollout, and spectral-penalty training. Sweep the learned Jacobian gain and measure the divergence boundary. The quantitative prediction is that affine error grows quadratically with radius around the expansion state, while long-horizon error sharply worsens when measured $\rho(A_k)$ crosses $1$; the empirical boundary should be within roughly 20% of this predicted threshold.