# Self-Supervised Amortized Mean-Field Controller

- ID: 2508
- Canonical URL: https://synthcore.org/idea/2508/self-supervised-amortized-mean-field-controller
- API JSON: https://synthcore.org/api/idea/2508.json
- API Markdown: https://synthcore.org/api/idea/2508.md
- Verification status: mech_ok_no_baseline
- Source: [arXiv:2608.18282](https://arxiv.org/abs/2608.18282)
- Category: training
- Solves: speedup, sample-efficiency, scalability
- ML areas: diffusion, transformer, world-model, training
- Math tags: probability, pde, optimization
- Ratings: usefulness 7/10; difficulty 6/10; novelty 6/10

## Idea description

Train one prompt-conditioned controller to solve a distribution of stochastic control tasks directly from the control objective, instead of generating an optimal trajectory dataset for every task. Use the probability-flow velocity to evolve particles deterministically, evaluate running and terminal costs on those particles, and backpropagate through the rollout to learn a reusable operator.

## Mathematical statement

The controlled representative-agent dynamics are \(\mathrm dX_t=v(X_t,t)\,\mathrm dt+\sqrt{2\gamma}\,\mathrm dW_t\), with diffusion coefficient \(\gamma>0\). The population density \(p(x,t)\) satisfies \(\partial_t p+\nabla\cdot(pv)=\gamma\Delta p\), with initial density \(p_0\). The objective is \(J(p,v)=\int_0^T\int_\Omega L(x,v(x,t))p(x,t)\,\mathrm dx\,\mathrm dt+\int_0^T I(p(\cdot,t))\,\mathrm dt+M(p(\cdot,T))\). Since \(\gamma\Delta p=\nabla\cdot(\gamma p\nabla\log p)\), the equivalent continuity equation has deterministic probability-flow velocity \(u(x,t)=v(x,t)-\gamma\nabla_x\log p(x,t)\): \(\partial_t p+\nabla\cdot(pu)=0\). A neural operator \(v_\theta(x,t;c)\), or a flow producing it, is conditioned on task specification \(c\), such as cost parameters, target distributions, obstacle parameters, or a particle prompt. Monte Carlo particles evolved by \(\dot X=u_\theta(X,t;c)\) estimate the objective without requiring precomputed optimal controls.

## Key formulas

- $$\mathrm dX_t=v(X_t,t)\,\mathrm dt+\sqrt{2\gamma}\,\mathrm dW_t$$
- $$\partial_t p(x,t)+\nabla\!\cdot\!(p(x,t)v(x,t))=\gamma\Delta p(x,t),\qquad p(\cdot,0)=p_0$$
- $$J(p,v)=\int_0^T\!\int_\Omega L(x,v(x,t))p(x,t)\,\mathrm dx\,\mathrm dt+\int_0^T I(p(\cdot,t))\,\mathrm dt+M(p(\cdot,T))$$
- $$u_\theta(x,t;c)=v_\theta(x,t;c)-\gamma\nabla_x\log p_\theta(x,t;c),\qquad \dot X_t=u_\theta(X_t,t;c)$$

## Implementation notes

Integrate this at training-time rollout and use the resulting controller for fast inference across many related tasks. Sample a minibatch of tasks \(c_b\), each containing compact parameters or a particle-cloud prompt, and sample initial particles \(X_{b,0}^{(n)}\sim p_{0,b}\). A transformer encoder maps \(c_b\) to a context vector; a pointwise controller head outputs \(v_\theta(x,t;c_b)\). If the density is represented by the invertible flow above, evaluate \(s_\theta=\nabla_x\log p_\theta\); otherwise estimate the score with a denoising score head and include a score-matching auxiliary loss. Use Euler or Heun integration: `for m=0..M-1: t=m*T/M; s=score(X,t,c); v=controller(X,t,c); u=v-gamma*s; X=X+dt*u`. Accumulate `dt * mean(L(X,v))`, estimate \(I(p_t)\) from flow log density or a kernel/entropy proxy, and add the terminal cost \(M(X_T)\). Backpropagate the total Monte Carlo cost; do not use target trajectories. The exact mathematical transformation is the probability-flow drift \(u=v-\gamma\nabla\log p\); numerical integration, finite-particle estimates, and score estimation are empirical approximations. First experiment on 2-D Gaussian initial and target distributions with quadratic control cost and three task parameters: target mean, target covariance, and \(\gamma\). Compare one shared prompt-conditioned model against solving each task independently with the same MLP and against supervised imitation from an optimizer. Evaluate objective value, terminal mismatch, constraint violations, and wall-clock time for 100 new tasks. Success is lower amortized cost after accounting for pretraining, plus substantially reduced per-task optimization time without a significant increase in terminal error.

## Verification

- Status: mech_ok_no_baseline
- Mechanism evidence: yes
- Mechanism confirmed: yes
- Practical verdict: harms
- Verdict: Built a direct-cost, prompt-conditioned probability-flow controller and verified the core mechanism numerically. Across gamma={0,.05,.1,.2,.4}, the measured Gaussian variance derivative matched 2gamma to less than 1.2e-16 absolute error, and the Euler excess variance matched gamma^2 dt^2/S0 to less than 2.4e-16 absolute error; the zero-gamma correction was approximately 1.1e-16. The amortized controller had cost 0.728 versus 0.313 for short independent per-task optimization, so this toy run confirms the mechanism but does not demonstrate an overall amortized objective win after pretraining.

### Mechanism check

- Verdict: Built a direct-cost, prompt-conditioned probability-flow controller and verified the core mechanism numerically. Across gamma={0,.05,.1,.2,.4}, the measured Gaussian variance derivative matched 2gamma to less than 1.2e-16 absolute error, and the Euler excess variance matched gamma^2 dt^2/S0 to less than 2.4e-16 absolute error; the zero-gamma correction was approximately 1.1e-16. The amortized controller had cost 0.728 versus 0.313 for short independent per-task optimization, so this toy run confirms the mechanism but does not demonstrate an overall amortized objective win after pretraining.
- Confidence: 9/10
- Limitations: The controller comparison is a small 2-D diagonal-Gaussian toy study, uses an empirical Gaussian score from finite particles, and does not test obstacles, non-Gaussian densities, supervised imitation, 100-task wall-clock accounting, Heun integration, or a rigorous amortized break-even analysis. Independent optimization was evaluated on 8 tasks rather than the full task distribution to limit runtime.

### Practical benchmark

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

## Artifacts

- [bench_report.json](https://synthcore.org/code/965/bench_report.json)
- [bench_stage2.py](https://synthcore.org/code/965/bench_stage2.py)
- [experiment.py](https://synthcore.org/code/965/experiment.py)
- [report.md](https://synthcore.org/code/965/report.md)
- [report_bench_2026-09-01T170449.md](https://synthcore.org/code/965/report_bench_2026-09-01T170449.md)
- [results.json](https://synthcore.org/code/965/results.json)
- [results_fixed.json](https://synthcore.org/code/965/results_fixed.json)
- [Download all files as ZIP](https://synthcore.org/download/965)

## Disclaimer

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