# Stochastic Optimal Control Matching

arXiv: https://arxiv.org/abs/2312.02027 · HTML: https://mlsanity.com/paper/2312.02027/

## Key facts

- Cited by 39 papers in the map (data through 7 Oct 2026).
- Builds on 29 of its 88 references that are in the map.
- First posted on arXiv 4 Dec 2023.

## Builds on (most cited first, 29 of 29)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Solving high-dimensional partial differential equations using deep learning](https://mlsanity.com/paper/1707.02568/) (2017)
- [Scalable Gradients for Stochastic Differential Equations](https://mlsanity.com/paper/2001.01328/) (2020)
- [Safe and Efficient Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1606.02647/) (2016)
- [Deep learning-based numerical methods for high-dimensional parabolic partial differential equations and backward stochastic differential equations](https://mlsanity.com/paper/1706.04702/) (2017)
- [Path Integral Sampler: a stochastic control approach for sampling](https://mlsanity.com/paper/2111.15141/) (2021)
- [Neural SDEs as Infinite-Dimensional GANs](https://mlsanity.com/paper/2102.03657/) (2021)
- [An optimal control perspective on diffusion-based generative modeling](https://mlsanity.com/paper/2211.01364/) (2022)
- [Denoising Diffusion Samplers](https://mlsanity.com/paper/2302.13834/) (2023)
- [Improved sampling via learned diffusions](https://mlsanity.com/paper/2307.01198/) (2023)
- [Sticking the Landing: Simple, Lower-Variance Gradient Estimators for Variational Inference](https://mlsanity.com/paper/1703.09194/) (2017)
- [Solving high-dimensional Hamilton-Jacobi-Bellman PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space](https://mlsanity.com/paper/2005.05409/) (2020)
- [Algorithms for Solving High Dimensional PDEs: From Nonlinear Monte Carlo to Machine Learning](https://mlsanity.com/paper/2008.13333/) (2020)
- [Stochastic Optimal Control for Collective Variable Free Sampling of Molecular Transition Paths](https://mlsanity.com/paper/2207.02149/) (2022)
- [Convergence Analysis of Machine Learning Algorithms for the Numerical Solution of Mean Field Control and Games: II -- The Finite Horizon Case](https://mlsanity.com/paper/1908.01613/) (2019)
- [Deep Relaxation: partial differential equations for optimizing deep neural networks](https://mlsanity.com/paper/1704.04932/) (2017)
- [Actor-Critic Method for High Dimensional Static Hamilton--Jacobi--Bellman Partial Differential Equations based on Neural Networks](https://mlsanity.com/paper/2102.11379/) (2021)
- [Adaptive importance sampling for control and inference](https://mlsanity.com/paper/1505.01874/) (2015)
- [A Neural Network Approach for High-Dimensional Optimal Control Applied to Multi-Agent Path Finding](https://mlsanity.com/paper/2104.03270/) (2021)
- [Deep Fictitious Play for Finding Markovian Nash Equilibrium in Multi-Agent Games](https://mlsanity.com/paper/1912.01809/) (2019)
- [High-Dimensional Stochastic Optimal Control using Continuous Tensor Decompositions](https://mlsanity.com/paper/1611.04706/) (2016)
- [Variational approach to rare event simulation using least-squares regression](https://mlsanity.com/paper/1901.09195/) (2019)

## Built on by (most cited first, 30 of 39)

- [Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control](https://mlsanity.com/paper/2409.08861/) (2024)
- [Sequential Controlled Langevin Diffusions](https://mlsanity.com/paper/2412.07081/) (2024)
- [ResearchTown: Simulator of Human Research Community](https://mlsanity.com/paper/2412.17767/) (2024)
- [Q-learning with Adjoint Matching](https://mlsanity.com/paper/2601.14234/) (2026)
- [Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference](https://mlsanity.com/paper/2508.12511/) (2025)
- [Stochastic Optimal Control for Diffusion Bridges in Function Spaces](https://mlsanity.com/paper/2405.20630/) (2024)
- [Proximal Diffusion Neural Sampler](https://mlsanity.com/paper/2510.03824/) (2025)
- [Conditioning Diffusions Using Malliavin Calculus](https://mlsanity.com/paper/2504.03461/) (2025)
- [Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization](https://mlsanity.com/paper/2510.08233/) (2025)
- [Value Gradient Guidance for Flow Matching Alignment](https://mlsanity.com/paper/2512.05116/) (2025)
- [A Taxonomy of Loss Functions for Stochastic Optimal Control](https://mlsanity.com/paper/2410.00345/) (2024)
- [Bridge Matching Sampler: Scalable Sampling via Generalized Fixed-Point Diffusion Matching](https://mlsanity.com/paper/2603.00530/) (2026)
- [Deep Learning for Continuous-Time Stochastic Control with Jumps](https://mlsanity.com/paper/2505.15602/) (2025)
- [Control Consistency Losses for Diffusion Bridges](https://mlsanity.com/paper/2512.05070/) (2025)
- [A brief review of the Deep BSDE method for solving high-dimensional partial differential equations](https://mlsanity.com/paper/2505.17032/) (2025)
- [Iterative Tilting for Diffusion Fine-Tuning](https://mlsanity.com/paper/2512.03234/) (2025)
- [Discrete Adjoint Schr"odinger Bridge Sampler](https://mlsanity.com/paper/2602.08243/) (2026)
- [Adjoint Matching through the Lens of the Stochastic Maximum Principle in Optimal Control](https://mlsanity.com/paper/2604.08580/) (2026)
- [Scalable Maximum Entropy Reinforcement Learning for Diffusion Policies via Adjoint Matching](https://mlsanity.com/paper/2606.22630/) (2026)
- [Conditional Lagrangian Wasserstein Flow for Time Series Imputation](https://mlsanity.com/paper/2410.07550/) (2024)
- [High-dimensional Mean-Field Games by Particle-based Flow Matching](https://mlsanity.com/paper/2512.01172/) (2025)
- [Supervised Guidance Training for Infinite-Dimensional Diffusion Models](https://mlsanity.com/paper/2601.20756/) (2026)
- [Trust Region Q Adjoint Matching](https://mlsanity.com/paper/2605.27079/) (2026)
- [Flow-based generative models as iterative algorithms in probability space](https://mlsanity.com/paper/2502.13394/) (2025)
- [Improved Constrained Generation by Bridging Pretrained Generative Models](https://mlsanity.com/paper/2603.06742/) (2026)
- [Inference-Time Attribute Distribution Alignment for Unconditional Diffusion](https://mlsanity.com/paper/2605.07456/) (2026)
- [Unsupervised Diffusion Solver for Combinatorial Optimization via Combinatorial Adjoint Matching](https://mlsanity.com/paper/2605.30920/) (2026)
- [An Efficient On-Policy Deep Learning Framework for Stochastic Optimal Control](https://mlsanity.com/paper/2410.05163/) (2024)
- [Stein Diffusion Guidance: Training-Free Posterior Correction for Sampling Beyond High-Density Regions](https://mlsanity.com/paper/2507.05482/) (2025)
- [HOTA: Hamiltonian framework for Optimal Transport Advection](https://mlsanity.com/paper/2507.17513/) (2025)

## Abstract (arXiv)

Stochastic optimal control, which has the goal of driving the behavior of noisy systems, is broadly applicable in science, engineering and artificial intelligence. Our work introduces Stochastic Optimal Control Matching (SOCM), a novel Iterative Diffusion Optimization (IDO) technique for stochastic optimal control that stems from the same philosophy as the conditional score matching loss for diffusion models. That is, the control is learned via a least squares problem by trying to fit a matching vector field. The training loss, which is closely connected to the cross-entropy loss, is optimized with respect to both the control function and a family of reparameterization matrices which appear in the matching vector field. The optimization with respect to the reparameterization matrices aims at minimizing the variance of the matching vector field. Experimentally, our algorithm achieves lower error than all the existing IDO techniques for stochastic optimal control for three out of four control problems, in some cases by an order of magnitude. The key idea underlying SOCM is the path-wise reparameterization trick, a novel technique that may be of independent interest. Code at https://github.com/facebookresearch/SOC-matching