Primal-Dual Flow Matching for Sample-Wise Constrained Generation

1 Introduction

• Scientific and molecular samples may be useful only when they satisfy application-specific constraints as well as distributional requirements.Examples include physical consistency with boundary conditions, conservation laws, or governing PDEs, and chemically useful molecular properties.Source: §1 · p. 1
• The related-work discussion organizes constrained generation by constraint type and by whether methods operate at inference time or during training.It distinguishes sample-wise from distributional constraints and identifies membership-oracle access as one sample-wise setting [Huan et al., 2026].Source: §1 · p. 1
• For sample-wise constraints available only through a membership oracle, the paper develops theoretical guarantees and a practical primal-dual flow-matching method.The practical method uses only binary feedback and provides direct control over the constraint satisfaction rate.Source: §1 · p. 2
• PDFM achieves competitive constraint satisfaction and distributional fidelity, with reported superior performance on chaotic PDE-constrained dynamics and improvements over FM in molecule generation.The molecule-generation comparison concerns stability and distributional fidelity.Source: §1 · p. 2

2 Constrained Generation Setup

• The population-level problem requires the generated distribution to assign at least the target rate b to the application-specific constraint set C.The setup defines C through its indicator and gives examples such as trajectories satisfying a conservation-law tolerance.Source: §2 · p. 2
• For 0 < Q(C) < b, the population-level optimum reweights Q inside and outside C.The section states that this characterization follows from a lemma.Source: §2 · p. 3
• Training on the truncated feasible distribution can reduce effective sample size, degrade distribution matching, and still fail to reach the target with finite-capacity models.These drawbacks are especially relevant when training data are noisy.Source: §2 · p. 3

2.1 Problem Regimes and Learning Objective

• When Q(C) < b, the objective is to approximate the constrained optimum; when b ≤ Q(C), the distribution-level optimum is Q, but model learning can still violate the constraint.At b = 1 with Q(C) < 1, every feasible distribution must place zero mass outside C, while the stated objective becomes infinite.Source: §2 · p. 2, §2.1 · p. 3
• The learning objective uses membership-oracle access to C, without differentiable distances, projections, or structural properties such as convexity.In the regime b ≤ Q(C), the model is trained to match Q while meeting the specified constraint rate.Source: §2.1 · p. 3

3 Constraint Guarantees with FM Parameterization

• The FM model represents distributions through a parameterized velocity field and its induced flow map.The parameterized Lagrangian and dual function define the corresponding constrained optimization over velocity fields.Source: §3 · pp. 3–4
• The analysis assumes regular target paths and sufficiently smooth, bounded true and parameterized velocity fields, with parameterization error bounded by ν.It also assumes bounded optimal dual variables and positivity conditions on the relevant densities.Source: §3 · p. 4
• When ν is sufficiently small, Theorem 3.5 bounds the FM solution’s total-variation distance from the distribution-level optimum and its constraint-rate error.The bound separates FM approximation error, A1ν + A2ν^2, from dual mismatch controlled through the strong-concavity constant k; the rate error vanishes as ν decreases.Source: §3 · p. 4

4.1 Theoretical PDFM Algorithm

• The theoretical algorithm minimizes the parameterized Lagrangian over velocity fields and updates the dual variable using the constraint-satisfaction gap.With exact primal minimization and dual supergradient updates, the dual variable converges to an optimum; further conditions yield convergence of the primal distributions.Source: §4.1 · pp. 4–5

4.2 Practical Implementation of PDFM

• The practical objective replaces likelihood minimization with conditional flow matching and uses independent random couplings with linear interpolants.The conditional-FM objective is motivated by tractability and a result relating bounded velocity error to bounded KL divergence under regularity assumptions.Source: §4.2 · p. 5
• A learned predictor trained on terminal membership labels supplies differentiable constraint-satisfaction estimates, addressing the indicator-gradient obstacle and avoiding Gaussian velocity randomization.The paper's strategy avoids the Gaussian randomization of the velocity field proposed in Huan et al. [2026].Source: §4.2 · pp. 5–6

4.3 Constraint Guarantees on PDFM

• The asymptotic generated distribution remains within an explicit TV neighborhood of the distribution-level optimum, with constraint satisfaction at least b minus that error.The radius depends on optimization error ε and FM parameterization error ν; the theorem assumes boundedness, approximation, and exact dual-gradient conditions.Source: §4.3 · p. 7

5.1 Illustrative 2-D Constrained Generation (Q(C) < b)

• 99.51% constraint satisfaction and 0.1096 SWD: PDFM meets the 99.5% target and has the lowest SWD in the 2-D example.Vanilla FM achieves 96.37% satisfaction and 0.1384 SWD; FM trained on truncated Q′ achieves 98.74% and 0.1156.Source: §5.1 · p. 8

5.2 Constrained Generation with ℓ2-Ball Constraints (Q(C) ≥b)

• In the 8-D and 20-D tasks, PDFM reaches constraint satisfaction close to the target while achieving smaller SWD than FM-RE.PDFM uses membership-oracle feedback and directly controls P(C) through the target b; FM-RE influences it indirectly through a constraint-term weight.Source: §5.2 · p. 8, Table 1 · p. 9

5.3 Unconditional PDE-Constrained Generation of Spatiotemporal Dynamics

• In the PDE task, PDFM achieves 97.74% satisfaction in Case 1 and 90.48% in Case 2, while the respective vanilla-FM rates are 14.19% and 9.82%.For Case 2, FM on truncated Q′ reaches 33.24%, still below the 90% target; rejection removes approximately 60% of training data and increases histogram error.Source: §5.3 · pp. 9–10, Table 2 · p. 9

5.4 Additional Applications

• PDFM reaches 99.28% fingerprint constraint satisfaction against a 99% target, but molecule stability improves only to 48.75%, below its 60% target.For fingerprints, FM reaches 72.74% and FM-RE 88.14%; for molecules, PDFM also improves NLL from −19.74 to −20.70 relative to FM.Source: §5.4 · p. 10

6 Conclusion

• PDFM requires only membership-oracle feedback, directly controls constraint satisfaction rate, and admits theoretical guarantees.The conclusion also proposes a non-stochastic constraint-guided velocity-field update to use oracle feedback during primal optimization.Source: §6 · p. 10
• PDFM's limitation is the need for sufficient positive constraint-satisfaction feedback during training to learn fϕ reliably.Source: §6 · p. 10

A Related work

• The paper concerns sample-wise constraints and is a training-time method that uses standard Euler sampling with the learned velocity field.Training-time methods incorporate constraint information during training; inference-time methods instead modify sampling and typically add inference computation.Source: A Related work · p. 14

A.1 Training-Time Methods

• Strict-feasibility training methods can require convexity, explicit boundaries, or a closed-form bijection for the constraint set.Reflection methods require convexity and boundary knowledge; MDM requires convexity and a closed-form bijective mapping.Source: §A.1 · p. 14
• Among methods that relax strict feasibility, NAMM needs differentiable distances, while FM-RE uses membership-oracle feedback but lacks direct control of P(C) and theoretical feasibility guarantees.PDFM retains membership-oracle access while providing direct control over P(C) and theoretical guarantees.Source: §A.1 · p. 14

A.2 Inference-Time Methods

• Inference-time methods can avoid retraining but typically require guidance, projection, or correction and add computational burden during sampling.Examples include PDM, which requires projection, and PCFM, which requires a differentiable residual and its Jacobian.Source: §A.2 · pp. 14–15

B Notations

• Appendix B provides a table of notation definitions and descriptions for reference.Source: Table 3 · p. 15

C Proofs for Lemma 2.2 and Lemma 3.1

• Strong duality follows from convexity and Slater’s condition, and strict convexity of the objective yields a unique optimizer.The optimizer is unique on the support where q(x) > 0.Source: C Proofs for Lemma 2.2 and Lemma 3.1 · pp. 15–16
• A total-variation distance δ between two distributions bounds their constraint probabilities within δ.The resulting interval is max{0, P2(C) − δ} ≤ P1(C) ≤ min{1, P2(C) + δ}.Source: C Proofs for Lemma 2.2 and Lemma 3.1 · p. 16

D Proofs for Constraint guarantees with FM Parameterization

• The FM-parameterized solution error is decomposed into velocity-family approximation error and distributional error from dual-variable mismatch.The approximation bound uses regularity constants under Assumption 3.2; the dual-variable component is bounded under additional stated assumptions.Source: D Proofs for Constraint guarantees… · p. 16, §D.1 · p. 16, §D.2 · p. 20

D.2 Bound ∥P (¯λ∗) −P (λ∗)∥TV

• The dual-variable mismatch bound uses strong concavity of g and a bound on total-variation distance between distributions at two dual variables.The appendix states that g is strongly concave on [0, λmax] under Assumptions 2.1 and 3.3.Source: §D.2 · pp. 18–20

• The total-variation bound follows by applying equation (31) and Lemmas D.1, D.2, and D.5.Source: §D.3 · p. 21

E Proof for Proposition 4.1

• The proof assumes the velocity-parameter gradient exists and is uniformly bounded, then expands derivatives across flow steps and exchanges integration with differentiation by dominated convergence.Source: E Proof for Proposition 4.1 · p. 21

F Proofs for Constraint guarantees on PDFM

• The PDFM gap is decomposed into FM-class and minimization approximation error, plus distributional error from mismatch between the current dual variable and its optimum.Source: F Proofs for Constraint guarantees on… · p. 22

F.1 Bound

• Lemma F.1 bounds the stated gap using the approximation error and terms involving μ(λ), ν, and constants A1 and A2.Source: §F.1 · p. 22

F.2 Bound ∥P (λ(i)) −P (λ∗)∥TV

• Under the specified step-size conditions, the dual-iteration analysis bounds the distance to the optimal dual variable as iterations increase.Source: §F.2 · pp. 23–24

• The proof combines Lemma F.1 and Lemma D.2 to obtain the stated constraint-satisfaction bound.Source: §F.3 · p. 24

G Convergence under Unbiased Constraint Estimation

• The analysis replaces exact constraint-gap estimates with conditionally unbiased Monte Carlo estimates and states a convergence guarantee under its listed assumptions.Source: G Convergence under Unbiased… · pp. 24–25

H Experiment Details

• The experiments use standard flow matching with independent random coupling and linear interpolants, with the learning rate decaying from 10^-4 to 10^-5.Source: H Experiment Details · p. 26

H.1 Illustrative 2-D Constrained Generation (Q(C) < b)

• The 2-D illustration compares histograms for the clean distribution, training distribution, constrained optimum, two FM baselines, and PDFM.Source: §H.1 · p. 26, Figure 1 · p. 27

H.2 Constrained Generation with ℓ2-Ball Constraints (Q(C) ≥b)

• For ℓ2-ball tasks in dimensions 8 and 20, the velocity field and feasibility predictor use separate parameterized networks.Source: §H.2 · p. 27

H.3 Unconditional PDE-Constrained Generation of Spatiotemporal Dynamics

• For the Kuramoto–Sivashinsky task, feasibility is based on residuals between generated next states and numerical-stepper predictions; PDFM's objective pairs attractor matching with local dynamical consistency.Source: §H.3 · p. 28

H.4 SOCOFing Fingerprint Generation with Ridge-Connectivity Constraint

• PDFM reaches the 99% target constraint satisfaction rate while maintaining NLLs comparable to FM; FM-RE improves feasibility over FM but has a degraded induced distribution.Source: Table 6 · p. 29, §H.4 · p. 29
• The fingerprint task uses a non-differentiable oracle, with feasibility defined by fewer than 50 connected ridge components; the training-set feasibility rate is approximately 67.84%.Source: §H.4 · p. 29

H.5 QM9 3-D Molecule Generation

• On QM9, 95.2% of samples are stable, and feasibility is determined by valid valency rules after bond inference using membership-oracle feedback.The constraint has no differentiable distance to the feasible set or available projection; the study tests improvement without molecule-specific designs.Source: §H.5 · p. 30
• 48.75 ± 1.56% P(C) versus FM’s 32.77 ± 1.22%, and NLL −20.70 ± 0.25 versus −19.74 ± 0.44; PDFM remains below the 60% target.The authors attribute the shortfall to design limitations, including a generic backbone and probability path without domain-specific knowledge.Source: §H.5 · p. 30

I.1 Why Flow Matching?

• The procedure can extend to DDPM and score-based models, but DDPM’s fixed stochasticity and binary feedback make optimization difficult, with distribution degradation reported for more complex constraints.The authors contrast this with FM’s more stable trajectories, which they say adapt to more complex constraints while preserving distributional fidelity.Source: §I.1 · pp. 30–31

I.2 Implementation Tips and Relationship with FM-RE

• For θ-updates, fϕ avoids full-trajectory sampling; intermediate states may use linear interpolants when the path stays close, and single-iteration training is about 2–3× faster than FM-RE.Training fϕ still requires complete trajectories, while the θ-update can use its fixed predictions; multiple θ-updates between ϕ-updates are suggested only if constraint behavior does not change substantially.Source: §I.2 · p. 31

I.3 Limitations

• PDFM needs both positive and negative feedback; overwhelmingly negative feedback can prevent reliable classifier boundaries, so the initial velocity field must generate some feasible samples.This requirement applies at the start of iteration i = 1 in Algorithm 2.Source: §I.3 · p. 31
• When the feasible region is not sufficiently represented or reachable, especially when Q deviates significantly from C, PDFM may not learn a near-optimal constrained distribution.In that setting, the flow-matching objective can be severely misaligned with constraint satisfaction, producing a large optimization error in Assumption 4.3.Source: §I.3 · p. 31

2. Limitations

• The authors mark LLM usage N/A because LLMs do not contribute as important, original, or non-standard components of the core method, and refer readers to the NeurIPS handbook policy.Source: §2 · p. 42