Constrained Flow Matching via Lagrangian Dual Flows

Abstract: Flow matching is a powerful tool for generative modeling, but emerging applications in robotics, planning, and physics require inference-time constraints on generated outputs. Such constraints are often complex and highly nonlinear. As a result, methods designed for linear constraints like image inpainting are rarely sufficient, and projection or optimization-based alternatives can be prohibitively expensive. In this paper, we introduce Lagrangian Dual Flows, a new family of constrained generation techniques based on Lagrangian dual dynamics. By simply flowing a dual co-state alongside generated samples, we can guarantee nonlinear constraint satisfaction without expensive optimization subproblems, pseudoinverses, or projection steps during the denoising process. The resulting constrained generation algorithms are simple, effective, and open new theoretical connections between flow matching and primal-dual methods in numerical optimization.

• The paper introduces Lagrangian Dual Flows (LDF) to integrate dual variable dynamics into flow matching, eliminating expensive post-hoc projections.
• It provides theoretical convergence guarantees by ensuring constraint residual decay and asymptotic feasibility for both nonlinear equality and inequality constraints.
• Empirical results show that LDF offers significant computational efficiency and improved constraint satisfaction in applications like image inpainting and robotic trajectory generation.

Lagrangian Dual Flows for Constrained Flow Matching

Overview

The paper "Constrained Flow Matching via Lagrangian Dual Flows" (2607.04513) proposes a new approach for enforcing nonlinear equality and inequality constraints in generative flow models at inference time. The Lagrangian Dual Flows (LDF) framework augments the standard flow matching ODE with dual variables, effectively transforming constraint satisfaction in generative modeling from a challenging optimization or projection problem into the joint integration of an augmented ODE system. The method is efficient, general (supports arbitrary nonlinear constraints), theoretically well-justified, and provides explicit convergence guarantees.

Background and Motivation

Flow matching learns a time-dependent vector field $v_\theta(x, t)$ to transport samples from a reference distribution (typically Gaussian) to the target data distribution by integrating the ODE $\dot{x}_t = v_\theta(x_t, t)$. This construction enables tractable sample generation, but many downstream applications—robotics, physical simulation, image inpainting—require satisfying hard constraints on generated outputs.

Existing post-hoc constraint mechanisms are either:

• Limited to linear constraints (e.g., pseudoinverse-guided methods [pokle2024training])
• Computationally expensive for nonlinear constraints due to repeated projection/optimization at each integration step (e.g., physics-constrained flow matching, PCFM [utkarsh2025pcfm])
• Incurring only approximate satisfaction via penalty or guidance schemes, with exact satisfaction requiring ill-conditioned or stiff ODEs

LDF is designed to overcome these limitations by ensuring guaranteed asymptotic constraint feasibility for general nonlinear constraints while preserving the efficiency of standard ODE-based generation.

Lagrangian Dual Flows: Formulation and Theoretical Guarantees

LDF augments the state space of the generative ODE by adding dual variables $\lambda$ (for equality constraints) and, if needed, slack variables $s$ (for inequalities). The modified ODE for equality constraints $g(x) = 0$ is:

$$\begin{aligned} \dot{x}_t &= v_\theta(x_t, t) - J_g(x_t)^\top \lambda_t - c J_g(x_t)^\top g(x_t) \ \dot{\lambda}_t &= \frac{g(x_t)}{(1-t)^p} \end{aligned}$$

where $J_g(x)$ is the Jacobian of $g$, $c > 0$ is a penalty parameter, and $p=2$ is the theoretically justified exponent controlling the singular accumulation as $\dot{x}_t = v_\theta(x_t, t)$0. For inequality constraints $\dot{x}_t = v_\theta(x_t, t)$1, slack variables $\dot{x}_t = v_\theta(x_t, t)$2 are introduced to recast the problem via $\dot{x}_t = v_\theta(x_t, t)$3, and an additional soft-projection ODE keeps $\dot{x}_t = v_\theta(x_t, t)$4 nonnegative.

This design ensures:

• No explicit pseudoinverse, projection, or optimization at each step
• Full support for general nonlinear constraints, including those only known at inference time
• Compatibility with arbitrary pretrained flow models

The paper rigorously proves that, under standard regularity and constraint qualification conditions, the constraint residual decays as $\dot{x}_t = v_\theta(x_t, t)$5 with $\dot{x}_t = v_\theta(x_t, t)$6 (affine constraints are fully covered), ensuring that feasible samples are produced as $\dot{x}_t = v_\theta(x_t, t)$7.

Comparison to Baselines and Empirical Evaluation

LDF is compared to pseudoinverse-based guidance [pokle2024training] and physically-constrained flow matching (PCFM) [utkarsh2025pcfm], both theoretically and empirically. LDF avoids the computational overhead of repeated matrix inversions (pseudoinverse guidance) and expensive step-wise projections/optimizations (PCFM), instead requiring only a single vector-Jacobian product that is accessible via autodiff.

Empirical studies include:

• Nonlinear equality constraint: Pretrained generative models producing "star" shaped distributions subsequently constrained to lie on a unit circle.
• Nonlinear inequality constraint: The same model constrained to one half-plane.
• Linear equality constraint (image inpainting): Pretrained model conditioned to match observed pixels in MNIST digit images.

Figure 1: Generated star examples subject to a unit circle constraint. Compares the proposed LDF, pseudoinverse guidance, and PCFM.

In all settings, LDF achieves constraint violations comparable to or better than soft penalty or pseudoinverse guidance methods, and is significantly more computationally efficient than projection/optimization-based baselines. For example, Table 1 (in the paper) shows that LDF achieves low constraint violation without the high inference time or stability issues that afflict PCFM for nonlinear constraints.

Figure 2: A pre-trained model generates MNIST handwritten digits; applying pseudoinverse guidance, PCFM, and LDF achieves image inpainting by matching the upper half of each image.

Figure 3: The original PCFM algorithm (left) produces noisy generation paths under nonlinear constraints; repeated linearization/projection (right) improves stability, but increases cost.

Numerical and Implementation Considerations

Empirical investigations quantify the tradeoff between the penalty $\dot{x}_t = v_\theta(x_t, t)$8 and ODE stiffness. While larger $\dot{x}_t = v_\theta(x_t, t)$9 improves constraint tightness, it increases ODE stiffness, necessitating careful integrator tuning (e.g., higher-order, adaptive solvers). The dual variable dynamics critically reduce stiffness compared to naive penalty-only approaches—enabling tight feasibility with manageable integration steps.

Furthermore, LDF requires only modifying the ODE right-hand-side evaluation to include dual variable dynamics, making it immediately applicable to any pretrained flow matching model (as shown by open-source JAX+Flax+Diffrax implementation).

Implications and Future Research

Practical implications: LDF gives a computationally affordable pathway to infuse arbitrary hard constraints into generative policies—immediately benefiting fields such as:

• Robotics: Hard safety constraints for action/trajectory generation
• Physics-informed learning: Conservation laws, boundary constraints
• Inverse imaging/generative design: Arbitrary mask or domain constraints

Theoretical implications: LDF bridges flow matching with the continuous-time theory of primal-dual and saddle-point optimization, introducing the possibility to port mature convergence and optimality guarantees into the constrained generative modeling domain.

Limitations and Open Directions:

• LDF guarantees only feasibility, not full conditional distribution matching—samples may cluster near constraint boundaries.
• Increasing $\lambda$0 improves tightness but induces ODE stiffness; optimization of integration strategies and adaptive schemes is warranted.
• Direct extension to efficient conditional generation (density adjustment to the constraint manifold) remains an open problem.
• Systematic integration with structured control-theoretic methods and application to high-dimensional, real-world systems can further amplify LDF's impact.

Conclusion

Lagrangian Dual Flows constitute a rigorous, flexible, and efficient framework for enforcing general nonlinear constraints in pretrained flow-matching generative models. By executing primal-dual ODE integration at inference time, LDF circumvents key drawbacks of projection and pseudoinverse-based approaches and achieves strong empirical and theoretical guarantees for constraint satisfaction. The technical bridge to continuous-time optimization theory opens substantial new avenues in both applied generative modeling and algorithmic theory—suggesting future work in distributional conditioning and broader application domains.