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

TL;DR

This work reframes sampling with pre-trained flow-based models as a terminal-constrained optimal control problem, deriving a Hamilton–Jacobi–Bellman framework and identifying an energy interpretation in a co-moving frame. It develops practical, geometry-aware proximal approximations—GD, GN, and the proposed TOCFlow—via a constant-metric reduction that yields a closed-form scalar damping factor, enabling scalable, constraint-preserving guidance at the cost of standard gradient methods. The paper proves asymptotic behaviors (convergence to the reference flow for large penalties and Wasserstein projections for vanishing penalties) and validates TOCFlow across three high-dimensional scientific tasks (Darcy flow, trajectory planning, turbulence) where it consistently improves constraint satisfaction while maintaining generative quality. Collectively, TOCFlow offers a principled, efficient approach to enforcing complex terminal constraints in flow-based generative modeling, with broad applicability to physics-informed sampling and high-dimensional inference.

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.

Figures (9)

• Figure 1: Terminal standard deviation $\tilde{\sigma}$ against the weight schedule $\lambda_t$ under different guidance schemes. Near $\lambda=0$, the GD scheme decays too fast compared with the optimal-control solution, whereas GN and TOCFlow exhibit a comparatively more accurate contraction.
• Figure 2: Qualitative comparison of the generated samples using different methods. The panels display the permeability field $\bm K$ (left) and the pressure field $\bm p$ (middle) for a representative sample. The rightmost plot in each block shows the histogram of PDE residuals aggregated over $512$ generated samples, with the $y$-axis representing density. Top row: Vanilla (left) and GD (right). Bottom row: TOCFlow (left) and terminal projection (right). Both GD and TOCFlow shift the residual distribution toward lower values, indicating better constraint satisfaction.
• Figure 3: Quantitative evaluation of constraint violation and hyperparameter sensitivity.Top: Violin plots of the terminal cost $H(\bm{x})$ (log scale) for $512$ generated samples across all methods. The results show two performance tiers: vanilla/terminal projection (high error) vs. GD/TOCFlow (low error). Bottom left/center: Ablation of the number of lookahead explicit Euler steps $k$ for GD and TOCFlow. Both methods show similar sensitivity, with increased $k$ slightly increasing the cost. Bottom right: Ablation of the number of iterations for terminal projection. The error remains high and constant regardless of the iteration count, indicating that the solver cannot bridge the gap from the unguided distribution to the constraint manifold.
• Figure 4: Visualization of $32$ independent realizations from the synthetic trajectory dataset in \ref{['subsec:geometric_inequality_constraints_trajectory_planning']}.
• Figure 5: Qualitative comparison of the generated trajectories using different methods. Thin lines display $32$ individual samples for clarity. Top row: Samples generated by the vanilla and GD methods. Bottom row: Samples from the TOCFlow, approximated GN and terminal projection methods. The grey shaded regions denote the permissible safety corridors. The vanilla method frequently violates the boundaries. Both GD and TOCFlow successfully steer the trajectory into the valid region while preserving the geometric regularity (smoothness) of the reference distribution. In contrast, the projection-based methods (GN and terminal projection), while achieving strict constraint satisfaction, exhibit non-physical artifacts characterized by sharp kinks at the boundaries where the trajectory is "snapped" to the safety corridors.

Theorems & Definitions (36)

• Proposition 1: Consistency of the conditional mean field albergo2023stochastic
• Remark 1: Validity of the reference dynamics
• Definition 1: Admissible controls
• Definition 2: Objective functional and control energy
• Proposition 2: Hamilton--Jacobi--Bellman equation and optimal control
• Proposition 3: Dynamics in the initial co-moving frame and energy equivalence
• Proposition 4: Energy bounds and Wasserstein distance
• Theorem 1: Large penalty limit and convergence to the reference flow
• Theorem 2: Small penalty limit and selection principle
• Proposition 5: Exact solution for the Gaussian model