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

TL;DR

This paper formulates terminally constrained sampling with flow-based models as an optimal control problem, deriving a new guidance method (TOCFlow) that improves constraint satisfaction while maintaining generative quality.

Contribution

It introduces TOCFlow, a geometry-aware control-based sampling method for flow models, with a theoretical foundation and practical efficiency for high-dimensional constrained sampling.

Findings

TOCFlow improves constraint satisfaction over baselines.

It preserves the generative quality of pre-trained flow models.

Demonstrated effectiveness on scientific high-dimensional tasks.

Abstract

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.