Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Zeyang Li, Kaveh Alim, Navid Azizan

TL;DR

This work tackles hard-constrained sampling in flow-matching and diffusion models by reframing it as terminal-time trajectory optimization. HardFlow uses model-predictive-control-inspired recursion and a posterior-mean terminal-state proxy to solve a tractable surrogate that enforces feasibility at $t=1$ while optionally minimizing distribution shift and improving sample quality. The authors provide theoretical bounds on the surrogate's suboptimality and demonstrate superior constraint satisfaction and sample quality across robotics, PDE control, and vision tasks compared to projection-based and soft-guidance baselines. The approach offers a training-free, flexible framework for controllable generation with hard guarantees at inference time.

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.

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

TL;DR

This work tackles hard-constrained sampling in flow-matching and diffusion models by reframing it as terminal-time trajectory optimization. HardFlow uses model-predictive-control-inspired recursion and a posterior-mean terminal-state proxy to solve a tractable surrogate that enforces feasibility at while optionally minimizing distribution shift and improving sample quality. The authors provide theoretical bounds on the surrogate's suboptimality and demonstrate superior constraint satisfaction and sample quality across robotics, PDE control, and vision tasks compared to projection-based and soft-guidance baselines. The approach offers a training-free, flexible framework for controllable generation with hard guarantees at inference time.

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

Paper Structure

This paper contains 17 sections, 7 theorems, 57 equations, 8 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Background
  4. Problem Formulation
  5. Algorithm Design
  6. Theoretical Analysis
  7. Experiments
  8. Robotic Manipulation
  9. Maze Navigation
  10. PDE Control
  11. Text-Guided Image Editing
  12. Conclusion
  13. Experiment Details
  14. Robotic Manipulation
  15. Maze Navigation
  16. ...and 2 more sections

Key Result

Lemma 1

Let $Z=(X_0,X_1)\sim\pi_{0,1}$. Consider the affine conditional path $X_t\mid Z=\alpha_t X_1+\beta_t X_0$ for $t\in[0,1]$ with differentiable scheduler $(\alpha_t,\beta_t)$. The conditional velocity field is therefore $v_{t\mid Z}(X_t\mid Z)=\dot{\alpha}_t X_1+\dot{\beta}_t X_0$. Suppose $v_t(x)$ is Additionally, the following identity holds:

Figures (8)

  • Figure 1: Roadmap of problem transformations. Double arrows indicate equivalence; single arrows indicate approximation.
  • Figure 2: The robotic manipulation task.
  • Figure 3: Visualization of HardFlow in the robotic manipulation task. The green region denotes the target area, the red circles represent original obstacles, and the purple region denotes the novel obstacles introduced at test time. The robot successfully weaves around all obstacles to reach the target.
  • Figure 4: Visualization of HardFlow in the maze navigation task. The red region denotes the introduced obstacles. The robot successfully navigates from the start position to the goal position while avoiding obstacles.
  • Figure 5: Heatmaps of controlled state $u$ and predicted control $f$ produced by HardFlow at one trial in the PDE control task. The grid uses $m=10$ time steps and $n=128$ spatial points; the axis ticks reflect this discretization.
  • ...and 3 more figures

Theorems & Definitions (20)

  • Remark 1
  • Lemma 1: Posterior mean feng2025on
  • Proposition 1
  • proof
  • Remark 2
  • Remark 3
  • Remark 4
  • Proposition 2
  • proof
  • Proposition 3
  • ...and 10 more