Signature-conditioned cylindrical law head
Implementation & benchmark of arXiv:2608.08040 — A cylindrical neural approximation theorem for conditional laws of McKean-Vlasov equations with common noise
Source paper: A cylindrical neural approximation theorem for conditional laws of McKean-Vlasov equations with common noise arXiv:2608.08040 ⓘ · analyzed Sep 1, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Add a conditional-law head that maps a compact representation of an initial distribution and a shared-noise trajectory to a Gaussian mixture, then computes downstream predictions as analytic expectations under that mixture. This can replace expensive particle rollouts or particle pooling in stochastic world models and conditional diffusion systems while retaining multimodality.
Formulas
Mathematical statement
The paper's central construction is \(\hat V(x,\mu_0,w)=\Psi\!\left(x,\langle\varphi,\hat N(E_L(\mu_0),S_M(r_t w))\rangle\right)\). Here \(x\in\mathbb R^d\) is a query state, \(\mu_0\in\mathcal P_2(\mathbb R^d)\) is the initial probability law, \(w\) is the common-noise path, \(r_t w\) is its restriction to \([0,t]\), \(E_L(\mu_0)\) is a finite vector of Fourier moments, \(S_M(r_t w)\) is the truncated signature of the time-augmented common-noise path through level \(M\), \(\hat N\) is a mixture-density network returning a Gaussian-mixture approximation of the conditional law \(\mu_t\), \(\varphi\) is a test function, \(\langle\varphi,\mu\rangle=\int\varphi(y)\,\mu(dy)\), and \(\Psi\) is a cylindrical neural network. The theorem's mechanism is that Fourier moments separate probability measures, signatures separate paths up to the relevant equivalence, Gaussian mixtures are dense in Wasserstein space, and neural networks approximate the resulting continuous map. For implementation, use Fourier features \(E_L(\mu_0)=[\mathbb E_{Y\sim\mu_0}e^{i\xi_j^\top Y}]_{j=1}^L\), and signature coordinates \(S_M(X)=(1,\int dX,\int_{s_1<s_2}dX_{s_1}\otimes dX_{s_2},\ldots)\), where \(X_s=(s,w_s)\) is the time-augmented path. The learned law is \(\hat\mu_t=\sum_{k=1}^K\pi_k\mathcal N(m_k,\Sigma_k)\), with \(\pi_k=\operatorname{softmax}(a)_k\) and \(\Sigma_k=L_kL_k^\top+\epsilon I\). For a test function \(\varphi\), approximate \(\langle\varphi,\hat\mu_t\rangle=\sum_k\pi_k\,\mathbb E_{Z\sim\mathcal N(m_k,\Sigma_k)}[\varphi(Z)]\) analytically when possible or by reparameterized Monte Carlo.
Implementation notes
(1) Integration point: implement this as a reusable conditional-distribution head after an encoder of an initial particle set, replay buffer, latent state distribution, or batch of stochastic trajectories. The head receives initial samples \(y_1,\ldots,y_B\), a shared-noise path sampled at times \(t_0<\cdots<t_T\), and optionally a query \(x\). It outputs mixture logits \(a_k\), means \(m_k\), and Cholesky factors \(L_k\) for the law at time \(t\); downstream value, reward, or transition predictions consume expectations under the mixture rather than a single point estimate. (2) Pseudocode: compute E = concat([mean(cos(xi[j] @ y_b)), mean(sin(xi[j] @ y_b)) for j in 1..L]); form the time-augmented path X[t]=(t,w[t]); compute the truncated signature S using iterated-integral or log-signature software up to level M; set h = MLP(concat(E,S,query_x)); produce a,m,L = heads(h); set pi=softmax(a), Sigma=LL^T+eps*I; for each requested test function evaluate sum_k pi[k]*mean_r phi(m[k]+L[k]*normal()). Train with conditional negative log likelihood of observed terminal samples, plus optional downstream task loss. (3) Compute from the mathematics: Fourier feature construction, signature truncation, normalized mixture weights, positive-definite covariance parameterization, and differentiable expectation. Estimate empirically: choose frequencies \(\xi_j\), mixture count \(K\), signature level \(M\), covariance floor \(\varepsilon\), and Monte Carlo count \(R\) using validation likelihood and calibration. For irregular or noisy paths, use piecewise-linear signatures; for Brownian data, use a consistent Stratonovich-style geometric lift or explicitly compare Ito and Stratonovich preprocessing. (4) First cheap experiment: use a 1D or 2D stochastic dynamical system with shared Brownian noise and nonlinear drift, such as \(dY_t=f(Y_t)dt+g(Y_t)dW_t^{\mathrm{common}}+\sigma dB_t\). Compare the proposed head against (a) empirical particle plug-in pooling and (b) a standard MLP receiving raw noise samples. Train on 1,000 noise scenarios with 32 particles each, evaluate held-out conditional NLL, Wasserstein-2 distance, mean/variance calibration, and downstream expected reward. The expected signal is lower conditional NLL and Wasserstein error at the same particle budget, or comparable accuracy with at least a 4-8x reduction in particles and memory.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Not reproduced ✗
Stage 2 · Practical benchmark result: Inconclusive
Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works
Stage 1 — Mechanism check agent confidence 7/10
Built a reusable Fourier-feature plus truncated-signature Gaussian-mixture law head and a seeded stochastic toy experiment. The mechanism checks passed: level-2 signature recovered the cross-integral with maximum error 2.22e-16, Fourier features separated nonzero translations while giving zero distance for identical laws, and Cholesky covariances remained positive definite with increasing floors. In the toy comparison, the idea improved held-out NLL over the raw-path Gaussian baseline (-1.83 vs 0.68), but had worse mean and variance calibration; therefore the mechanism is real, while broad predictive superiority is not established.
- Agent confidence
- 7/10
- Baseline
- Raw-path Gaussian: NLL=0.6822, mean RMSE=0.0603, variance absolute error=0.0584
- Idea
- Signature MDN, level 2: NLL=-1.8266, mean RMSE=0.1098, variance absolute error=0.2252; level 1 NLL=-1.8016 and level 3 NLL=-1.8293. With 8 initial particles: NLL=-0.3701, mean RMSE=0.1568, variance absolute error=0.1918.
Limitations: Only a small 1D synthetic system and one fixed seed were tested. No real world-model or diffusion task, downstream cylindrical expectation loss, Wasserstein metric, independent particle-pooling baseline, runtime/FLOP/memory measurement, GPU execution, or multi-seed statistical comparison was evaluated.
How to run: python3 experiment.py
Stage 2 — Benchmark vs. tuned baseline (arXiv:2608.08040)
Ahead of the baseline (−33.6%) but not statistically significant (p=0.1792, wins 5 / 8 paired seeds).
- Benchmark
- Dynamics forecasting
- Model
- rnn_small
- Paired seeds
- 8
- Baseline mean
- 0.0112
- Idea mean
- 0
- Effect (Δ)
- -0.0038 (−33.6%; negative = idea better)
- Wins
- 5 / 8 paired seeds
- p-value
- 0.1792 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±51.3%
- Mechanism
- Not confirmed ✗
- Practical effect
- Inconclusive
- Baseline tuning
- swept over 3 configs
The completed protocol used 400 training samples, 200 test samples, 6 epochs, and batch size 128 because the initial larger run exceeded the time budget. Only the built-in dynamics track was tested; particle-rollout replacement, Wasserstein distance, downstream analytic reward, and larger-scale world-model experiments were not tested.
How to run: python3 bench_experiment.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 benchmark comparing a standard 64-unit GRU point predictor with the same GRU augmented by Fourier-conditioned level-2 path signatures and a Gaussian-mixture law head. The idea achieved lower test MSE (0.00744951 vs 0.01121620; delta -0.00376669; 5/8 wins), but the paired permutation test was not significant (p=0.17915), so it does not satisfy the benchmark success criterion. Mixture normalization and covariance-floor sanity checks passed, but trained-model variance calibration was not confirmed.
- Baseline
- dynamics/rnn_small, tuned lr=0.01: MSE mean=0.01121620, std=0.00535185
- Idea
- signature-conditioned MDN, lr=0.003: MSE mean=0.00744951, std=0.00325745; paired delta=-0.00376669, p=0.17915, 5/8 wins
Benchmark result
Ahead of the baseline (−33.6%) but not statistically significant (p=0.1792, wins 5 / 8 paired seeds).
- Benchmark
- Dynamics forecasting
- Model
- rnn_small
- Paired seeds
- 8
- Baseline mean
- 0.0112
- Idea mean
- 0
- Effect (Δ)
- -0.0038 (−33.6%; negative = idea better)
- Wins
- 5 / 8 paired seeds
- p-value
- 0.1792 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±51.3%
- Mechanism
- Not confirmed ✗
- Practical effect
- Inconclusive
- Baseline tuning
- swept over 3 configs
The completed protocol used 400 training samples, 200 test samples, 6 epochs, and batch size 128 because the initial larger run exceeded the time budget. Only the built-in dynamics track was tested; particle-rollout replacement, Wasserstein distance, downstream analytic reward, and larger-scale world-model experiments were not tested.
How to run: python3 bench_experiment.py
Artifacts
- bench_experiment.py 8.2 KB View Raw
- bench_report.json 3.8 KB View Raw
- experiment.py 7.4 KB View Raw
- report.md 1.5 KB View
- report_bench_2026-09-01T024822.md 4.5 KB View
- results.json 1.6 KB View Raw