Constraint-Aware Flow Matching: Decision Aligned End-to-End Training for Constrained Sampling

Abstract: Deep generative models provide state-of-the-art performance across a wide array of applications, with recent studies showing increasing applicability for science and engineering. Despite a growing corpus of literature focused on the integration of physics-based constraints into the generation process, existing approaches fail to enforce strict constraint satisfaction while maintaining sample quality. In particular, training-free constrained sampling methods, while providing per-sample feasibility guarantees, introduce a fundamental mismatch between the training objective and the constrained sampling procedure, often leading to performance degradation. Identifying this training-sampling misalignment as a central limitation of current constrained generative modeling approaches, this paper proposes Constraint-Aware Flow Matching, a novel end-to-end framework that explicitly incorporates constraint projections into the training objective. By aligning the model's learned dynamics with the constrained sampling process, the proposed method mitigates distributional shift induced by projection-based corrections, enabling high-quality constrained generation. The proposed approach is evaluated on three challenging real-world benchmarks, illustrating the generality and efficacy of the method.

• The paper introduces CAFM, a novel end-to-end framework that integrates differentiable constraint projections to align generative flows with feasibility requirements.
• It reformulates the flow matching loss to target the projected state, ensuring high-density sampling and improved data fidelity across scientific tasks.
• Empirical evaluations on PDE-constrained generation and material design demonstrate CAFM's superior accuracy and robust constraint satisfaction compared to prior methods.

Constraint-Aware Flow Matching: End-to-End Training for Constrained Generative Sampling

Motivation and Limitations of Existing Approaches

Constrained generative modeling remains essential in scientific and engineering domains where strict adherence to physical or domain constraints is non-negotiable. Despite the success of flow matching and diffusion models in diverse generative tasks, their conventional training produces models that naturally assign probability mass to infeasible regions, making post-hoc correction necessary at inference. Existing methods are categorized into three classes: (1) constraint guidance via conditional generation (lacking feasibility guarantees), (2) physics-informed training with residual losses (offering distributional alignment but not per-sample guarantees), and (3) constrained sampling, which projects generated endpoints onto feasible manifolds during inference for strict constraint satisfaction.

A fundamental technical gap is illustrated: training-free constrained sampling methods leverage pretrained models, with hard projections only at inference. This introduces a misalignment between training and deployment, resulting in suboptimal data fidelity as projections push samples into low-density regions, away from the natural support of the learned distribution.

CAFM: Technical Contributions and Objective Formulation

The paper presents Constraint-Aware Flow Matching (CAFM), a novel end-to-end framework that directly integrates constraint projections into the training objective, fully aligning the learned dynamics with the downstream constrained sampling process. CAFM is inspired by decision-focused learning and differentiable optimization, using the principle that optimization layers can be embedded within neural architectures and differentiated through, allowing constraints to be enforced during training and not just inference.

CAFM reformulates the flow matching loss to operate on the projected state. While classical flow matching trains the velocity field to minimize regression error to the reference endpoint, CAFM redefines the loss so that the target is the feasible projection of the generative trajectory. For constraints described by equality and inequality conditions $h(z_1) = 0$ and $g(z_1) \leq 0$, the projection operator $P_c$ is implemented as a differentiable layer via iterative solvers (e.g., sequential quadratic programming), unrolled into the computational graph.

The constraint-aware loss is given by:

$$\mathcal{L}_{\mathrm{CAFM}} = \| P_c(z_0 + v_\theta(Z_t, t)) - Z_1 \|^2$$

where the projection $P_c$ is defined through an iterative process compatible with gradient-based optimization. This objective function ensures that the training process produces generative flows whose terminal projections reside in high-density regions of the data distribution, optimizing for data fidelity post-projection rather than solely unconstrained accuracy.

Differentiable Projection Layers and Implementation

Unlike classical proximal operators, which may only apply to convex constraint sets, CAFM generalizes to nonlinear and nonconvex constraints by employing iterative optimization techniques such as SQP. The projection layer operates by integrating the solution process into the neural network's forward pass, allowing efficient gradient computation via implicit differentiation of fixed-point conditions. The approach leverages vector-Jacobian products for scalable backward propagation.

Empirically, CAFM enables efficient warm-starting from pretrained unconstrained flow matchers, mitigating increased training complexity introduced by projection unrolling. Ablation analyses confirm that a small number of unrolling steps are sufficient for tight feasibility while maintaining fidelity gains.

Empirical Evaluations

CAFM was benchmarked on three challenging real-world scientific tasks:

1. PDE-Constrained Generation

Evaluations on Navier-Stokes, Reaction-Diffusion, and Burgers equations impose increasingly complex linear and nonlinear constraints. CAFM consistently achieves lower mean squared error and constraint violation metrics than unconstrained flow matching and prior constrained sampling baselines such as PCFM and ECI. Notably, CAFM demonstrates robust generalization to out-of-distribution initial and boundary conditions, outperforming physics-informed and sampling-based models even under strong structural constraints.

2. Microweather Wind Velocity Field Estimation

For wind field prediction requiring spatiotemporal coherence, CAFM outperforms unconstrained and projected sampling baselines in MMSE and coherence violation metrics. Results indicate that CAFM reduces coherence residuals by an order of magnitude relative to state-of-the-art sampling methods, without significant compromise in variance structure. The constraint-aware objective demonstrates a trade-off in suppressing high-variance features to stabilize projection, but achieves dominant reconstruction and feasibility.

3. Microstructure Inverse-Design

On a material science task with strict porosity constraints, CAFM reliably generates samples with perfect constraint satisfaction and lower reconstruction error than unconstrained and projected diffusion alternatives. This validates CAFM’s applicability to convex count constraints and its potential for practical deployment in inverse design applications with scarce imaging data.

Discussion: Theoretical Implications and Future Directions

CAFM signals a paradigm shift in constrained generative modeling, asserting that post-hoc correction is insufficient for high-fidelity scientific generation. Directly embedding the constraint enforcement in the training objective is critical for maintaining sample quality and feasibility simultaneously. The connection to multistage optimization is highlighted, suggesting extensions to sequential decision processes and policy induction.

Theoretical implications include improved robustness to out-of-distribution constraints and enhanced generalization, as empirical evidence demonstrates superior performance on unseen constraint regimes. The methodology accommodates additional sampling-time projections, enabling hybrid approaches for domains requiring flexible constraint imposition.

The computational overhead of differentiable optimization layers is acknowledged, but practical mitigations such as warm starting and reduced projection iterations are demonstrated.

Conclusion

Constraint-Aware Flow Matching provides an end-to-end solution for constrained generative sampling by embedding differentiable projectors into the training process. Empirical results across diverse scientific tasks support its efficacy in achieving both high accuracy and strict constraint satisfaction. CAFM addresses the central misalignment between training and constrained sampling, marking substantial progress towards developing physically consistent, practically reliable generative models for scientific and engineering domains. Future developments may extend CAFM to multistage, policy-driven optimization architectures and broader constraint families, closing the gap between theory and deployment in constrained generative modeling (2605.12754).