Support-Sparse Koopman World Model
Source paper: Sparse Koopman Autoencoders Identify Local Dynamical Regimes in Multibasin Systems arXiv:2608.29057 ⓘ · analyzed Sep 1, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace a dense Koopman autoencoder latent with a sparse code whose active-coordinate support can represent the local dynamical regime or basin. Train reconstruction, latent linear prediction, and multi-step rollout losses jointly; use the learned support as a label-free regime variable and optionally select a local transition matrix for forecasting.
Formulas
Mathematical statement
The paper models sampled dynamics by the discrete recurrence x_{k+1}=F_{Delta t}(x_k), where x_k is the observed state in R^{d_x} and F_{Delta t} is the unknown time-Delta t flow map. Introduce an encoder E_theta from R^{d_x} to R^{d_z}, a decoder D_phi from R^{d_z} to R^{d_x}, and a latent linear operator K in R^{d_z times d_z}, with sparse latent code z_k=E_theta(x_k). The adapted objective is L=L_rec+alpha L_koop+gamma L_roll+lambda L_sparse+eta L_support. Here L_rec=||D_phi(z_k)-x_k||_2^2, L_koop=||z_{k+1}-K z_k||_2^2, L_roll=(1/H) sum from h=1 to H of ||D_phi(K^h z_k)-x_{k+h}||_2^2, and L_sparse=||z_k||_1. Define the soft support s_j(z)=|z_j|/(|z_j|+epsilon), where epsilon is positive, and encourage support persistence along a trajectory with L_support=||s(z_{k+1})-s(z_k)||_1. At inference, the discrete support is m_j(z)=1 if |z_j| exceeds tau and 0 otherwise, where tau is a validation-selected threshold. The intended mechanism is that different basins use different active coordinate subsets while each subset admits approximately linear evolution under K.
Implementation notes
Integrate this into a small state-space or world-model network, not into a language model. Use an MLP encoder E_theta mapping state vectors to a latent width d_z such as 32 or 64, an MLP decoder D_phi, and a trainable matrix K. Training batches should contain contiguous windows (x_k,...,x_{k+H}) from trajectories; use H=5 initially and increase to 10 after reconstruction stabilizes. For each window, compute z_k=E_theta(x_k), decode it, compute the one-step latent residual E_theta(x_{k+1})-K z_k, recursively generate K^h z_k, decode every rollout, and add the L1 latent penalty. Compute s(z) with absolute value divided by absolute value plus epsilon, using epsilon=10^{-6}, and add the support-persistence penalty for adjacent states from the same trajectory. Pseudocode: z=E(x[:,0]); rec=mean((D(z)-x[:,0])^2); pred_lat=0; pred_obs=0; sparse=mean(abs(z)); support=0; for h in 1..H: z=K @ z; pred_lat += mean((z-E(x[:,h]))^2); pred_obs += mean((D(z)-x[:,h])^2); support += mean(abs(s(E(x[:,h]))-s(E(x[:,h-1])))); loss=rec+alpha*pred_lat/H+gamma*pred_obs/H+lambda*sparse+eta*support. Select tau on validation data to produce an average active fraction of 10 to 30 percent, then cluster binary masks to test basin separation. The first experiment should use a two-basin nonlinear system such as a double-well oscillator or a mixture of two stable systems, with 1,000 training trajectories and held-out trajectories from each basin. Compare against a dense Koopman autoencoder with identical encoder width, parameter count, optimizer, and rollout loss. Success requires lower 20-step state MSE, less long-rollout divergence, and support-mask classification accuracy above the dense baseline on held-out basin-interior states. Include ablations without support persistence and with only one-step loss; the method is supported only if sparse supports remain stable along trajectories and forecasting improves at equal latent width.
Verification
This idea has not been verified yet.
Verification happens in two stages: Stage 1 — a mechanism check on a toy system confirms the claimed mathematical phenomenon reproduces; Stage 2 — a benchmark implements the idea on a real (small) neural network task and compares it against a tuned baseline over 8 paired seeds with a permutation test.
Artifacts
Artifacts unavailable.