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

2026-05-12

AI Summary

This work addresses the misalignment between training objectives and sampling requirements in existing constrained generative models, which often struggle to simultaneously ensure strict constraint satisfaction and high sample quality. To resolve this issue, the authors propose an end-to-end constraint-aware flow matching framework that explicitly embeds a constraint projection operator into the training objective. This design aligns the generative dynamics with the constrained sampling process, thereby avoiding the distributional shift typically induced by post-hoc projection. The method enables, for the first time, fully differentiable end-to-end optimization across both training and constrained sampling stages. Extensive experiments on three real-world benchmark tasks demonstrate the approach’s generality and effectiveness, achieving significant improvements in both sample quality and constraint satisfaction rates.

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.

Problem

• constrained sampling
• generative models
• training-sampling misalignment
• constraint satisfaction
• sample quality

Innovation

• Constraint-Aware Flow Matching
• constrained sampling
• end-to-end training
• distributional alignment
• generative modeling