Covariance-Conditioned Neural Rollouts
Implementation & benchmark of arXiv:2607.15949 — Gaussian behaviors and stochastic data-driven control
Source paper: Gaussian behaviors and stochastic data-driven control arXiv:2607.15949 ⓘ · analyzed Aug 30, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Augment a neural latent or sequence model with a Gaussian behavior head that predicts an entire future trajectory jointly from the observed prefix and planned inputs. Instead of recursively applying only a point predictor, condition the learned joint trajectory covariance on the available prefix, producing a corrected future mean and uncertainty that incorporates temporal correlations.
Formulas
Mathematical statement
Let w ∈ Rᵈʷ contain the observed prefix, initial latent state, and planned inputs, and let y ∈ Rᵈʸ contain the future outputs or latent states. Assume the joint vector [w; y] is Gaussian with mean [μw; μy] and covariance blocks Σww, Σwy, Σyw, and Σyy. The conditional mean is μy + Σyw Σww⁻¹(w − μw), and the conditional covariance is the Schur complement Σyy − Σyw Σww⁻¹Σwy. In implementation, replace the inverse with a ridge-regularized inverse using λ > 0. The Schur-complement term is positive semidefinite and quantifies uncertainty remaining after conditioning on the prefix.
Implementation notes
Integrate this as a parallel uncertainty-aware head on a GRU, transformer, or neural state-space model. First train a deterministic model h[t+1] = fθ(h[t], u[t]) using ordinary teacher forcing. Collect overlapping trajectory windows containing an observed prefix, planned inputs, and a complete future sequence. Define w as the concatenation of the prefix and planned inputs, and define y as the future latent states or observations. Estimate μw, μy, and the covariance blocks from the training buffer, optionally after subtracting the deterministic network prediction so that the Gaussian head models residuals. At inference, compute M̂ = Σ̂yw solve(Σ̂ww + λI, I), then return mean = μy + M̂(w − μw) and covariance = Σ̂yy − Σ̂yw solve(Σ̂ww + λI, Σ̂wy). Use a low-rank-plus-diagonal covariance if the trajectory dimension is large. The Gaussian conditioning equations are taken from the paper; the feature map, ridge parameter, covariance rank, and residual estimation are empirical choices. The first cheap experiment is a small GRU on noisy Lorenz-63 or pendulum trajectories, comparing recursive mean-only prediction with the joint conditional head at horizons 10, 25, and 50. The predicted signature is that increasing prefix length decreases trace(Σ̂y|w), apart from finite-sample noise, because the removed term Σyw Σww⁻¹Σwy is positive semidefinite. Evaluate 95 percent ellipsoidal calibration using eᵀ(Σ̂y|w + εI)⁻¹e ≤ χ²[d_y, 0.95]; empirical coverage should approach 95 percent, and conditional negative log-likelihood should improve over an unconditional diagonal Gaussian head.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Confirmed ✓
Stage 2 · Practical benchmark result: Significantly worse ✗
Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works
Stage 1 — Mechanism check agent confidence 8/10
Built a covariance-conditioned joint rollout MVP in mvp.py, including a Schur-complement PSD check and held-out AR(1) rollout comparison. The conditional head achieved nearly identical mean MSE to recursive prediction (2.72084 vs 2.72124), but substantially improved Gaussian NLL (14.23 vs 18.92) and 95% coverage (94.6% vs 86.9%) using the full correlated covariance. The effect is real for uncertainty quality, not a demonstrated mean-accuracy gain; redundant older prefix history did not materially help after observing the latest state.
- Agent confidence
- 8/10
- Baseline
- Recursive AR rollout: MSE 2.72124, NLL 18.9187, 95% coverage 0.8687
- Idea
- Joint covariance-conditioned rollout: MSE 2.72084, NLL 14.2336, 95% coverage 0.9458, conditional covariance trace 26.8955, minimum eigenvalue 0.2972; Schur complement minimum eigenvalue 0.2894
Limitations: The mini-experiment uses a stationary Gaussian AR(1) process rather than a trained GRU, transformer, Lorenz-63 system, or neural latent residual head. It evaluates a fixed horizon of 10 and does not measure FLOPs, runtime, low-rank covariance scaling, or performance under distribution shift.
How to run: python3 mvp.py
Stage 2 — Benchmark vs. tuned baseline (arXiv:2607.15949)
Significantly WORSE than the tuned baseline (+116.3%, p=0.0386, wins 1 / 8 paired seeds).
- Benchmark
- Dynamics forecasting
- Model
- rnn_small
- Paired seeds
- 8
- Baseline mean
- 0
- Idea mean
- 0.0001
- Effect (Δ)
- 0 (+116.3%; negative = idea better)
- Wins
- 1 / 8 paired seeds
- p-value
- 0.0386 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±175.7%
- Mechanism
- Confirmed ✓
- Practical effect
- Hurts
- Baseline tuning
- swept over 3 configs
The benchmark used the built-in one-step actuated-pendulum dynamics task rather than a multi-step trajectory-output task, because that is the available matched dynamics track. The Gaussian head predicted scalar next-state uncertainty, not a full future trajectory covariance; uncertainty metrics were measured on seed 0 only and were not the primary comparison metric. No low-rank covariance scaling or distribution-shift evaluation was performed.
How to run: python3 bench_runner.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Stage 2 — Benchmark latest Failed ✗
Agent confidence: 9/10
Implemented a matched dynamics-track comparison using the shared rnn_small GRU backbone: deterministic MSE baseline versus an end-to-end Gaussian behavior head trained with Gaussian NLL. The Schur-complement sanity check was positive semidefinite, and the trained Gaussian model had 95% interval coverage of 1.00 on the measured seed, confirming an uncertainty-calibration signal within the declared tolerance. However, across 8 paired seeds the idea had higher test MSE (8.59209e-05 versus 3.97188e-05), with a significant paired loss (delta_mean=4.62020e-05, p=0.03865), so it did not win.
- Baseline
- dynamics/rnn_small; best lr=0.006; mean test MSE=3.971884348175081e-05; std=7.221693826148878e-06
- Idea
- GaussianGRU with the same GRU backbone; best lr=0.001; mean test MSE=8.592086487624329e-05; std=7.508073923840228e-05
- Mechanism signature
- seed-0 trained-model Gaussian NLL=-2.5754659175872803, 95% interval coverage=1.0, mean predicted variance=0.0008512148051522672, confirmed=true
- Paired comparison
- delta_mean=4.6202021394492476e-05, p_value=0.03865, idea_wins=1/8, verdict=idea worse (significant)
Benchmark result
Significantly WORSE than the tuned baseline (+116.3%, p=0.0386, wins 1 / 8 paired seeds).
- Benchmark
- Dynamics forecasting
- Model
- rnn_small
- Paired seeds
- 8
- Baseline mean
- 0
- Idea mean
- 0.0001
- Effect (Δ)
- 0 (+116.3%; negative = idea better)
- Wins
- 1 / 8 paired seeds
- p-value
- 0.0386 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±175.7%
- Mechanism
- Confirmed ✓
- Practical effect
- Hurts
- Baseline tuning
- swept over 3 configs
The benchmark used the built-in one-step actuated-pendulum dynamics task rather than a multi-step trajectory-output task, because that is the available matched dynamics track. The Gaussian head predicted scalar next-state uncertainty, not a full future trajectory covariance; uncertainty metrics were measured on seed 0 only and were not the primary comparison metric. No low-rank covariance scaling or distribution-shift evaluation was performed.
How to run: python3 bench_runner.py
Artifacts
- bench_report.json 3.7 KB View Raw
- bench_runner.py 4.7 KB View Raw
- experiment.py 4.8 KB View Raw
- mechanism_results.json 2.7 KB View Raw
- mvp.py 4.6 KB View Raw
- mvp_results.json 1.6 KB View Raw
- report.md 1.5 KB View
- report_bench_2026-09-02T160812.md 4.8 KB View
- results.json 3.4 KB View Raw
- verify_mechanism.py 3.5 KB View Raw