Primal-Dual Flow Matching for Sample-Wise Constrained Generation

summary

Video file (mp4)

The gist

The gist The Primal-Dual Flow Matching (PDFM) provides a theoretically grounded approach for learning to generate data while satisfying a prescribed sample-wise constraint with provable guarantees on

In short

Primal-Dual Flow Matching (PDFM) is a method for learning data generation while enforcing specific sample-wise constraints. It uses a flow-matching parameterization to derive theoretical guarantees, requiring only simple binary membership feedback for the constraint set. This allows direct control over how well generated samples satisfy the constraint, offering superior performance compared to existing methods.

Key concepts

Flow Matching (FM)
Flow matching is a way to parameterize the process of transforming noise into data. It allows researchers to define a continuous path or 'flow' that guides the generation process from a simple starting point (like random noise) toward the desired complex data distribution. This parameterization is key to deriving guarantees for how well the generated samples adhere to constraints.
Primal-Dual Optimization
This refers to an optimization technique where two related problems—a primal problem and a dual problem—are solved simultaneously. In PDFM, the primal step minimizes a loss function, while the dual step adjusts a variable (the dual variable $\lambda$) to maximize another objective. This coupled approach is used to find an optimal solution that satisfies both the data modeling goal and the sample constraint.
Membership Oracle Feedback
This is a simple type of feedback mechanism where the algorithm only needs to know if a specific point or sample belongs inside or outside a defined constraint set. Unlike complex methods, this feedback does not require calculating distances or projections; it only asks for a binary yes/no answer regarding feasibility, making the method highly flexible for various constraints.

Terminology used across episodes

This episode discusses

The paper

Primal-Dual Flow Matching for Sample-Wise Constrained Generation · Read on arXiv

Zhengyan Huan, Peter Y. Lu, Shuchin Aeron

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

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.

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom.

Jane: And I'm Jane, here with Lu, Meng and our in-house large language model, Lalam.

Tom: Today's paper: "Primal-Dual Flow Matching for Sample-Wise Constrained Generation".

Jane: The gist The Primal-Dual Flow Matching (PDFM) provides a theoretically grounded approach for learning to generate data while satisfying a prescribed sample-wise constraint with provable guarantees on the probability of…

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So we're looking at this paper called Primal-Dual Flow Matching for Sample-Wise Constrained Generation. It’s by Zhengyan Huan, Peter Y. Lu, and Shuchin Aeron at Tufts University Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: It sounds complicated because it deals with learning to generate data while making sure the samples meet a specific constraint probability we choose ourselves for them, and this paper introduces a way to get provable guarantees on how likely those samples are to satisfy the constraint Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: Exactly what it does is start with a population-level problem that minimizes the negative log-likelihood subject to some inequality involving the distribution P and the constraint set C Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Lu: The core idea they use is flow matching parameterization to derive these guarantees on constraint satisfaction probability for a corresponding primal-dual optimization algorithm Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Jane: What's really interesting here is that they propose something called Primal-Dual Flow Matching PDFM which has direct control over the constraint satisfaction rate and only needs binary membership oracle feedback for constraint satisfaction Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Tom: That sounds like a big deal because it avoids needing differentiable distances or projections or any special structure of the constraint set like convexity Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Jane: It seems they are trying to make this approach more accessible for real-world constraints where you might not have that nice smooth structure Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Tom: Right, and the theoretical framework involves defining an optimal population-level problem OPT that minimizes P Ex∼Q

− log p(x): subject to b − Z C p(x)dx ≤ zero Zhengyan <ref:2610.12699#pg3,C p(x)dx ≤ 0>.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: And under Assumption two point one they show that the optimal population-level solution corresponds to re-weighting the density of Q inside and outside of C Zhengyan <ref:2610.12699#pg1>.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg3

Tom: Now let's talk about their practical side which is PDFM Algorithm one which is an idealized primal-dual procedure where the primal step minimizes the parameterized Lagrangian L¯(uθ, λ) over UΘ for a current dual variable lambda Zhengyan <ref:2610.12699#pg1>.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Jane: And the dual update is a projected supergradient ascent method for maximizing a concave function g¯ Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: The practical implementation they show in PDFM Algorithm two involves alternating updates for the velocity field uθ and the dual variable λ Zhengyan <ref:2610.12699#pg1>.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Jane: A key innovation there is building differentiable surrogates for constraint satisfaction probability under the FM parameterization so it can work with only membership oracle access to the constraint set Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: So if I'm thinking about what this means for me as an engineer on a team shipping models, this suggests we can handle constraints without needing to know the exact distance metric between points Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: That's right because they show that Theorem three point five demonstrates that under Assumptions two point one, three point two, three point three and three point four the FM-parameterized constrained solution induces a distribution P(λ¯∗)uθ that remains close in total variation to the distribution-level optimum P(λ∗) Zhengyan <ref:2610.12699#pg2>.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: And they give us this bound Epara which separates the two main sources of error: the FM approximation error controlled by A1ν + A2ν squared and the dual mismatch between λ¯∗ and λ* Zhengyan <ref:2610.12699#pg1>.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Jane: The proof shows that P(λ¯∗)uθ satisfies the target constraint rate b up to an explicit error term that vanishes as ν decreases Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: Now for the experimental results they show on two-D constrained generation PDFM achieves a constraint satisfaction rate of ninety-nine point five one percent and also the lowest Sliced Wasserstein Distance SWD of zero point one zero nine six Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Jane: In higher-dimensional tasks they say PDFM achieves P(C) close to the target b while yielding a smaller SWD than FM-RE Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: They also show in the context of PDE-constrained spatiotemporal dynamics that PDFM can produce samples of higher dynamical quality than the data it was trained on by leveraging a solver-based residual oracle to suppress non-physical regions of the learned distribution Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Jane: And in three-D molecule generation they improve mean molecule stability from thirty-two point seven seven percent to forty-eight point seven five percent without relying on task-specific designs Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: So what does this all mean for the world right now? It means we have a primal-dual framework for constrained generation that requires only membership oracle feedback provides direct control over the constraint satisfaction rate and admits theoretical guarantees Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: It's competitive against constraint-aware baselines while maintaining distributional fidelity Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: The limitation they point out is that they need to obtain sufficient positive constraint satisfaction feedback during training to learn the flow fϕ reliably Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: And also PDFM requires the feasible region to be sufficiently represented or reachable under the data and model distributions Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Tom: Well that’s the picture for now on Primal-Dual Flow Matching for Sample-Wise Constrained Generation Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: We've covered the title, the mechanics of PDFM, and what they found in 2D and three dee generation Zhengyan <ref:2610.12699#pg2>.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: It’s a method that seems to simplify how we handle constraints by using membership checks instead of relying on complex distances Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: That’s the core idea here: direct control over the constraint satisfaction rate without needing differentiable projections Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: We'll take a quick break and then we'll look at how this method compares to other ways of doing constrained generation Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

The paper's summary: Tom: So, to recap this paper, they’re taking something called flow matching and turning it into a primal-dual framework for generating data that has constraints we specify on each individual sample Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: Right, so instead of just making the whole dataset fit some average constraint, they’re getting it right down to the level of every single sample's probability of meeting that constraint Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: Exactly. The big deal is they get these provable guarantees on how likely those samples are to satisfy the constraint, and they don't need any fancy math structures for the constraint set itself Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: What that means for us is that we can set a target rate of satisfaction—say, ninety-nine point five percent—and the paper gives us a mathematical way to know if our generation model is actually hitting that mark Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: And they do this by using something called binary membership oracle feedback Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: That's the key part for implementation, because it means we don't need to calculate some complex distance between points to check if a sample is inside or outside the constraint region Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: Right, so we just need an oracle that tells us yes or no whether a sample is in the set C Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: And the paper shows that this approach works well with flow matching parameterization, creating these surrogates to handle that binary feedback without needing smooth differentiable distances Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: The theoretical guarantee they’re showing is that the generated distribution P(λ¯∗)uθ stays close to the ideal distribution P(λ*) Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: They give a bound for this error Epara which tells us how far off we are from the perfect solution based on two main sources: the approximation error from flow matching and the mismatch between our dual variable and the optimal one Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: And they show that this error term goes to zero as some parameter ν decreases Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: It's pretty strong because it shows we can directly control the constraint satisfaction rate b, which is something methods like FM-RE only do indirectly by adjusting a penalty weight Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: And experimentally they’re getting solid numbers on two-D generation, hitting ninety-nine point five one percent constraint satisfaction and the lowest sliced Wasserstein distance of zero point one zero nine six Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: And in higher dimensions, it still beats FM-RE for the sliced Wasserstein distance while keeping the constraint level close to what we want Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: It's also showing some cool things in physics applications too, like spatiotemporal dynamics where they use a solver-based residual oracle to clean up the samples and get higher dynamical quality Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: So basically, this method is proving it can handle complex constraints without needing those complicated structures we usually have to build into our models Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: The main thing to watch is that they need a decent amount of positive constraint satisfaction feedback during training for the flow function fϕ to learn reliably Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

The paper's improvements: Tom: So, moving on from what they just described about the core method, I want to talk about what they suggest as improvements to make PDFM even better Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: Okay so if we look at the math behind it, what’s the next step they propose for refinement Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: They suggest that they focus on making sure the positive constraint satisfaction feedback is robust enough during training so that the flow function fϕ learns reliably Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: That sounds like a practical hurdle because in real training runs you don't always get perfect signals on every sample Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: And they also flag that the feasible region itself needs to be sufficiently represented or reachable by the data and the model distribution Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: So it’s not just about the algorithm itself but making sure your setup—the data and how you represent things—is actually good enough for this to work well Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: The implication there is that if we can get that reliable feedback loop going then the theoretical bounds on the error term Epara become even tighter Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: And those tighter bounds mean we can expect even more precise control over the final distribution P(λ¯∗)uθ Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: It’s about pushing the limits of how accurately we can model these constraints using only binary membership checks Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: And if they nail that then PDFM becomes a really flexible tool for any generation task with sample-wise rules without needing those complicated differentiable structures Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: It means the practical barrier isn't just the math but getting that initial training stability down Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: And it changes how we think about constrained generation by showing that you can achieve a specific constraint satisfaction level with high confidence even when you only have membership oracle access Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: So the path forward seems to be getting that stable training signal and then applying the framework to more complex constraint types Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Conclusion: Tom: So to wrap up, this paper is about Primal-Dual Flow Matching for Sample-Wise Constrained Generation Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: It basically shows that we can use flow matching to generate data where we directly control the probability that each individual sample meets a specific constraint Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: That direct control without needing complex differentiable distances is the big win here Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: It really shifts the focus from just fitting a general distribution to explicitly managing sample-wise properties which is super useful for quality control in generation Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Lu: I think the theoretical part is solid because they’re linking the optimal population problem directly to re-weighting the density inside and outside of C Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg3

Meng: From an engineering standpoint, if we can control the constraint rate like this it means we don't have to guess how much penalty weight to use in a loss function Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Lalam: I’m confident this is a solid direction because it tackles the sample-wise constraint problem with provable guarantees using established flow matching techniques Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: And experimentally they showed they hit ninety-nine point five one percent for two-D generation which is pretty close to what we need Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: It’s also showing up in spatiotemporal dynamics where they use a solver oracle to suppress non-physical parts of the distribution Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Lu: The caveat is definitely that they need good training feedback for the flow function fϕ to work well and the constraint set needs to be reachable by the model’s data Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Meng: If we can get that stable training signal then this could actually make deployment easier because we’re not wrestling with implicit penalty tuning Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Lalam: I’m pretty sure this is the most direct way to achieve what people want when they ask for sample-wise constraints using flow matching Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: So Primal-Dual Flow Matching for Sample-Wise Constrained Generation is a framework that requires only binary feedback to get provable constraint satisfaction guarantees Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

Jane: It's a really neat way to put the power of flow matching behind hard constraints without needing all those complex projection steps Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg2

Tom: Next up we’re looking at how this compares to other constrained generation methods and whether the hardware cost is reasonable Zhengyan.Huan,Peter.Lu,Shuchin.Aeron @tufts.edu #pg1

More episodes

← Home