Neural Koopman Power-Iteration Latent Space

(1) Integration point: use an encoder with a mode bank $z_r=f_{r,\theta_r}(x)$ before a small latent dynamics module. For an image world model, $x$ is the encoder feature or observation embedding; for a state-space model, it is the current hidden state. Train the mode bank on batches of paired transitions $(x_i,y_i)$, where $y_i$ is the next state or next embedding. Keep the downstream predictor simple, initially $\widehat z_{t+1,r}=\lambda_r z_{t,r}$, and optionally add a learned small matrix after the diagonal baseline. (2) Pseudocode: initialize an empty list of mode outputs. For mode $r=1,\ldots,R$, compute $u_i=f_{r-1}(y_i)$, subtract its empirical mean, and set $u_i\leftarrow u_i-\sum_{j<r}(N^{-1}\sum_i u_i f_j(x_i))f_j(x_i)$. Normalize $u$ by $s=(N^{-1}\sum_i u_i^2+\varepsilon)^{1/2}$. Update neural parameters $\theta_r$ for several minibatch epochs using $L_r=N^{-1}\sum_i[f_{r,\theta_r}(x_i)-u_i/s]^2$, while adding $\rho(N^{-1}\sum_i f_{r,\theta_r}(x_i))^2$ for centering and $\eta\sum_{j<r}(N^{-1}\sum_i f_{r,\theta_r}(x_i)f_j(x_i))^2$ for orthogonality. Recompute $\lambda_r$ with the Rayleigh quotient above. (3) The formulas determine normalization, deflation, and eigenvalue estimation; empirically estimate all inner products from minibatches and use an exponential moving average to reduce noise. Stop a mode if its residual regression error is high or its post-deflation norm is below a threshold. (4) First experiment: compare a GRU, an unconstrained MLP latent transition, and this $R=16$-mode model on a small Lorenz-63 or MuJoCo trajectory dataset. Train all models for equal parameter count and one-step loss, then evaluate 10-, 50-, and 100-step rollout MSE. Success is lower long-horizon error, slower latent norm growth, and a stable estimated spectrum across random seeds; a useful initial target is at least 20% lower 50-step MSE than the unconstrained latent baseline at equal one-step error.