Research map / Learning theory and algorithms / Prediction and decision losses

Stochastic Optimal Control Matching

NeurIPS 2024 · Learning theory and algorithms › Prediction and decision losses · approach: Losses for decision-making · 1 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes Stochastic Optimal Control Matching (SOCM), an iterative diffusion optimization method that learns the optimal control by least-squares regression onto a matching vector field whose reparameterization matrices are also learned, and reports lower control L2 error than existing IDO losses on four benchmark settings.

Problem

Learning optimal feedback controls for high-dimensional stochastic optimal control problems, where existing IDO losses (adjoint/relative entropy, cross-entropy, variance-based) suffer from non-convex landscapes, high gradient variance and unstable training.

Approach

Replaces the non-convex control objective with a least-squares loss that regresses the control onto a random matching vector field, obtained via a novel path-wise reparameterization trick (introducing learned reparameterization matrices M_t that are optimized to reduce variance). Bias term equals the cross-entropy loss, so SOCM is a variance-reduced version of it.

Key results

• On Quadratic OU (easy), SOCM reaches the lowest L2 control error, about ten times lower than the best existing method (adjoint).
• On Linear OU, SOCM's error is around five times smaller than any existing method.
• In Gaussian mixture sampling at d=128, control objective estimates are 0.146 (Adjoint), 7.49 (SOCM) and 12.61 (cross-entropy), so SOCM fails at high dimension.

Measured effects reported in the paper

Surprising findings

• SOCM's advantage disappears at high dimension: at d=128 SOCM is worse than the adjoint method, due to exploding variance of the importance weight.
• On Double Well, the adjoint method reaches slightly lower error than SOCM but shows instabilities that persist for small learning rates.

Limitations

• The importance weight variance explodes when f or g are large or when dimension is high, limiting SOCM to moderate problems.
• Solving for optimal reparameterization matrices via the Fredholm equation is expensive, so M is parameterized by a neural network instead.

Experimental setup

Datasets: simulated SDE trajectories (Quadratic Ornstein-Uhlenbeck, Linear Ornstein-Uhlenbeck, Double Well, Gaussian mixture sampl

Benchmarks: Quadratic Ornstein-Uhlenbeck (easy) Quadratic Ornstein-Uhlenbeck (hard) Linear Ornstein-Uhlenbeck Double Well Two-mode Gaussian mixture sampling

Metrics: L2 control error control objective time per iteration gradient norm squared

Baselines: Adjoint (relative entropy loss) Cross-entropy loss Log-variance loss Variance loss Moment loss SOCM with M_t=I SOCM-Adjoint

Builds on: cross-entropy loss for control adjoint method (relative entropy loss) path integral sampler Girsanov theorem

Closest papers by meaning

• Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference (NeurIPS 2025 · similarity 0.86)
• Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control (ICLR 2025 · similarity 0.86)
• Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence (ICML 2025 · similarity 0.86)
• Scalable Maximum Entropy Reinforcement Learning for Diffusion Policies via Adjoint Matching (NeurIPS 2026 · similarity 0.85)
• Stochastic Optimal Control for Diffusion Bridges in Function Spaces (NeurIPS 2024 · similarity 0.84)
• Stochastic Optimal Control and Estimation with Multiplicative and Internal Noise (NeurIPS 2024 · similarity 0.84)