Technical Summary: Constrained Diffusion Models with Primal-Dual Inference

Problem Statement

Standard diffusion models excel at sampling from complex, high-dimensional distributions in unconstrained settings (e.g., image and audio synthesis). However, they lack robust mechanisms to enforce statistical (average) constraints during the sampling process. Many real-world optimization problems—such as wireless resource allocation and portfolio management—require solutions that are probabilistic rather than deterministic, where the optimal policy is a distribution satisfying constraints in expectation (e.g., per-user rate guarantees or risk budgets).

Existing approaches to constrained sampling often rely on:

1. Pointwise constraints: Enforcing constraints on every individual sample (e.g., via projections or boundary reflections), which can be overly conservative and degrade optimality.
2. Dual Training (DT): Estimating the optimal Lagrange multiplier (dual variable) before sampling and freezing it. This approach is fragile because the optimal multiplier is defined by intractable expectations under the very distribution being learned, requiring expensive retraining for each new problem instance or constraint shift. Furthermore, a frozen sampler cannot adapt to constraint violations observed during generation.

Methodology: Primal-Dual Inference (PDI)

The paper proposes Primal-Dual Inference (PDI), a framework that treats the dual variable not as a pre-computed parameter, but as a dynamic state of the sampler that evolves jointly with the reverse diffusion process.

1. Theoretical Formulation

The problem is cast as an entropy-regularized optimization in the Lagrangian dual domain. The optimal distribution $\mu^*$ is shown to be a Gibbs distribution indexed by the optimal dual variable $\lambda^*$:
$$\mu^*(x) \propto \exp\left(-\frac{1}{\beta}(f_0(x) + \lambda^{*\top}f(x))\right)$$
where $f_0$ is the objective, $f$ represents the constraints, and $\beta$ is the temperature.

2. The PDI Algorithm

PDI interleaves primal denoising (reverse diffusion steps) with dual ascent (updating the multiplier) at every inference step $t$:

• Primal Step (PDI-P): The sampler updates the noisy sample $x_t$ using a score field conditioned on the current dual variable $\lambda_t$.
• Dual Step (PDI-D): The dual variable is updated via gradient ascent based on the estimated constraint violation of the denoised samples. Crucially, the constraint violation is estimated using Tweedie's formula (posterior mean estimates) rather than clean samples, allowing the update to occur without access to ground truth.
  $$\lambda_{t+1} = \left[ \lambda_t + \eta_t \mathbb{E}_{x_{t+1}} \left[ f(\hat{y}_0) \right] \right]_+$$
  where $\hat{y}_0$ is the Tweedie estimate of the clean sample.

3. Score Network Training

To enable this dynamic coupling, the authors train a single dual-conditioned score network $s_\theta(x_t, t, \lambda_t)$. Unlike standard models trained on a single target distribution, this network learns the family of Gibbs score fields indexed by the range of dual variables encountered during the inference trajectory. The training procedure combines:

• Exploitation: Rolling out the current model to collect trajectories of $(x_t, \lambda_t)$ pairs.
• Exploration: Sampling dual variables from a prior distribution to cover the space of potential multipliers.
• Monte Carlo Estimation: Using MC estimators to approximate the true score of the Gibbs distribution for regression targets.

Key Contributions

The paper outlines five primary contributions:

1. Saddle-Point Formulation: Casting constrained sampling with average constraints as a saddle-point problem in the Lagrangian dual domain, identifying the optimal distribution as a Gibbs distribution indexed by the dual variable.
2. PDI Algorithm: Introducing a reverse-diffusion sampler where primal denoising and dual ascent evolve jointly. The multiplier is updated at every step based on constraint violations estimated from denoised samples.
3. Generalized Score Network: Training a single network conditioned on the dual variable, allowing the model to generalize across a family of Gibbs distributions rather than a single fixed target.
4. Convergence Analysis: Proving that the time-averaged dual iterates converge to a neighborhood of the optimal multiplier at a rate of $O(1/\sqrt{T})$. The authors also bound the effect of residual dual mismatch on the terminal distribution, showing that early high-noise steps dampen errors while later steps refine the solution.
5. Empirical Validation: Demonstrating the method's efficacy on three distinct stochastic optimization problems: Gaussian mixture sampling, wireless power control, and portfolio management.

Experimental Results

The authors evaluate PDI against baselines including Projected Diffusion Models (PDM), Primal-Dual Langevin (PDL), Diffusion Posterior Sampling (DPS), and Dual Training (DT).

• Mixture of Gaussians: PDI achieves better objective values (closer to the unconstrained optimum) while maintaining full constraint feasibility. Unlike PDM and PDL, which enforce constraints on every sample (leading to conservative solutions), PDI satisfies constraints in expectation, allowing for more diverse and optimal samples.
• Wireless Resource Allocation: PDI achieves the best balance between mean sum-rate (optimality) and feasibility (satisfying minimum rate constraints for users). It significantly outperforms PDM and PDL in feasibility metrics (e.g., 5th-percentile rates) while maintaining high mean rates. It also demonstrates robustness to out-of-distribution (OOD) constraint shifts (e.g., different $r_{min}$ values) compared to DT.
• Portfolio Management: PDI provides the best trade-off between expected return, risk feasibility, and portfolio diversification (entropy). It outperforms unconstrained models (which violate risk budgets) and overly conservative baselines (which sacrifice return).

Ablation Studies:

• Fixed Multipliers: Freezing the dual variable (either the final iterate or the time-average) during inference results in inferior performance compared to the full PDI dynamics, confirming that the path-dependent evolution of the multiplier is critical.
• Noise Schedule: The choice of noise schedule is vital. Schedules that maintain a low Signal-to-Noise Ratio (SNR) for longer (e.g., standard cosine) allow the dual variable to stabilize before the reverse process becomes sensitive to errors. Schedules that reconstruct clean samples too early amplify dual mismatches, leading to constraint violations.
• Dual Training vs. PDI: While DT can reach informative dual variables, it fails to match PDI's performance because its score model is an "averaged" field over the training trajectory and cannot adapt to specific constraint residuals during inference.

Significance and Claims

The paper claims that PDI bridges the gap between diffusion-based generative modeling and constrained stochastic optimization. Its significance lies in:

• Dynamic Adaptation: Unlike prior methods that require pre-computing and freezing the dual multiplier, PDI treats the multiplier as an inference-time state, allowing the sampler to react instantaneously to constraint violations.
• Efficiency and Robustness: By training a single network over a family of distributions, PDI avoids the expensive retraining required by Dual Training for new problem instances or constraint shifts.
• Theoretical Guarantees: The work provides convergence guarantees for the time-averaged dual variables and stability bounds for the reverse process, offering a theoretical foundation for constrained diffusion sampling.

The authors note limitations, specifically that convergence guarantees apply to time-averaged dual iterates rather than the final primal distribution directly, and that the method requires problem-specific tuning of hyperparameters (e.g., entropy regularization strength $\beta$ and dual priors). Nevertheless, PDI is presented as a robust framework for sampling from optimal distributions in constrained optimization problems where solutions must be probabilistic.