# Terminally constrained flow-based generative models from an optimal control perspective

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

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 14 Jan 2026.

## Built on by (most cited first, 4 of 4)

- [Adjoint Matching through the Lens of the Stochastic Maximum Principle in Optimal Control](https://mlsanity.com/paper/2604.08580/) (2026)
- [ProFlow: Zero-Shot Physics-Consistent Sampling via Proximal Flow Guidance](https://mlsanity.com/paper/2601.20227/) (2026)
- [A Quantitative Approximation Framework for Flow Distillation in Diffusion Models](https://mlsanity.com/paper/2606.03820/) (2026)
- [Ensemble Controlled-Flow Filtering for Implicit Data Assimilation](https://mlsanity.com/paper/2607.12975/) (2026)

## Abstract (arXiv)

We address the problem of sampling from terminally constrained distributions with pre-trained flow-based generative models through an optimal control formulation. Theoretically, we characterize the value function by a Hamilton-Jacobi-Bellman equation and derive the optimal feedback control as the minimizer of the associated Hamiltonian. We show that as the control penalty increases, the controlled process recovers the reference distribution, while as the penalty vanishes, the terminal law converges to a generalized Wasserstein projection onto the constraint manifold. Algorithmically, we introduce Terminal Optimal Control with Flow-based models (TOCFlow), a geometry-aware sampling-time guidance method for pre-trained flows. Solving the control problem in a terminal co-moving frame that tracks reference trajectories yields a closed-form scalar damping factor along the Riemannian gradient, capturing second-order curvature effects without matrix inversions. TOCFlow therefore matches the geometric consistency of Gauss-Newton updates at the computational cost of standard gradient guidance. We evaluate TOCFlow on three high-dimensional scientific tasks spanning equality, inequality, and global statistical constraints, namely Darcy flow, constrained trajectory planning, and turbulence snapshot generation with Kolmogorov spectral scaling. Across all settings, TOCFlow improves constraint satisfaction over Euclidean guidance and projection baselines while preserving the reference model's generative quality.