Primal-Dual Flow Matching for Sample-Wise Constrained Generation

AI summary

Overview

Research area: Generative modeling — specifically constraint-aware generation with flow-matching (FM) models, combined with primal-dual optimization from convex analysis.

Technical level: Advanced. The paper is theory-forward: it builds on Lagrangian duality, total variation (TV) bounds, KL divergence control of velocity-field approximation error, and stochastic approximation theory for projected supergradient ascent.

Scope in one sentence: The paper derives population-level and finite-approximation guarantees for a primal-dual flow-matching method that trains a generative model to hit a user-specified sample-wise constraint satisfaction rate using only binary membership-oracle feedback.

What This Paper Is About

Generative models used in science need to produce samples that are not merely realistic but also valid — satisfying a boundary condition, a conservation law, a governing PDE, or a chemical property. The paper studies how to train a flow-matching model so that the probability its samples land inside a designated constraint set reaches a user-chosen target rate b ∈ (0,1), while keeping the generated distribution faithful to the data. Crucially, the only access to the constraint set is a membership oracle: a function that answers "yes/no, is this point in the constraint set?" — no differentiable distance, no projection operator, and no structural assumptions such as convexity are required.

Key Contributions

1. A constrained population-level formulation with a closed-form optimum. Starting from a constrained negative log-likelihood minimization (equivalent to minimizing KL(Q||P) subject to b − ∫_C p(x)dx ≤ 0), the authors show via Lemma 2.2 that under Assumption 2.1 the optimal density re-weights the data density q inside versus outside the constraint set, and give the exact optimal dual variable λ* = (1−Q(C))/(1−b) − Q(C)/b.

2. Constraint guarantees under flow-matching parameterization (Theorem 3.5). The authors bound the total variation distance between the FM-parameterized constrained solution P_{u_θ}^{(λ̄*)} and the distribution-level optimum P^{(λ*)} by an explicit error term E_para, and show this implies P_{u_θ}^{(λ̄*)}(C) ≥ b − E_para.

3. A theoretical primal-dual algorithm and a practical implementation. Algorithm 1 is the idealized primal-dual template (exact primal minimization, projected supergradient ascent on the dual). Algorithm 2, PDFM, replaces the intractable likelihood objective with the tractable conditional flow-matching (CFM) objective and builds a differentiable neural surrogate for constraint satisfaction probability, inspired by deterministic policy gradient methods.

4. An asymptotic guarantee for the practical algorithm (Theorem 4.5). Under boundedness, ε-suboptimality, and exact dual-gradient assumptions, the practical algorithm's iterates stay within an explicit TV neighborhood of the distribution-level optimum and satisfy the constraint at rate at least b − E_PDFM asymptotically.

5. Empirical evaluation reported in the paper's abstract. PDFM is claimed to improve constraint satisfaction over vanilla FM, achieve competitive performance against constraint-aware baselines while maintaining distributional fidelity, and to improve molecule stability and distributional fidelity over FM without task-specific designs. The paper also states code is available at https://github.com/ZhengyanHuan/PDFM.

Main Findings

• Closed-form constrained optimum: Lemma 2.2 gives p^{(λ*)}(x) = (b q(x)/Q(C)) 1_C(x) + ((1−b)q(x)/(1−Q(C))) 1_{C^c}(x). The optimum is a re-weighting of the data density inside and outside the constraint set.

• Two problem regimes: When 0 < Q(C) < b < 1, the goal is to approximate the constrained optimum P^{(λ*)}; this regime is studied both theoretically and empirically. When 0 < b ≤ Q(C) ≤ 1, the distribution-level optimum is Q itself, so the goal becomes matching Q while still achieving the target constraint rate; this regime is studied empirically only.

• The b = 1 boundary is degenerate: If Q(C) < 1, every feasible P must satisfy P(C) = 1 and P(C^c) = 0, while Q(C^c) > 0 yields an infinite objective value.

• Rejection sampling is not a substitute: Training on the truncated distribution Q' = Q(·|C), with density q'(x) = q(x)1_C(x)/Q(C), can substantially reduce the effective sample size and degrade the training distribution, especially with noisy data. Even with enough truncated samples, finite-capacity models may still miss the target constraint rate — the paper states this is shown in Sections 5.1 and 5.3.

• TV converts into constraint-rate control: Lemma 3.1 states that if ||P1 − P2||_TV = δ, then max{0, P2(C) − δ} ≤ P1(C) ≤ min{1, P2(C) + δ}.

• Theorem 3.5 decomposition: The bound E_para = 1/(((1−Q(C))^{1/3} + Q(C)^{1/3})^3) · sqrt((2/k) μ(λ̄*)(A1 ν + A2 ν²)) + sqrt((1/2)(A1 ν + A2 ν²)) separates the flow-matching parameterization error (controlled by A1 ν + A2 ν², where ν bounds the mean-squared velocity error) from the dual mismatch between λ̄* and λ*, which is governed by the strong concavity constant k of the dual function. Here μ(λ) = (1/2)(λ + 1 + sqrt((λ+1)² − 4λ(1 − Q(C)))).

• Dual updates are well-behaved: The dual objective ḡ is concave as the pointwise infimum of affine functions, and the supergradients ∇_λ L̄(u_θ^{(λ)}, λ) = b − P_{u_θ}^{(λ)}(C) are uniformly bounded in [b−1, b], so Algorithm 1 is projected supergradient ascent with standard step-size conditions (η_i > 0, Ση_i = ∞, Ση_i² < ∞).

• Theorem 4.5 error term: With e(λ) := sqrt((ε + μ(λ)(A1 ν + A2 ν²)) / (2(μ(λ) − λ))) and e_max = max_{λ∈[0,λ_max]} e(λ), the practical algorithm obeys E_PDFM = e_max + 2e_max / (k((1−Q(C))^{1/3} + Q(C)^{1/3})^3), subject to 0 ≤ η_i ≤ 2/(kw). The neighborhood radius shrinks with the optimization error ε and the parameterization error.

• Exactness assumption and its relaxation: The guarantee assumes the empirical estimator of b − P_{u_θ'}(C) is exact at every dual update (Assumption 4.4). If relaxed to unbiased Monte Carlo estimation, the guarantee holds in expectation over liminf, bounding liminf_{i→∞} E[P_{u_θ'}^{(λ(i))}(C)] rather than the pathwise quantity; these statements and proofs are in Appendix G.

• Empirical results as reported: The paper claims PDFM improves constraint satisfaction over vanilla FM and is competitive with constraint-aware baselines at comparable distributional fidelity, and for chaotic spatiotemporal PDE dynamics achieves superior performance while needing only binary feedback. Specific numeric results — benchmark scores, dataset sizes, or model names — are not reported in the available paper content (the experiments section is truncated).

Methodology in Plain English

The authors begin from a standard probabilistic goal: learn a distribution P that stays close to the data distribution Q but is required to place at least probability b of its mass inside a constraint set C. Writing this as a constrained optimization problem and introducing a Lagrange multiplier λ turns it into a min-max problem: for a fixed λ, minimize a penalized objective; then adjust λ upward when the constrained mass is too low and downward when it is high enough.

Because actual distributions are represented by flow-matching velocity fields rather than densities, the theory has to account for the fact that a neural network can only approximate the ideal velocity field. The authors formalize this with a parameterization error ν and show that if ν is small enough, the generated distribution stays close in total variation to the ideal constrained distribution, and the constraint is satisfied up to a matching error term.

For the practical algorithm, two obstacles appear. First, the exact negative log-likelihood objective is replaced with the conditional flow-matching regression objective, which is tractable and has an existing result linking bounded velocity error to bounded KL divergence. Second, differentiating through the indicator function 1_C is impossible — its gradient is zero almost everywhere and undefined on the constraint boundary. The authors' fix is to train a separate neural network f_φ that predicts, from an intermediate point (x_t, t) along a discretized Euler trajectory, whether the trajectory will end inside C. It is trained by binary cross-entropy against the terminal label y = 1_C(x_{NΔt}), using an N-step forward Euler discretization with step Δt = 1/N. Because f_φ is differentiable, its gradient can be back-propagated into the velocity field, giving an estimable gradient for the constraint satisfaction probability (Proposition 4.1, with the practical gradient formulas in equations 17 and 18). This design deliberately avoids the Gaussian randomization of the velocity field used in prior work by Huan et al. (2026); the authors describe the intuition as coming from deterministic policy gradient methods.

In practice the algorithm alternates: update θ (the velocity field) and φ (the predictor) to convergence, then update the dual variable λ' by λ'(i+1) = [λ'(i) + η_i (b − P_{u_θ}(C))]_+, where P_{u_θ}(C) is estimated by Monte Carlo — generating terminal samples and computing the fraction that land in C.

Why This Matters

Impact on research. The paper connects constrained generative modeling to classical convex duality with explicit, non-asymptotic-style error decompositions, showing that the constraint satisfaction rate can be controlled — not merely hoped for — as a function of approximation error. It also expands the class of usable constraints: because only a membership oracle is needed, the method applies where differentiable distances, projections, and convexity are unavailable, a regime previous constraint-aware training methods could not address directly. The work sits within a literature the paper organizes along two axes: type of constraint (distributional constraints such as Khalafi et al., 2024, versus sample-wise constraints) and intervention point (inference-time methods such as Christopher et al., 2024, Utkarsh et al., 2025, Li et al., 2026, Zhu et al., 2026, versus training-time methods such as Lou and Ermon, 2023, Fishman et al., 2023, Xie et al., 2024, Feng et al., 2025, Giannone et al., 2023, and Huan et al., 2026). PDFM is positioned as a training-time, membership-oracle-only method.

Real-world applications (as motivated by the paper):

• Chaotic spatiotemporal system emulation. Flow- and diffusion-based models are increasingly used for chaotic dynamics, where high sensitivity to initial conditions makes trajectories probabilistically modeled (Medio and Lines, 2001; Lienen et al., 2024; Du et al., 2024; Price et al., 2025). Distributional fidelity alone is insufficient — samples must satisfy boundary conditions, conservation laws, or a governing PDE (Beucler et al., 2021; Utkarsh et al., 2025).

• Molecular generation. Generating new molecules that conform to structures with chemically useful, numerically quantifiable properties (Bickerton et al., 2012; Giaginis et al., 2018; Hoogeboom et al., 2022; Song et al., 2023). The paper reports that PDFM improves molecule stability and distributional fidelity over FM without task-specific designs.

• Scientific machine learning generally. Any setting where validity is checked numerically rather than expressed analytically — the paper frames constraint specification via membership oracle as "the most flexible" form of access when constraints are complex.

Industry relevance. A method that requires only a yes/no validity checker is attractive because in industrial pipelines the validity check is often a simulator, a rules engine, or an expensive numerical PDE solver, not a differentiable function. A single tunable knob b that directly trades off constraint satisfaction against distributional fidelity is also operationally convenient, and the requirement of a known constraint boundary or differentiable distance is dropped.

Future Directions

• Relaxing the exact dual-gradient assumption. The paper's Theorem 4.5 assumes the empirical estimator of b − P_{u_θ'}(C) is exact. The authors note that with unbiased Monte Carlo estimation the guarantee must be stated in expectation; tightening and generalizing these expectation-based results is a natural next step.

• Finite-sample and finite-iteration theory. The current guarantees are population-level or asymptotic (i → ∞). Quantifying how the Monte Carlo estimate of P_{u_θ}(C) and finite dual iterations affect the achievable constraint rate would connect the theory more directly to the practical Algorithm 2.

• Choosing the surrogate quality ρ. The bound in equation 14 depends on how well f_φ approximates the indicator, with tolerance ρ. Characterizing how neural approximator quality, discretization step Δt = 1/N, and trajectory length interact with ρ — and hence with the overall error — is left open.

• The practical optimization error ε. Assumption 4.3 postulates ε-suboptimality of the primal solution at each dual variable without specifying how ε arises from the alternating θ/φ updates. Deriving ε from the actual optimization dynamics, or improving it with better surrogate training, is an open question.

Target Audience

This paper is aimed at researchers and graduate students in machine learning, generative modeling, and scientific machine learning who work on constrained or physics-informed generation. It will be most valuable to readers comfortable with convex optimization and duality, flow matching, and total variation / KL divergence arguments — the derivations lean on regularity conditions borrowed from Su et al. (2025) and on supergradient ascent analysis. Practitioners who need constraint-satisfying generative models in scientific domains (PDE emulation, molecular design) but who only have access to a binary validity checker will also find the practical algorithm and its oracle-only requirement directly relevant. Readers looking for extensive benchmark tables should note that the quantitative experimental details are not contained in the available portion of the paper.

Authors’ 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.