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

2026-01-14

AI Summary

This work addresses the challenge of efficiently sampling from pretrained flow models under terminal constraints. To this end, the authors propose TOCFlow, a method grounded in optimal control theory that defines a value function via the Hamilton–Jacobi–Bellman equation and solves the control problem in a terminal comoving coordinate system to enable geometry-aware guided sampling. A key innovation is the introduction of a closed-form scalar damping factor that explicitly incorporates curvature information along the Riemannian gradient direction on the manifold, achieving higher-order geometric consistency without increasing the computational cost of standard gradient-based guidance. Experiments demonstrate that TOCFlow significantly outperforms Euclidean guidance and projection-based baselines on tasks including Darcy flow simulation, constrained trajectory planning, and turbulent flow snapshot generation, while simultaneously ensuring high constraint satisfaction and sample quality.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Reasoning under Uncertainty: Stochastic Optimization
• Constraint Satisfaction and Optimization: Constraint Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Responsible Web: Human-perceived consequences of algorithmic deployment on the web

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.

Problem

• terminally constrained sampling
• flow-based generative models
• optimal control
• constraint satisfaction
• terminal distribution

Innovation

• optimal control
• flow-based generative models
• Hamilton-Jacobi-Bellman equation
• Riemannian geometry
• terminal constraints