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

2025-11-11

AI Summary

Existing diffusion and flow-matching models struggle to maintain generation quality under strict hard constraints—e.g., obstacle avoidance in robot trajectory planning—because conventional projection-based methods impose constraints too rigidly across the entire sampling trajectory. This work proposes a trajectory-optimization-based framework for hard-constrained generation: it reformulates constraint satisfaction as an optimal control problem with exact terminal-state controllability, jointly optimizing an integral cost and a terminal objective. Leveraging a flow-matching architecture, the method integrates model predictive control (MPC) with numerical optimal control to construct a tractable surrogate optimization problem. Evaluated on robotic motion planning, PDE boundary control, and image editing, the approach significantly improves both constraint satisfaction rates and sample fidelity. To our knowledge, it is the first method to simultaneously guarantee high-quality generation and rigorous terminal-state hard constraints.

Technology Category

• Constraint Satisfaction and Optimization: Constraint Optimization
• Search and Optimization: Sampling/Simulation-based Search
• Computer Vision: Diffusion Models for Vision

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization

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 $ extit{HardFlow}$, substantially outperforms existing methods in both constraint satisfaction and sample quality.

Problem

• Enforcing hard constraints on samples from flow-matching models
• Reformulating constrained sampling as trajectory optimization problem
• Ensuring precise constraint satisfaction while maintaining sample quality

Innovation

• Reformulates hard-constrained sampling as trajectory optimization
• Leverages numerical optimal control for terminal constraint satisfaction
• Transforms complex optimization into tractable surrogate via MPC