✓✓ Beats tuned baseline 2026

Conjugate Bayesian Latent Dynamics Head

Implementation & benchmark of arXiv:2607.26345 — MetaKoopman: Bayesian Meta-Learning of Koopman Operators for Modeling Structured Dynamics under Distribution Shifts

Usefulness7/10

Difficulty5/10

Novelty6/10

Dynamics Stability Sample efficiency Generalization World model Meta learning Uncertainty Rl Dynamical systems Probability Statistics Linear algebra

Source paper: MetaKoopman: Bayesian Meta-Learning of Koopman Operators for Modeling Structured Dynamics under Distribution Shifts arXiv:2607.26345 ⓘ · analyzed Aug 31, 2026

AI-generated research hypothesis, automatically tested. Not peer-reviewed.

Idea description

Replace the final nonlinear transition network of a latent world model with a linear Koopman-style transition whose coefficients have a Matrix Normal-Inverse Wishart prior. Meta-learn the prior across tasks, then adapt only closed-form sufficient statistics from a few recent transitions; this should be more data-efficient and uncertainty-aware than gradient fine-tuning under distribution shift.

Formulas

$$x_{t+1}=f(x_t,u_t)$$

$$\bar{\mathcal K}\psi(x_t,u_t)=\psi(f(x_t,u_t),u_{t+1})=\psi(x_{t+1},u_{t+1})$$

$$K_N=K_0+QQ^\top,\qquad M_N=(M_0K_0+Z^+Q^\top)K_N^{-1}$$ $$\nu_N=\nu_0+N,\qquad S_N=S_0+Z^+Z^{+\top}+M_0K_0M_0^\top-M_NK_NM_N^\top$$

$$z_{t+1}\mid q_t,\mathcal D\sim t_{\nu_N-d+1}\left(M_Nq_t,\frac{1+q_t^\top K_N^{-1}q_t}{\nu_N-d+1}S_N\right)$$

Mathematical statement

The paper starts from the nonlinear controlled system $x_{t+1}=f(x_t,u_t)$ and uses the Koopman relation $\bar{\mathcal K}\psi(x_t,u_t)=\psi(f(x_t,u_t),u_{t+1})=\psi(x_{t+1},u_{t+1})$, where $x_t\in\mathbb R^n$ is the state, $u_t\in\mathbb R^m$ is the control, $f$ is unknown, and $\psi$ is an observable. Implement a finite learned observable $z_t=\phi_\theta(x_t)\in\mathbb R^d$ and controlled linear model $z_{t+1}=Fq_t+\epsilon_t$, with $q_t=[z_t;u_t;1]\in\mathbb R^r$, $F\in\mathbb R^{d\times r}$, and $\epsilon_t\sim\mathcal N(0,\Sigma)$. Use the conjugate prior $F\mid\Sigma\sim\mathrm{MN}(M_0,\Sigma,K_0^{-1})$ and $\Sigma\sim\mathrm{IW}(S_0,\nu_0)$. For a context of $N$ transitions, stack $Q=[q_1,\ldots,q_N]\in\mathbb R^{r\times N}$ and $Z^+=[z_2,\ldots,z_{N+1}]\in\mathbb R^{d\times N}$. The exact posterior parameters are $K_N=K_0+QQ^\top$, $M_N=(M_0K_0+Z^+Q^\top)K_N^{-1}$, $\nu_N=\nu_0+N$, and $S_N=S_0+Z^+Z^{+\top}+M_0K_0M_0^\top-M_NK_NM_N^\top$. A one-step predictive distribution is multivariate Student-$t$: $z_{t+1}\mid q_t,\mathcal D\sim t_{\nu_N-d+1}(M_Nq_t,\frac{1+q_t^\top K_N^{-1}q_t}{\nu_N-d+1}S_N)$. The term $q_t^\top K_N^{-1}q_t$ explicitly increases uncertainty for regressors far from the adaptation data.

Implementation notes

Integrate this as the transition head of a small latent world model, after an encoder and before a decoder or policy. Let an encoder produce $z_t=\phi_\theta(o_t)$, where $o_t$ is an observation, and construct $q_t=[z_t;u_t;1]$. For each meta-training task, initialize prior parameters $(M_0,K_0,S_0,\nu_0)$ from a shared learned prior, collect a context window, and run: Z = encoder(obs[1:N]); Znext = encoder(obs[2:N+1]); Q = concat(Z[:-1], actions[:-1], ones); K = K0 + Q @ Q.T; B = M0 @ K0 + Znext @ Q.T; M = solve_right(B, K); nu = nu0 + N; S = S0 + Znext @ Znext.T + M0 @ K0 @ M0.T - M @ K @ M.T; q = concat(z_t, action_t, 1); mean = M @ q; scale = (1 + q.T @ solve(K, q)) * S / (nu-d+1); sample or score StudentT(mean, scale, df=nu-d+1) . Use Cholesky solves instead of explicit inverses, symmetrize $S$ as $(S+S^\top)/2$, and add $10^{-5}I$ for numerical stability. Train the encoder, decoder, and prior parameters with multi-step negative log likelihood across tasks; freeze encoder weights during online adaptation so only matrix sufficient statistics change. The exact mathematical quantities are the conjugate posterior and Student-$t$ predictive distribution; estimate empirically whether latent linearity holds using held-out rollout error and estimate calibration using prediction-interval coverage. First test on CartPole, a synthetic switched nonlinear system, or HalfCheetah state prediction with task-specific friction. Compare against an MLP transition model and an MLP adapted by gradient descent using the same context size and compute budget. Success means lower 10-step prediction NLL, improved 90% interval coverage, and faster adaptation after 5-20 context transitions without online backpropagation.

Verification

Beats tuned baseline

Stage 1 · Toy mechanism gate: Passed ✓

Stage 2 · Mechanism transferred to benchmark: Confirmed ✓

Stage 2 · Practical benchmark result: Significant win ✓

Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works

Stage 1 — Mechanism check agent confidence 9/10

Built a reproducible Matrix Normal-Inverse Wishart conjugate dynamics head in experiment.py with stable posterior and Student-t predictive calculations. The normal-equation residual was 2.7e-15; predictive scale increased by the predicted 2.4402x between near and far regressors, observed exactly 2.4402x; degrees of freedom increased exactly by one per context transition; and posterior coefficient RMSE decreased from 0.0927 at N=1 to 0.0131 at N=80. On a shifted synthetic task, Bayesian predictive NLL was better than the ridge baseline at N=2, 5, and 10, but not clearly at N=20, so the mechanism is real while the practical win is mainly in the low-data regime.

Agent confidence: 9/10

Baseline: Ridge adaptation shifted-task NLL: N=2 75.105, N=5 4.499, N=10 0.157, N=20 -0.966
Idea: Conjugate Bayesian Student-t predictive NLL shifted task: N=2 -0.235, N=5 -0.526, N=10 -0.966, N=20 -1.104; leverage scale ratio 2.440 observed vs 2.440 predicted; df exact for N=0,2,5,10,20,40; adaptation RMSE 0.0927 -> 0.0131 from N=1 -> 80

Limitations: This is a synthetic linear controlled-dynamics regression test, not a learned encoder/world model. It does not test meta-learning across many tasks, multi-step rollout error, CartPole or other physical environments, CUDA behavior, or a carefully compute-matched gradient-adapted MLP. The ridge comparison uses a simple residual variance estimate and is not a full calibrated uncertainty baseline.

How to run: python3 experiment.py

Stage 2 — Benchmark vs. tuned baseline (arXiv:2607.26345) Sep 2, 2026

Beats the tuned baseline by 98.8% (p=0.0081, wins 8 / 8 paired seeds; metric: lower is better).

Benchmark: Dynamics forecasting
Model: shared_gru_latent_linear_head
Paired seeds: 8
Baseline mean: 0.0326
Idea mean: 0.0004
Effect (Δ): -0.0322 (−98.8%; negative = idea better)
Wins: 8 / 8 paired seeds
p-value: 0.0081 (permutation test, 20 000 shuffles)
Smallest detectable effect: n/a
Mechanism: Confirmed ✓
Practical effect: Helps
Baseline tuning: swept over 3 configs

Limitations:

This is a supervised fixed dynamics benchmark, not explicit online few-shot adaptation across separated meta-training tasks. Long-horizon rollout NLL, wall-clock adaptation speed, full multivariate Student-t likelihood, and a separately gradient-adapted MLP baseline were not tested.

How to run: python3 stage2_bench.py

Verdict computed by deterministic test code from paired-seed statistics — not by the language model.

Stage 2 — Benchmark Aug 31, 2026 latest Worked ✓

Agent confidence: 7/10

The conjugate Bayesian latent dynamics head significantly improved test MSE on the registered dynamics track. It beat the tuned baseline on all 8 paired seeds with delta_mean -0.0321840645 and permutation p=0.0081. The trained-model signature confirmed leverage-dependent uncertainty scaling, though the uncertainty effect was small and coverage was conservative.

Baseline: Best lr=0.01; full 8-seed mean MSE 0.032576876.
Comparison: delta_mean=-0.032184065; idea_wins=8/8; p_value=0.0081.
Idea: Best lr=0.01; full 8-seed mean MSE 0.000392811.

Benchmark result