Control constraint satisfaction in data generation with primal dual flow matching
Primal-Dual Flow Matching for Sample-Wise Constrained Generation
Machine Learning
Summary
Generating new data samples often needs to meet specific rules or constraints, but making sure these rules are met can be tricky. The authors present a method called Primal-Dual Flow Matching (PDFM) that learns to generate data while ensuring the rules hold with a certain user-chosen probability. Their method only needs to know if a sample passes the rule or not, without needing complicated math about the rule's shape. They show that PDFM better satisfies these constraints compared to previous methods without losing the quality of the generated data.
What this means in practice
- •For machine learning engineers: Train generative models that meet specific constraints during sample creation with guaranteed satisfaction rates using limited feedback.
- •For synthetic data providers: Produce synthetic datasets that comply with legal or ethical constraints without needing complex distance calculations or projections.$Commercial implications: Enables building commercial synthetic data generation tools that reliably enforce constraints, appealing to regulated industries.
Authors
Zhengyan Huan, Peter Y. Lu, Shuchin Aeron
Abstract
We consider the problem of learning to generate data while satisfying a prescribed \emph{sample-wise} constraint with a user-specified probability of constraint satisfaction. Starting with a population-level constrained negative log-likelihood minimization and using a flow-matching (FM) parameterization, under regularity assumptions and sufficient expressivity of the velocity field family, we derive guarantees on the probability of constraint satisfaction for the corresponding primal-dual optimization algorithm. We provide a practical implementation of this approach, called Primal-Dual Flow Matching (PDFM), which has direct control over constraint satisfaction rate and requires only binary membership oracle feedback for constraint satisfaction, avoiding the need for differentiable distances, projections, or special structure of the constraint set, such as convexity. We show that PDFM improves constraint satisfaction over vanilla FM and achieves competitive performance against constraint-aware baselines while maintaining distributional fidelity. Our code is available at https://github.com/ZhengyanHuan/PDFM.