Primal-Dual Flow Matching for Sample-Wise Constrained Generation

Zhengyan Huan; Peter Y. Lu; Shuchin Aeron

cs.LGOct 8, 2026

Primal-Dual Flow Matching for Sample-Wise Constrained Generation

Authors: Zhengyan Huan, Peter Y. Lu, Shuchin Aeron

Organizations: Department of Electrical and Computer Engineering, Tufts University · NSF AI Institute for Artificial Intelligence and Fundamental Interactions

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.

Figures & tables

		FM	MDM	NAMM	PDM	PCFM	FM-RE	PDFM
$8$ -D $\ell_{2}$ ball	$P(\mathcal{C})$ $(\uparrow)$	$97.633\%$	$100\%$	$99.986\%$	$100\%$	$100\%$	$99.923\%$	$99.906\%$
$8$ -D $\ell_{2}$ ball	SWD $(\downarrow)$	$0.0193$	$0.0217$	$0.0228$	$0.0200$	$0.0284$	$0.0194$	$0.0172$
$20$ -D $\ell_{2}$ ball	$P(\mathcal{C})$ $(\uparrow)$	$90.918\%$	$100\%$	$99.711\%$	$100\%$	$100\%$	$99.748\%$	$99.755\%$
$20$ -D $\ell_{2}$ ball	SWD $(\downarrow)$	$0.0087$	$0.0144$	$0.0205$	$0.0257$	$0.0148$	$0.0132$	$0.0101$

Table 1: Results on

8

-D and

20

-D constrained generation tasks.

P(\mathcal{C})

denotes the mean constrained-satisfaction rate of

10000

samples, averaged over

100

generation trials. SWD is computed with

10000

generated samples and samples from

Q

, averaged over

100

generation trials.

Method	Case 1 $(Q(\mathcal{C})=1)$			Case 2 $(Q(\mathcal{C})\approx 39.47\%)$
Method	$P(\mathcal{C})$ $(\uparrow)$	Avg $\delta_{L_{n}}$ $(\downarrow)$	HE $(\downarrow)$	$P(\mathcal{C})$ $(\uparrow)$	Avg $\delta_{L_{n}}$ $(\downarrow)$	$\mathrm{HE}_{\mathrm{clean}}$ $(\downarrow)$	$\mathrm{HE}_{\mathrm{deg}}$ $(\downarrow)$
FM ( $Q$ )	14.19%	0.0409	(0.0087,0.0453)	9.82%	0.0622	(0.0094,0.0465)	(0.0086,0.0565)
FM ( $Q^{\prime}$ )	N/A	N/A	N/A	33.24%	0.0539	(0.0135,0.0477)	(0.0114,0.0489)
PCFM	6.03%	0.0484	(0.0156,0.1547)	0.06%	0.1022	(0.0173,0.0729)	(0.0186,0.0823)
FM-RE	43.17%	0.0339	(0.0117,0.0329)	46.33%	0.0513	(0.0118,0.0558)	(0.0097,0.0570)
PDFM	97.74%	0.0235	(0.0107,0.0269)	90.48%	0.0406	(0.0118,0.0435)	(0.0103,0.0457)

Table 2: Comparison across the considered methods for Case 1 and Case 2. Results are means over

10

trials, using

2000

generated samples per method per trial.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Definition	Description
$P,Q$		Measures
$p,q$		Probability densities
$P(\mathcal{C}),Q(\mathcal{C})$	$P(\mathcal{C})=\int_{\mathcal{C}}p(x)dx,Q(\mathcal{C})=\int_{\mathcal{C}}q(x)dx$	constraint satisfaction rates of $P$ and $Q$
$L(p,\lambda)$	$L(p,\lambda)=\mathbb{E}_{x\sim Q}[-\log p(x)]+\lambda\left(b-\int_{\mathcal{C}}p(x)dx\right)$	The Lagrangian
$u$		The velocity field
$g(\lambda)$	$g(\lambda)=\inf_{p}L(p,\lambda)$	The dual function

Appendix

Table 3: Definitions and descriptions of the notations used in this paper.

Panel of Figure 1 — Figure 1: Histogram plots of $50000$ samples from the compared distributions.

		FM	MDM	NAMM	PDM	PCFM	FM-RE	PDFM
$8$ -D $\ell_{2}$ ball	$P(\mathcal{C})$ $(\uparrow)$	$97.633\%$ [-1pt] $\pm\,0.143\%$	$100\%$ [-1pt] $\pm\,0$	$99.986\%$ [-1pt] $\pm\,0.008\%$	$100\%$ [-1pt] $\pm\,0$	$100\%$ [-1pt] $\pm\,0$	$99.923\%$ [-1pt] $\pm\,0.025\%$	$99.906\%$ [-1pt] $\pm\,0.026\%$
$8$ -D $\ell_{2}$ ball	SWD $(\downarrow)$	$0.0193$ [-1pt] $\pm\,0.0029$	$0.0217$ [-1pt] $\pm\,0.0030$	$0.0228$ [-1pt] $\pm\,0.0037$	$0.0200$ [-1pt] $\pm\,0.0031$	$0.0284$ [-1pt] $\pm\,0.0042$	$0.0194$ [-1pt] $\pm\,0.0029$	$0.0172$ [-1pt] $\pm\,0.0028$
$20$ -D $\ell_{2}$ ball	$P(\mathcal{C})$ $(\uparrow)$	$90.918\%$ [-1pt] $\pm\,0.318\%$	$100\%$ [-1pt] $\pm\,0$	$99.711\%$ [-1pt] $\pm\,0.053\%$	$100\%$ [-1pt] $\pm\,0$	$100\%$ [-1pt] $\pm\,0$	$99.748\%$ [-1pt] $\pm\,0.052\%$	$99.755\%$ [-1pt] $\pm\,0.0051\%$
$20$ -D $\ell_{2}$ ball	SWD $(\downarrow)$	$0.0087$ [-1pt] $\pm\,0.0008$	$0.0144$ [-1pt] $\pm\,0.0012$	$0.0205$ [-1pt] $\pm\,0.0008$	$0.0257$ [-1pt] $\pm\,0.0018$	$0.0148$ [-1pt] $\pm\,0.0009$	$0.0132$ [-1pt] $\pm\,0.0010$	$0.0101$ [-1pt] $\pm\,0.0007$

Appendix

Table 4: Results on

8

-D and

20

-D constrained generation tasks with standard deviations.

P(\mathcal{C})

denotes the mean constrained-satisfaction rate of

10000

samples, averaged over

100

generation trials. SWD is computed with

10000

generated samples and samples from

Q

, averaged over

100

generation trials.

Method	Case 1 $(Q(\mathcal{C})=1)$			Case 2 $(Q(\mathcal{C})\approx 39.47\%)$
Method	$P(\mathcal{C})$ $(\uparrow)$	Avg $\delta_{L_{n}}$ $(\downarrow)$	HE $(\downarrow)$	$P(\mathcal{C})$ $(\uparrow)$	Avg $\delta_{L_{n}}$ $(\downarrow)$	$\mathrm{HE}_{\mathrm{clean}}$ $(\downarrow)$	$\mathrm{HE}_{\mathrm{deg}}$ $(\downarrow)$
FM ( $Q$ )	$14.19\%$ [-1pt] $\pm\,0.732\%$	$0.0409$ [-1pt] $\pm\,0.0002$	$(0.0087,0.0453)$ [-1pt] $\pm\,(0.0002,0.0005)$	$9.82\%$ [-1pt] $\pm\,0.767\%$	$0.0622$ [-1pt] $\pm\,0.0005$	$(0.0094,0.0465)$ [-1pt] $\pm\,(0.0002,0.0004)$	$(0.0086,0.0565)$ [-1pt] $\pm\,(0.0002,0.0005)$
FM ( $Q^{\prime}$ )	N/A	N/A	N/A	$33.24\%$ [-1pt] $\pm\,0.470\%$	$0.0539$ [-1pt] $\pm\,0.0001$	$(0.0135,0.0477)$ [-1pt] $\pm\,(0.0002,0.0004)$	$(0.0114,0.0489)$ [-1pt] $\pm\,(0.0002,0.0004)$
PCFM	$6.03\%$ [-1pt] $\pm\,0.383\%$	$0.0484$ [-1pt] $\pm\,0.0001$	$(0.0156,0.1547)$ [-1pt] $\pm\,(0.0002,0.0008)$	$0.06\%$ [-1pt] $\pm\,0.054\%$	$0.1022$ [-1pt] $\pm\,0.0010$	$(0.0173,0.0729)$ [-1pt] $\pm\,(0.0002,0.0005)$	$(0.0186,0.0823)$ [-1pt] $\pm\,(0.0002,0.0005)$
FM-RE	$43.17\%$ [-1pt] $\pm\,1.496\%$	$0.0339$ [-1pt] $\pm\,0.0002$	$(0.0117,0.0329)$ [-1pt] $\pm\,(0.0002,0.0003)$	$46.33\%$ [-1pt] $\pm\,1.007\%$	$0.0513$ [-1pt] $\pm\,0.0003$	$(0.0118,0.0558)$ [-1pt] $\pm\,(0.0002,0.0004)$	$(0.0097,0.0570)$ [-1pt] $\pm\,(0.0002,0.0004)$
PDFM	$97.74\%$ [-1pt] $\pm\,0.531\%$	$0.0235$ [-1pt] $\pm\,0.0001$	$(0.0107,0.0269)$ [-1pt] $\pm\,(0.0003,0.0003)$	$90.48\%$ [-1pt] $\pm\,0.525\%$	$0.0406$ [-1pt] $\pm\,0.0002$	$(0.0118,0.0435)$ [-1pt] $\pm\,(0.0002,0.0004)$	$(0.0103,0.0457)$ [-1pt] $\pm\,(0.0002,0.0004)$

Appendix

Table 5: Comparison across the considered methods for Case 1 and Case 2, evaluated using

2000

generated samples per method in each of

10

trials. Results are reported as the mean and standard deviation.

Method	$P(\mathcal{C})(\uparrow)$	$\mathrm{NLL}_{\mathrm{train}}(\downarrow)$	$\mathrm{NLL}_{\mathrm{feas}}(\downarrow)$
FM	$72.74\%\pm 1.57\%$	$-4.0851\pm 0.0063$	$-4.1525\pm 0.0061$
FM-RE	$88.14\%\pm 1.37\%$	$-0.9254\pm 0.0039$	$-0.9368\pm 0.0057$
PDFM	$99.28\%\pm 1.33\%$	$-3.8904\pm 0.0122$	$-4.0768\pm 0.0130$

Appendix

Table 6: Constraint satisfaction is evaluated using 1000 generated samples per trial.

\mathrm{NLL}_{\mathrm{train}}

and

\mathrm{NLL}_{\mathrm{feas}}

are evaluated using 1000 samples randomly drawn from the full training set and its feasible subset, respectively, per trial. Results are reported as the mean and standard deviation over 10 trials.

Figure 2: Samples generated by FM in Case 1 ( $\tau=0.033,b=98\%$ ). The titles denote the residual $\delta_{L_{n}}$ . The constraint satisfaction rate is $14.19\%$ .

Figure 3: Samples generated by PDFM in Case 1 ( $\tau=0.033,b=98\%$ ). The titles denote the residual $\delta_{L_{n}}$ . The constraint satisfaction rate is $97.74\%$ .

Panel of Figure 4 — Figure 4: Single generated sample and its summary statistic in Case 1. For each three-image set, the first image shows sampled emulator dynamics, and the last two show histograms of summary statistic distributions. First row: samples generated by the generative models. Left: FM, $\delta_{L_{n}}=0.0421>\tau$ . Right: PDFM, $\delta_{L_{n}}=0.0256<\tau$ . Second row: simulator (truth) samples initialized from the same initial states in the first row.

Panel of Figure 5 — Figure 5: Single generated sample and its summary statistic in Case 2. For each three-image set, the first shows sampled emulator dynamics, and the last two show histograms of summary statistic distributions. First row: samples generated by the generative models. Left: FM, $\delta_{L_{n}}=0.0734>\tau$ . Right: PDFM, $\delta_{L_{n}}=0.0354<\tau$ . Second row: simulator (truth) samples initialized from the same initial states in the first row.

Panel of Figure 6 — Figure 6: Summary statistic over $2000$ samples in Case 1. First column: training samples. Second column: FM’s generated samples. Third column: PDFM’s generated samples.

Panel of Figure 7 — Figure 7: Summary statistic over $2000$ samples in Case 2. First column: clean samples. Second column: degraded samples used during training. Third column: FM’s generated samples. Fourth column: PDFM’s generated samples.

Panel of Figure 8 — Figure 8: Generated fingerprint samples from different methods. The title of each sample reports its number of connected components, with infeasible samples highlighted in red.

Figure 9: Visualization of the $3$ -D molecules generated by FM. The titles indicate whether the generated sample satisfies the molecule-stability constraint.

Figure 10: Visualization of the $3$ -D molecules generated by PDFM. The titles indicate whether the generated sample satisfies the molecule-stability constraint.

Constrained Generative Modeling

Constrained generative modeling develops methods that produce samples satisfying structural, physical, semantic, or user-defined requirements.

See all 182 papers

Constrained Flow Matching via Lagrangian Dual Flows
Jul 5, 2026Vince Kurtz, Alexander DavydovConstrained Generative Modeling Flow Matching
DePaul University · Rice University
Constraint-Aware Flow Matching: Decision Aligned End-to-End Training for Constrained Sampling
May 12, 2026Jacob K. Christopher, James E. Warner, Ferdinando FiorettoConstrained Generative Modeling Flow Matching
University of Virginia · NASA Langley Research Center
MintFlow: Minimal Trajectory Intervention for Constrained Flow Matching
Sep 30, 2026Yesom Park, Kelvin Kan, Qifan Chen +2Constrained Generative Modeling Constrained Sampling
Department of Mathematics, University of California, Los Angeles · Brookhaven National Laboratory

Jul 5, 2026J/K move · Enter open · S save

Constrained Flow Matching via Lagrangian Dual Flows

Vince Kurtz, Alexander Davydov

DePaul University · Rice University

Constrained Generative Modeling Flow Matching Generative Modeling Constrained Optimization Flow-Based Generative Modeling

Open paper AskPDF arXiv

Flow matching is a powerful tool for generative modeling, but emerging applications in robotics, planning, and control require inference-time constraints on generated outputs. Such constraints are often complex and highly nonlinear. As a result, methods designed for linear constraints like image inpainting are rarely sufficient, and projection or optimization-based alternatives can be prohibitively expensive. In this paper, we introduce Lagrangian Dual Flows, a new family of constrained generation techniques based on Lagrangian dual dynamics. By flowing a dual co-state alongside generated samples, we can guarantee nonlinear constraint satisfaction without expensive optimization subproblems, pseudoinverses, or projection steps during the denoising process. The resulting constrained generation algorithms are simple, effective, and open new theoretical connections between flow matching and primal-dual methods in numerical optimization.

Abstract

Constrained Generative Modeling

Explore similar work

Constrained Flow Matching via Lagrangian Dual Flows

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

MintFlow: Minimal Trajectory Intervention for Constrained Flow Matching