# Conjugate Bayesian Latent Dynamics Head

- ID: 1439
- Canonical URL: https://synthcore.org/idea/1439/conjugate-bayesian-latent-dynamics-head
- API JSON: https://synthcore.org/api/idea/1439.json
- API Markdown: https://synthcore.org/api/idea/1439.md
- Verification status: beats_baseline
- Source: [arXiv:2607.26345](https://arxiv.org/abs/2607.26345)
- Category: dynamics
- Solves: stability, sample-efficiency, generalization
- ML areas: world-model, meta-learning, uncertainty, rl
- Math tags: dynamical-systems, probability, statistics, linear-algebra
- Ratings: usefulness 7/10; difficulty 5/10; novelty 6/10

## 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.

## 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.

## Key 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)$$

## 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

- Status: beats_baseline
- Mechanism evidence: yes
- Mechanism confirmed: yes
- Practical verdict: helps
- Verdict: 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.

### Mechanism check

- Verdict: 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.
- Confidence: 9/10
- 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.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 0.0081
- Paired wins: 8/8
- Benchmark verdict: idea better (significant)

## Artifacts

- [bench_report.json](https://synthcore.org/code/562/bench_report.json)
- [experiment.py](https://synthcore.org/code/562/experiment.py)
- [report.md](https://synthcore.org/code/562/report.md)
- [report_bench_2026-08-31T181208.md](https://synthcore.org/code/562/report_bench_2026-08-31T181208.md)
- [results.json](https://synthcore.org/code/562/results.json)
- [stage2_bench.py](https://synthcore.org/code/562/stage2_bench.py)
- [Download all files as ZIP](https://synthcore.org/download/562)

## Disclaimer

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