HardFlow: Hard-Constrained Sampling for Flow-Matching Models via Trajectory Optimization

TL;DR

HardFlow introduces a trajectory optimization approach using numerical optimal control to enforce hard constraints in flow-matching generative models, improving sample quality and constraint satisfaction across various applications.

Contribution

The paper presents a novel trajectory optimization framework for hard-constrained sampling in flow-matching models, leveraging control theory to enhance flexibility and efficiency.

Findings

Outperforms existing methods in constraint satisfaction and sample quality.

Effectively applies to robotics, PDE boundary control, and image editing.

Provides theoretical bounds on approximation error.

Abstract

Diffusion and flow-matching have emerged as powerful methodologies for generative modeling, with remarkable success in capturing complex data distributions and enabling flexible guidance at inference time. Many downstream applications, however, demand enforcing hard constraints on generated samples (for example, robot trajectories must avoid obstacles), a requirement that goes beyond simple guidance. Prevailing projection-based approaches constrain the entire sampling path to the constraint manifold, which is overly restrictive and degrades sample quality. In this paper, we introduce a novel framework that reformulates hard-constrained sampling as a trajectory optimization problem. Our key insight is to leverage numerical optimal control to steer the sampling trajectory so that constraints are satisfied precisely at the terminal time. By exploiting the underlying structure of flow-matching models and adopting techniques from model predictive control, we transform this otherwise complex constrained optimization problem into a tractable surrogate that can be solved efficiently and effectively. Furthermore, this trajectory optimization perspective offers significant flexibility beyond mere constraint satisfaction, allowing for the inclusion of integral costs to minimize distribution shift and terminal objectives to further enhance sample quality, all within a unified framework. We provide a control-theoretic analysis of our method, establishing bounds on the approximation error between our tractable surrogate and the ideal formulation. Extensive experiments across diverse domains, including robotics (planning), partial differential equations (boundary control), and vision (text-guided image editing), demonstrate that our algorithm, which we name $\textit{HardFlow}$, substantially outperforms existing methods in both constraint satisfaction and sample quality.