⚡ electrical engineering

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

HardFlow is a novel framework that reformulates hard-constrained sampling in flow-matching models as a trajectory optimization problem, using numerical optimal control to ensure precise constraint satisfaction at the terminal time while maintaining high sample quality.

Original authors: Zeyang Li, Kaveh Alim, Navid Azizan

Published 2026-04-28

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Zeyang Li, Kaveh Alim, Navid Azizan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are using a high-tech GPS to drive a car from your house (the "starting point") to a beautiful destination (the "target data").

In the world of AI, "Flow-Matching" models are like these high-tech GPS systems. They know the general route to get you to a perfect destination—like a realistic image of a cat or a smooth path for a robot arm.

The Problem: The "Strict Traffic Cop"

Usually, when we want to add rules—like "don't drive through that park" or "don't hit that wall"—current AI methods use a "Strict Traffic Cop" approach (called Projection-based sampling).

This cop is incredibly annoying. Every single second you are driving, the cop grabs your steering wheel and forces you to stay on a very specific, narrow path. Because the cop is so obsessed with the rules during the whole trip, you end up driving in a weird, jerky, unnatural way. You might reach the destination, but you arrive exhausted, and the journey was so restricted that you missed all the scenic views (this is "degraded sample quality").

The Solution: The "Professional Race Car Driver" (HardFlow)

The authors of this paper, HardFlow, suggest a much smarter way. Instead of a nagging cop, they imagine you are a Professional Race Car Driver using Trajectory Optimization.

Here is how the "HardFlow" driver works:

Focus on the Finish Line: The driver doesn't care if you drift a little bit near a flowerbed in the middle of the trip. They only care about one thing: When you cross the finish line, you must be exactly where you are supposed to be, and you must not have hit anything.
The "Predictive" Brain: Instead of reacting to every tiny bump, the driver has a "mental simulator" (called Model Predictive Control). Before they turn the wheel, they quickly run a "what if" scenario in their head: "If I turn left now, where will I end up in ten minutes?"
The Smooth Correction: If the driver realizes, "Uh oh, if I keep going this way, I'm going to hit that wall at the end," they don't panic and jerk the wheel. They make tiny, smooth, calculated adjustments to the steering to ensure they glide perfectly into the target zone.

Why is this better?

By treating the problem like a "planned journey" rather than a "series of forced corrections," HardFlow achieves two things at once:

Perfect Obedience: It hits the "hard constraints" (like avoiding obstacles or keeping a face looking like the same person) almost perfectly.
High Quality: Because the "path" isn't being constantly mangled by a strict cop, the final result looks natural, beautiful, and high-quality.

Real-World Examples from the Paper:

Robotics: A robot arm learns to reach a target without ever bumping into a pillar.
Maze Navigation: A digital ball zips through a complex maze to a goal without ever touching the walls.
Image Editing: You can tell an AI, "Make this person smile," but the "HardFlow" driver ensures the person still looks like the same person (it keeps the "identity" constraint hard), whereas older methods might accidentally turn them into a completely different human!

In short: HardFlow stops micromanaging the journey and starts optimizing the destination.

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

HardFlow is a novel framework designed to address the challenge of enforcing hard constraints (inviolable conditions) during the inference stage of generative models, specifically Flow-Matching (FM) and Diffusion models.

1. The Problem: Limitations of Current Methods

Existing methods for constrained sampling generally fall into two categories, both of which have fundamental flaws:

Projection-Based Methods: These methods (e.g., projecting the sample back to a feasible set at each step or at the end) often enforce constraints on the entire sampling path. This is overly restrictive because intermediate states are merely algorithmic artifacts; only the final sample needs to satisfy the constraint. Constraining the path prematurely prunes the search space, leading to degraded sample quality and distribution shift.
Soft Guidance (Gradient-based): These methods treat constraints as penalties in a reward function. While they encourage desirable attributes, they cannot guarantee that hard constraints (like obstacle avoidance in robotics) are strictly met.

2. Methodology: Trajectory Optimization & MPC

The authors propose reformulating constrained sampling as a local optimal control problem. Instead of forcing the path to stay within a manifold, they allow the path to move freely, only requiring the terminal state to satisfy the constraints.

Key Mathematical Formulation:
The problem is framed as finding a control input $u_t$ that perturbs the learned velocity field $v_\theta(x)$ to minimize a terminal cost $C(x_1)$ and a control effort (to stay close to the original distribution), subject to the terminal constraint $h(x_1) \leq 0$ .

The HardFlow Algorithm (The "Direct Method" Approach):
Solving this continuous-time optimization is computationally intractable due to the non-linear neural dynamics. The authors derive a scalable algorithm through several transformations:

Discretization: Using the forward Euler method to turn the ODE into a discrete-time problem.
MPC Decomposition: Using Model Predictive Control (MPC) principles to break the long-horizon problem into a sequence of single-step subproblems.
Posterior Mean Proxy: Leveraging the unique structure of Flow-Matching to use the posterior mean $M_\theta(x)$ as an efficient proxy for the predicted terminal state.
Reverse Reparameterization: To avoid the complex, non-convex feasible sets created by propagating constraints backward through the neural network, they transform the decision variable from the current state to the predicted terminal state.
One-Step Fixed-Point Iteration: To keep the algorithm efficient, they use a linear approximation of the inverse mapping, requiring only a single forward pass of the neural network per step.

3. Key Contributions

New Paradigm: Shifts the perspective of constrained sampling from "path projection" to "trajectory optimization."
Unified Framework: A single formulation that jointly handles hard constraints (feasibility), terminal costs (sample quality), and integral costs (minimizing distribution shift).
Tractable Surrogate: Derives a highly efficient algorithm (HardFlow) that avoids the need for expensive second-order derivatives (Jacobians) of the neural network during optimization.
Theoretical Guarantees: Provides a control-theoretic analysis that quantifies the approximation error between their efficient surrogate and the ideal optimal control formulation.

4. Experimental Results

The authors evaluated HardFlow across four diverse domains, consistently outperforming state-of-the-art baselines (including Projection-based and Gradient-guidance methods):

Robotics (Manipulation & Maze Navigation): HardFlow achieved a 100% safety rate (zero collisions/violations) while maintaining higher sample quality (shorter paths/higher scores) and lower computational overhead than most baselines.
PDE Control (Burgers' Equation): Successfully enforced time-varying boundary constraints with minimal control energy.
Vision (Text-Guided Image Editing): Achieved 100% identity preservation (satisfying LPIPS constraints) while effectively following text prompts (high CLIP scores), avoiding the "over-editing" seen in gradient-based methods.

5. Significance

HardFlow represents a significant step forward in controllable generation. By treating generative sampling as a trajectory optimization problem, it provides a principled way to bridge the gap between high-quality generative modeling and the strict safety/physical requirements of real-world applications like robotics and physical simulations. It demonstrates that "direct methods" from optimal control can be made highly efficient for high-dimensional neural dynamics.