Primal-Dual Flow Matching for Sample-Wise Constrained Generation
cs.LG
Submitted: 2026-10-08
Updated: 2026-10-08
Comments: Accepted at NeurIPS 2026
Code: https://github.com/ZhengyanHuan/PDFM
License: http://creativecommons.org/licenses/by/4.0/
The gist: We consider the problem of learning to generate data while satisfying a prescribed sample-wise constraint with a user-specified probability of constraint satisfaction.
Terminology
Abstract
We consider the problem of learning to generate data while satisfying a prescribed 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.
Sources
- HardFlow: Hard-Constrained Sampling for Flow-Matching Models via Trajectory Optimization
- Learning to Emulate Chaos: Adversarial Optimal Transport Regularization
- On Flow Matching KL Divergence
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks