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

TL;DR

This paper introduces Constraint-Aware Flow Matching, an end-to-end training framework that explicitly incorporates constraints into generative models, improving the quality and feasibility of constrained sampling in scientific applications.

Contribution

It proposes a novel end-to-end approach that aligns training with constrained sampling, addressing the mismatch in existing methods and enhancing sample feasibility and quality.

Findings

Effective constrained generation on three real-world benchmarks.

Mitigates distributional shift caused by projection-based corrections.

Demonstrates generality and improved performance over existing methods.

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.