Constrained Diffusion Models with Primal-Dual Inference

arXiv:2606.17192 · cs.LG · Submitted 2026-06-15 · Read on arXiv

Listen

Transcript

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

Tom: Today's paper: "Constrained Diffusion Models with Primal-Dual Inference".

Jane: The gist This paper develops constrained diffusion models with primal-dual inference (PDI) to sample from optimal distributions of entropy-regularized optimization problems with average constraints.

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

Paper summary: Tom: So, we're talking about this paper called "Constrained Diffusion Models with Primal-Dual Inference." It looks like they're trying to sample from optimal distributions for optimization problems that have average constraints.

Jane: That sounds really technical, Tom. What's the main idea they're pushing? What does this paper actually claim is so important?

Lu: Basically, they formalize constrained sampling in the Lagrangian dual domain, and they show that the optimal distribution ends up being a Gibbs distribution indexed by an optimal dual variable. This is a big theoretical connection to make.

Meng: So, instead of just estimating that dual multiplier before you start sampling and freezing it, this paper suggests you figure out the primal distribution and its dual variable together as you go.

Lalam: It's like they are interleaving two things: the reverse diffusion process, which is denoising the samples, with dual ascent, which is updating that multiplier based on how much constraints are being violated.

Tom: Exactly. They propose primal-dual inference or PDI where at each step you use a score field related to the current multiplier to update your samples, and then you update that multiplier itself using estimates of constraint violation from those denoised samples <ref:2606.17192#pg2>.

Jane: And for training, they train just one dual-conditioned score network that can handle the whole family of Gibbs distributions indexed by those dual variables during inference. That sounds like a lot to learn at once.

Lu: The theoretical findings show that the optimal distribution is this Gibbs distribution parametrized by the optimal dual variable, which they also view as an exponential tilting of the implicit uniform prior on X <ref:2606.17192#pg2>. That connection comes from Proposition one showing that the optimal pair, mu star and lambda star, is a saddle-point of the Lagrangian <ref:2606.17192#pg3>.

Meng: And they prove that the time-average of those dual variables generated along the inference trajectory actually converges to a neighborhood around the dual optimum at a rate of one over square root T <ref:2606.17192#pg3>. That convergence is important for stability.

Tom: The results also show that the Wasserstein distance between the terminal distribution you get from PDI and the actual optimal distribution is controlled by that dual mismatch term <ref:2606.17192#pg3>. Specifically, it's bounded by E W2 pT, mu star <= T X−one t=zero Ψt,T bT −t √aT −t γtE∥λt − λ∗∥ + ϵapp(t) <ref:2606.17192#pg1>.

Jane: That formula looks intense. So, what does this mean practically for someone just listening to the show? What's the big picture here?

Paper summary: Lu: The empirical validation shows that PDI improves lower-rate users by updating those dual variables during denoising, steering samples toward feasibility <ref:2606.17192#pg3>. It also achieves the best overall balance between mean rates and feasibility, represented by the 5th percentile rate and mean and sum violations in problems like wireless power allocation <ref:2606.17192#pg3>.

Meng: It seems to be a real trade-off here. PDI gives you near-feasible performance while providing the best balance between return, feasibility, and diversification among all diffusion samplers <ref:2606.17192#pg3>. It's a better compromise than just aiming for the absolute highest rate or lowest violation alone.

Tom: They also found that PDI reduces this limitation by learning a lambda-conditioned family of score fields, which makes it more robust to shifts in parameters like rmin <ref:2606.17192#pg3>. It's less sensitive to those kinds of problem specifics.

Jane: So, compared to the unconstrained model—which completely ignores the constraints and achieves the lowest objective—PDI is clearly doing something different by incorporating that dual variable information into the sampling process <ref:2606.17192#pg3>.

Lu: The comparison with baselines shows that PDI exhibits the best objective and mode diversity compared to other methods like PDL or PDM, while still maintaining full feasibility <ref:2606.17192#pg3>. Even against unconstrained models, it maintains high feasibility and very small mean violations while achieving the highest return among all diffusion samplers <ref:2606.17192#pg3>.

Meng: That's practical stuff for deployment. If you need a solution that is both good and feasible, PDI seems to hit that sweet spot in terms of return and constraint satisfaction <ref:2606.17192#pg3>. But they do flag a limitation, though.

Jane: What is the limitation they explicitly state about this method? Where does it stop working right now?

Lu: They say the convergence guarantees for their algorithm apply to time-averaged dual iterates rather than final iterates or the primal distributions directly <ref:2606.17192#pg3>. It's about averages, not necessarily the final state of things.

Tom: And they also mention that PDI requires problem-specific tuning of the dual-variable prior used during training and the entropy regularization strength beta, for which principled selection rules are still open questions <ref:2606.17192#pg3>. That means you have to tweak things depending on your exact problem setup.

Jane: It sounds like it's a very powerful tool, but it still has those tuning knobs that need careful handling by the person using it. So, what's the final word on this paper and its title?

Paper summary: Lu: The paper "Constrained Diffusion Models with Primal-Dual Inference" is tackling how to sample from optimal distributions for optimization problems with average constraints by jointly inferring the primal distribution and its dual variable during inference <ref:2606.17192#pg1>. It shifts the focus from estimating the multiplier before sampling to making it part of the sampling state itself <ref:2606.17192#pg2>.

Tom: And what this means for us is that we can get better, more balanced solutions in complex systems where constraints are tricky, without having to freeze our variables too early <ref:2606.17192#pg3>. It’s about making the sampling process smarter by incorporating the dual information as it runs.

Jane: So, for anyone listening who only gets this show, think of it as an AI method that actively corrects its path toward satisfying constraints while trying to get the best possible outcome <ref:2606.17192#pg3>. That's what they’re showing us with constrained diffusion modeling.

Lu: Exactly. It’s about learning a lambda-conditioned family of score fields, which is a significant step in making diffusion models more flexible and better at handling real-world constraints <ref:2606.17192#pg3>.

Meng: From an engineering standpoint, it suggests that we can design diffusion samplers that are inherently more aware of the constraints they are trying to satisfy during the generation process, instead of just treating them as something to check at the end <ref:2606.17192#pg3>.

Jane: So, in summary, this paper is showing how primal-dual inference can lead to better trade-offs between performance and constraint satisfaction in sampling problems <ref:2606.17192#pg3>.

Tom: That's the gist of "Constrained Diffusion Models with Primal-Dual Inference." It’s a method that makes the dual variable an active part of the sampling process to guide it toward feasibility, and it does that by training a single score network that can handle different dual variables <ref:2606.17192#pg1>.

Lu: The connection between the optimal distribution being a Gibbs distribution indexed by the dual variable and an exponential tilting of a uniform prior is what underpins their theoretical structure <ref:2606.17192#pg2>.

Jane: And even with those caveats about training, the convergence guarantees for time-averaged iterates are solid at a rate of one over square root T <ref:2606.17192#pg3>.

Meng: It’s important to remember that they have to tune things like the dual variable prior and the entropy regularization strength beta depending on your specific problem <ref:2606.17192#pg3>.

Tom: So, it's not a plug-and-play solution yet, but it’s a sophisticated way to design constrained diffusion models that actively steer samples toward feasibility during inference <ref:2606.17192#pg3>.

Jane: And this work really shows how important it is to go beyond unconstrained models when real constraints are involved in the optimization problem you're trying to solve <ref:2606.17192#pg3>.

Conclusion: Tom: So we’ve been diving into how this paper uses primal-dual inference to sample from optimization problems with average constraints.

Jane: That's right, Tom, it’s about taking the dual variable—that extra piece of information—and using it actively during the sampling process instead of just checking it at the end.

Lu: The authors are showing that this primal-dual inference method lets us sample from optimal distributions that look like Gibbs distributions parameterized by that dual variable.

Meng: So, what’s the real practical meaning here for someone building something? Is it just theoretical math, or does this actually help engineers?

Lalam: It means we can design diffusion models that are inherently smarter about satisfying those constraints while they’re generating the output.

Tom: Exactly, Lalam. The title is "Constrained Diffusion Models with Primal-Dual Inference," and it points to this core idea of making the dual variable a dynamic part of the sampling itself.

Jane: And if you think about that, it suggests we move past just getting a final answer and start thinking about how the process gets there, moment by moment.

Lu: The paper proves that this dual ascent step interleaved with reverse diffusion actually converges to a region around the true optimum for the dual variables over time.

Meng: Convergence is one thing, but what about tuning this whole system? What are the practical hurdles when you try to implement something that depends on these dual iterates?

Lalam: The authors do admit that they need problem-specific tuning for two things: the prior used during training and how strong the entropy regularization strength beta needs to be.

Tom: Right, so it’s not a magic fix right now; you still have to tune those parameters based on your specific optimization setup.

Jane: It gives us a really solid framework for when we think about these complex sampling problems where feasibility matters as much as the final objective value.

Lu: This paper really pushes the idea that we can learn a family of score fields conditioned on that dual variable, which makes the model more robust to changes in things like those r parameters.

Meng: So, the implication is that we might be able to build diffusion samplers where they are naturally steering samples toward feasibility during generation.

Lalam: It’s about making the sampling process more aware of the constraints it’s trying to satisfy, rather than treating them as an afterthought you check at the end.

Tom: And that leads us nicely into how this approach compares to other methods we’ve seen for constrained sampling.

Jane: We'll be talking next on how this primal-dual inference actually stacks up against those other constraint handling techniques.

Department of Electrical and Systems Engineering, University of Pennsylvania

cs.LG

Submitted: 2026-06-15

Updated: 2026-10-10

Code: https://github.com/SMRhadou/PDI-Diffusion

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 90/100

The gist: The gist This paper develops constrained diffusion models with primal-dual inference (PDI) to sample from optimal distributions of entropy-regularized optimization problems with average constraints.

Key concepts

Constrained Sampling as Saddle-Point Problem: Sampling under constraints is framed as finding a saddle point in the Lagrangian dual domain. The optimal distribution for this problem is shown to be a Gibbs distribution indexed by the optimal dual variable. This mathematical framing allows the constrained sampling process to be integrated into a diffusion framework.
Primal-Dual Inference (PDI): PDI combines two processes: reverse diffusion, which generates samples, and dual ascent, which updates Lagrange multipliers. At each denoising step, PDI uses a score field related to the current multiplier to update samples while simultaneously updating the multiplier based on constraint violations.
Dual Variable-Conditioned Score Network: The model is trained using a single score network that depends on the dual variable. This network generalizes across different Gibbs distributions indexed by various dual variables encountered during inference, allowing the model to adapt its sampling strategy based on the current constraint context.

Terminology

Diffusion Model — 2× in this paper · explained in 10 episodes
Diffusion Models — 2× in this paper · explained in 9 episodes

Summary

The gist This paper develops constrained diffusion models with primal-dual inference (PDI) to sample from optimal distributions of entropy-regularized optimization problems with average constraints.

How it works

Constrained sampling is cast as a saddle-point problem in the Lagrangian dual domain, where the optimal distribution takes the form of a Gibbs distribution indexed by the optimal dual variable.
PDI interleaves dual ascent with reverse diffusion, where at each denoising step, it uses a score field associated with the current multiplier to update samples and then updates the multiplier through dual ascent using constraint violations.
A single dual-variable-conditioned score network is trained that generalizes across the family of Gibbs distributions indexed by the dual variable encountered during inference.
The primal update uses the score field associated with the current multiplier, and the dual step updates multipliers in the direction of average constraint violation estimated from Tweedie-based estimates of denoised samples.

Key Theoretical Findings

The optimal distribution is a Gibbs distribution parametrized by the optimal dual variable, which can equivalently be viewed as an exponential tilting of the implicit uniform prior on X.

This connection arises from Proposition 1, which shows that under Assumption 1, the optimal pair (µ∗,λ∗) is a saddle-point of the Lagrangian.

The time-average of the dual variables generated along the inference trajectory converges to a neighborhood of the dual optimum.

Theorem 1 proves this convergence at a rate of 1/√T to a region around the dual optimum.

The Wasserstein distance between the terminal distribution of PDI and the optimal distribution is controlled by the dual mismatch.

Theorem 2 establishes that E W2 pT, µ∗ ≤ T X−1 t=0 Ψt,T bT −t √aT −t γtE∥λt − λ∗∥ + ϵapp(t), where Ψs,T = QT s=s+1 ρs and γt = RαT−t βσ2 T−t p supλ∥Covπ t (y0x,λ)op.

Empirical Validation

PDI further improves the lower-rate users by updating the dual variables during denoising.

Figure 1 shows that PDI, unlike unconstrained samplers, steers samples toward (average) feasibility and lifts the worst-served users much closer to their targets.

PDI achieves the best overall balance between mean rates (objective) and feasibility represented by the 5th-percentile rate and mean and sum violations.

Table 1 shows that PDI achieves a near-feasible performance and provides the best balance between mean rates and constraint violations in the wireless power allocation problem.

PDI reduces this limitation by learning a λ-conditioned family of score fields.

In Figure 4, PDI is more robust to the shift in rmin because the score function ∇ log qT −t(xt λt) does not depend on rmin directly.

Comparison with Baselines

PDI exhibits the best objective and mode diversity.

Figure 2 shows that while PDI, PDL and PDM maintain full feasibility, PDI exhibits the best objective and mode diversity.

The unconstrained model ignores the constraints entirely and achieves the lowest objective.

The ablation on the dual variable through the unconstrained baseline shows that it achieves the lowest objective.

PDI maintains high feasibility and very small mean violations while achieving the highest return among the feasible diffusion-based methods.

Table 2 shows that PDI provides the best tradeoff between return, feasibility, and diversification among all diffusion samplers.

Limitations and Robustness

The convergence guarantees of our algorithm apply to time-averaged dual iterates rather than final iterates or the primal distributions directly.

The convergence guarantees of our algorithm apply to time-averaged dual iterates rather than final iterates or the primal distributions directly.

PDI requires problem-specific tuning of the dual-variable prior used during training and the entropy regularization strength β, for which principled selection rules remain an open question.

PDI requires problem-specific tuning of the dual-variable prior used during training and the entropy regularization strength β, for which principled selection rules remain an open question.

The noise schedule comparisons of Figure 14 show that the schedule is important for convergence.

Schedules that keep the process noisy for longer, such as the standard cosine and DDPM linear schedules, give more stable behavior and better feasibility.

The paper concludes by showing that PDI reduces the limitation of using an unconditional score model by learning a λ-conditioned family of score fields.

PDI reduces this limitation by learning a λ-conditioned family of score fields.

Improvements for AI systems

Bold dual variable-conditioned score network training allows PDI to generalize across a family of Gibbs distributions indexed by the dual variable encountered during inference. This enables the AI system to adapt effectively to shifted constraints and out-of-distribution (OOD) instances by learning a score field that is conditioned on the current multiplier, rather than estimating it beforehand.
Primal–dual inference (PDI) allows the sampler to steer trajectories toward feasibility by interleaving dual ascent with reverse diffusion, meaning the sampler follows a path-dependent sequence of Gibbs score fields. This enables the AI system to navigate complex constraints in problems like wireless resource allocation or portfolio management where solutions are probabilistic rather than deterministic.
The convergence analysis proves that the time average of the dual variables generated along the inference trajectory converges to a neighborhood of the dual optimum, which means the time-averaged multiplier λ¯ defines a near-optimal Gibbs target. This allows for robust sampling by using this time-averaged multiplier to define a fixed, high-quality target distribution.
The stability analysis bounding the Wasserstein distance shows that a dual mismatch at time t affects the final distribution through the future stability product QT s=t+1 ρt, allowing designers to control how early dual errors are damped rather than amplified by carefully selecting noise schedules.
PDI outperforms other methods because it keeps multiplier updates coupled to the denoising trajectory during inference, meaning the constraint correction and denoising dynamics remain coupled along the trajectory, allowing constraint violations to influence subsequent denoising steps more directly. This coupling leads to better performance than fixed-dual variants like DT.

Abstract

This paper develops constrained diffusion models with primal-dual inference (PDI) to sample from optimal distributions of entropy-regularized optimization problems with average constraints. We formalize constrained sampling in the Lagrangian dual domain, where the optimal distribution takes the form of a Gibbs distribution indexed by the optimal dual variable. Rather than estimating this dual multiplier before sampling and freezing it throughout generation, PDI jointly infers the optimal primal distribution and its parametrizing dual variable. Each reverse diffusion step denoises using the score field associated with the current multiplier and then updates the multiplier through dual ascent using the estimated constraint violation of the denoised samples. To enable this conditional score field, we train a single dual-conditioned score network over the family of Gibbs distributions induced by the dual variables encountered during inference. We prove that the dual iterates converge to a neighborhood of the dual optimum, and that the generated distribution approaches that of an oracle sampler that knows the optimal multiplier, up to a bias floor. We evaluate PDI on constrained sampling from a mixture of Gaussians, wireless resource allocation, and portfolio management.