Diffusion models as plug-and-play priors

Alexandros Graikos; Nikolay Malkin; Nebojsa Jojic; Dimitris Samaras

Source-linked AI summary

Diffusion models as plug-and-play priors

Alexandros Graikos, Nikolay Malkin, Nebojsa Jojic, Dimitris Samaras

arXiv:2206.09012v3Jun 17, 2022cs.LGcs.CV

TL;DR

The paper makes independently trained denoising diffusion models usable as plug-and-play priors under differentiable constraints, without additional training or finetuning. It applies this framework to conditional generation, domain-transfer segmentation, and relaxed combinatorial search, obtaining realistic constrained faces, comparable segmentation performance, and competitive TSP results.

Problem
Denoising diffusion models capture complex high-dimensional distributions but their costly inference has limited exploration as plug-and-play modules in systems combining multiple models.
Method
The authors perform approximate constrained inference by differentiating through a fixed denoising diffusion network at multiple noise levels alongside diverse differentiable constraints, without retraining model components.
Results
Across applications, the framework produces realistic attribute-conditioned faces, achieves segmentation performance comparable to state-of-the-art weakly supervised work, and yields competitive TSP solutions with better out-of-distribution generalization than baselines.
Takeaways & Limitations
Diffusion priors can support conditional generation, reduce sensitivity to conditioning-data domain shift, and generate candidates for combinatorial search.
Takeaways & Limitations
Plug-and-play inference inherits DDPMs’ high cost because generating samples requires many passes through the denoising network.

Abstract

from arXiv · show

We consider the problem of inferring high-dimensional data $\mathbf{x}$ in a model that consists of a prior $p(\mathbf{x})$ and an auxiliary differentiable constraint $c(\mathbf{x},\mathbf{y})$ on $x$ given some additional information $\mathbf{y}$. In this paper, the prior is an independently trained denoising diffusion generative model. The auxiliary constraint is expected to have a differentiable form, but can come from diverse sources. The possibility of such inference turns diffusion models into plug-and-play modules, thereby allowing a range of potential applications in adapting models to new domains and tasks, such as conditional generation or image segmentation. The structure of diffusion models allows us to perform approximate inference by iterating differentiation through the fixed denoising network enriched with different amounts of noise at each step. Considering many noised versions of $\mathbf{x}$ in evaluation of its fitness is a novel search mechanism that may lead to new algorithms for solving combinatorial optimization problems.

1 Introduction

The paper proposes using independently trained DDPMs as plug-and-play priors alongside differentiable constraints. This supports conditional generation, domain-transfer tasks, and continuous relaxations of combinatorial search.

1 Introduction: DDPMs model complex distributions over high-dimensional continuous data but require many denoising evaluations during generation.The computational depth complicates their use as off-the-shelf components in hierarchical generative models.
1 Introduction: The inference target combines a prior p(x) with a constraint c(x, y), yielding an approximation to p(x|y) ∝ p(x)c(x, y).The independently trained DDPM supplies the prior, making it a plug-and-play module.
1 Introduction: The paper addresses an unexplored use of DDPMs as plug-and-play modules without additional training or finetuning of model components.This distinguishes the approach from earlier plug-and-play methods that used additional training or inference networks.
1 Introduction: Applications include classifier-constrained conditional generation, segmentation under appearance or weak-label constraints, and continuous relaxations of combinatorial search.The constraint can come from an off-the-shelf classifier or encode deterministic combinatorial structure.
1 Introduction: Prior work used unconditional generative models or pretrained DDPMs for conditional generation and inverse imaging, whereas this paper targets broader differentiable constraints.The paper positions its independently trained, pixel-space components within this related work.

2 Method

The method minimizes a variational free energy under a fixed DDPM prior and differentiable constraint. It evaluates denoising consistency across noise levels, using annealed stochastic optimization to obtain point or uncertain posterior approximations.

2.1 Problem setting: The free energy is minimized when the approximate posterior q(x) is closest to p(x|y), though an expressive limitation can produce a modeseeking approximation.A point-mass q concentrates at a posterior mode.
2.1 Problem setting: The procedure minimizes the free energy with respect to an approximate posterior for any differentiable constraint when the prior is a DDPM.The prior and constraint can remain independently trained and operate directly in pixel space.
2.2 Denoising diffusion probabilistic models as priors: A DDPM reverses a Gaussian noising process through learned transitions from x_T to x_0, starting from a unit Gaussian and following a noise schedule.The latent trajectory h makes the prior likelihood involve an otherwise intractable integration over intermediate states.
2.2 Denoising diffusion probabilistic models as priors: The method replaces direct prior evaluation with denoising reconstruction terms evaluated on noised versions of x, together with the constraint term.The forward process supplies the noisy states, while the fixed denoiser is differentiated through during inference.
2.2 Denoising diffusion probabilistic models as priors: Algorithm 1 performs stochastic gradient updates from an initial Gaussian sample while traversing selected noise levels.Different time-step schedules can trade off optimization behavior and diversity in estimated point solutions.
2.2 Denoising diffusion probabilistic models as priors: Annealing noise levels from high to low first explores the search space coarsely and then approaches a nearby local posterior maximum.A Gaussian posterior approximation can additionally represent uncertainty around a mode.

3 Experiments: Conditional image generation

The experiments use unconditional diffusion priors with hand-crafted or independently trained constraints to generate MNIST digits and FFHQ faces with specified attributes. The diffusion objective preserves realistic image structure while constraint optimization enforces desired properties, with denoising used to reduce artifacts.

3.1 Simple illustration on MNIST: MNIST experiments generate digits by optimizing constraints for thickness, symmetry, digit class, and combinations of these properties.The constraints include image intensity, fold-wise L2 symmetry penalties, and an independently trained classifier likelihood.
3.1 Simple illustration on MNIST: An independently trained classifier can serve as the constraint, while the DDPM remains an unconditional prior.For example, a classifier distinguishes digit class y = 3 from other digits, and its likelihood guides inference.
3.2 Using off-the-shelf components for conditional generation of faces: FFHQ experiments combine an unconditional DDPM with a pretrained face-attribute classifier to enforce selected features such as smiling, blond hair, and gender.Only a small subset of attributes is strictly enforced, allowing convergence toward different modes exhibiting the desired features to varying degrees.
3.2 Using off-the-shelf components for conditional generation of faces: The diffusion objective counteracts classifier-driven adversarial solutions, guiding optimized samples toward realistic images that satisfy the requested attributes.Without denoising loss, optimization can produce adversarial noise that fools the classifier network.

4 Experiments: Semantic image segmentation

The segmentation experiments apply diffusion priors trained on land-cover labels to infer high-resolution semantic segmentations in geographically different cities from weak labels and color structure. The method reduces inference steps through weak-label initialization and supports domain transfer comparable to prior weakly supervised work.

4 Experiments: Semantic image segmentation: The EnviroAtlas task infers 5-class, 1m-resolution land-cover labels in Durham, Austin, and Phoenix using weak labels, with high-resolution labels available only for Pittsburgh.The cities provide geographically diverse domains for testing transfer beyond the labeled source region.
4 Experiments: Semantic image segmentation: A DDPM is trained on one-hot land-cover label representations and its samples are converted into per-pixel label probabilities for inference.The conversion uses a Gaussian-based probability model with user-defined σ; categorical diffusion is noted as an alternative.
4 Experiments: Semantic image segmentation: The segmentation constraint combines mutual information between local color clusters and predicted labels with negative KL divergence matching weak-label distributions.The mutual-information term operates in overlapping patches, while weak-label matching uses averages over non-overlapping blocks.
4 Experiments: Semantic image segmentation: Initializing inferred segmentations with weak labels reduces the number of optimization steps and permits starting from a smaller noise level.Figure 5 reports examples from Durham, Austin, and Phoenix; Phoenix retains overall structure despite a substantially different color-label joint distribution.
4 Experiments: Semantic image segmentation: The resulting domain-transfer approach is comparable with state-of-the-art weakly supervised training, whereas training only on Pittsburgh supervised labels fails to generalize geographically.Table 1 reports accuracy and class mean intersection-over-union across geographic domains.

5 Experiments: Continuous relaxation of combinatorial problems

The paper applies DDPM priors to TSP by optimizing a continuous latent adjacency matrix through image-space denoising, then extracting and refining tours. The method produces competitive solutions, generalizes to larger graphs, and uses fixed-size image computation.

Method: The procedure performs gradient descent through 256 decreasing-noise denoising steps, extracts a tour, and refines it with greedy edge insertion and 2-opt uncrossing.The figure describes tour fragments being recombined during inference before final refinement.
Results: For both displayed TSP problems, inferred tour lengths are within 1% of optimum, while the full evaluation uses test sets of 1280 graphs with 50 or 100 vertices.Table 2 reports mean tour length and optimality gap against ground-truth tours.
Method: DDPM-based TSP inference optimizes a continuous latent adjacency matrix represented as a 64 × 64 image, combining learned tour structure with differentiable relaxation.The adjacency matrix is parameterized through a stochastic matrix and symmetrized before applying the diffusion prior.
Method: Unconditional samples from the trained DDPM demonstrate that the model represents images of solved TSP tours.The model is trained on approximately 1.5m encoded solved instances across 50-node and 20-to-50-node datasets.
Results: DDPM inference is competitive with non-beam-search baselines, matches the reported 2.18% gap of a 5000-beam baseline on 100-vertex problems, and generalizes better to 100 vertices.The 50-node and mixed-size models perform almost identically, while both outperform baseline transfer from 50-node to 100-node instances.
Implications: Most inference computation runs on fixed-size images, giving sublinear cost scaling with graph size while the graph remains reasonably representable in 64 × 64 pixels.The image-space design supplies the method’s visual inductive bias and computational scaling property.

6 Conclusion

The paper concludes that constrained DDPM inference supports conditional generation, domain-shift-oriented segmentation, and combinatorial search. It also identifies denoising inference cost as a central limitation inherited by the proposed algorithms.

Conclusion: Pretrained-classifier constraints enable conditional generation, while clustering or weak-label constraints can reduce sensitivity to conditioning-data domain shift.These applications are presented as consequences of performing inference under DDPM priors and differentiable constraints.
Limitations: DDPM inference remains costly because generating a sample requires many passes through the denoising network.The authors expect fewer-step DDPM inference procedures to make plug-and-play use more appealing.
Conclusion: TSP results suggest that DDPMs can reason over uncertain hypotheses in a way that mimics human puzzle-solving behavior.The authors connect this result to future research on generating candidates for combinatorial search.

Checklist

The checklist reports that the paper describes its contributions, scope, assets, and reproducibility materials, while noting that main experiments were run only once without error bars. It also records no immediate expected societal impacts but flags deepfake risks.

Reporting: The authors report that contributions and scope are accurately described and that work limitations are discussed in the conclusion and discussion.These checklist responses point to the relevant paper sections rather than adding technical findings.
Societal impacts: No immediate negative societal impacts are expected, but the checklist flags flexible conditional image generation as a potential deepfake risk.The risk is stated as a consideration rather than as an observed effect of the experiments.
Reproducibility: The paper reports reproducibility materials and training details for most experiments, with compute information available in the Appendix.The checklist indicates that code, data, instructions, splits, and hyperparameters are addressed through the Appendix and experiment sections.
Reproducibility: Main experiments were run one time, so error bars with respect to random seed were not reported.This is the checklist’s explicit response to repeated-experiment uncertainty reporting.
Assets and data: Existing assets are cited, datasets are free to use for research purposes, and no new assets or human-subject data are reported.The checklist marks consent, personally identifiable information, offensive content, and crowdsourcing items as not applicable.

A Deriving a Gaussian approximation to the posterior

This appendix derives a Gaussian approximation to the posterior under a DDPM by expressing the noising process and conditional reverse distributions, then optimizing the resulting expectation with reparameterization.

Approximation: The appendix repeats the posterior derivation using a Gaussian approximation instead of a point estimate.This changes the posterior representation used in the inference objective.
Forward process: The DDPM likelihood uses the forward noising chain q(h|x0) as the variational distribution over latent variables xT through x1.The chain is written as a product of transition distributions q(xt|xt−1).
Posterior construction: A scalar-covariance Gaussian provides a computationally convenient approximate posterior over x0.The derivation also permits sampling xt at arbitrary time steps and extracting q(xt−1|xt, ψ, η).
Reverse process: The reparameterized derivation links the denoiser’s noise reconstruction ϵθ(xt,t) to the model expectation µθ(xt,t).This link is used in the Gaussian reverse-process formulation.
Optimization: The Gaussian objective includes an expectation over a Gaussian variable, which can be optimized through the reparameterization trick.The constraint term is rewritten using noise sampled from a standard normal distribution.

B Experiment details and extensions

The experiments use annealed optimization through diffusion time, with task-specific schedules and progressively reduced constraint influence to balance structure and sample quality.

Inference uses 1000 Adam optimization steps with a 10^-2 learning rate and a cosine-modulated, linearly decreasing t schedule.The schedule is designed to guide inference from coarse overall structure toward finer details.
The conditional-loss weight decreases linearly from λT = 10^-2 to λ1 = 0 as optimization proceeds from T to 1.This reduces the condition’s influence at later diffusion times to improve sample quality.
The t annealing schedules avoid optimizing all T = 1000 diffusion times when shorter task-specific schedules suffice.The appendix illustrates this design for MNIST and Land Cover, with analogous schedules for TSP and FFHQ.

B.2 FFHQ

FFHQ experiments show that classifier-guided inference can satisfy face-attribute constraints, but conflicting attributes and classifier artifacts can reduce fidelity and visual quality.

Low-temperature denoising can improve visual quality after classifier-guided optimization, but larger temperatures may erase the intended attributes.Choosing t involves a trade-off between removing optimization artifacts and preserving the inferred attributes.
Conflicting attributes tend to resolve toward one randomly chosen direction rather than satisfying both conditions simultaneously.Examples involving not male with male-correlated attributes satisfy only one condition; blonde plus black hair may appear mixed but is classified exclusively as blonde.
Some failures arise when the classifier paints desired features that the diffusion prior does not complete into a coherent face.Examples include eyeglasses outlines, misleading hair curves, and comically drawn mustaches.
The appendix presents additional conditional samples and failure cases both before and after diffusion denoising used to remove classifier-induced artifacts.The top images show inference results, while the bottom images show denoised outputs.

B.3 Land cover

Land-cover inference combines local image-structure matching with weak-label guidance, enabling transfer beyond Pittsburgh while retaining some geography-related biases.

The land-cover DDPM learns road continuity, suburban building planning, and Pennsylvania-specific structures from 64 × 64 label patches.It was trained with a U-Net, T = 1000 diffusion steps, and 105 batches of size 32.
Land-cover inference uses 200 optimization steps with a learning rate decreasing from 5×10^-3 to 5×10^-6 because initialization starts from weak labels.The procedure therefore need not begin at t = 1000.
The appendix includes unconditional samples, segmentation results, and comparisons with and without weak-label guidance.
The transferred semantic model remains superior to an image-based one despite artifacts such as a tendency to add uninterrupted forested areas.These artifacts reflect Pennsylvania-related biases learned by the diffusion model.
The segmentation constraint matches local color-clustering structure and weak-label segment distributions averaged over non-overlapping 31 × 31-pixel blocks.Mutual information is used for local structure, while KL-divergence matches predicted and weak-label distributions.
Weak labels are derived from 30m-resolution NLCD labels augmented with building footprints, road networks, and waterways or waterbodies.Without this guidance, the procedure can confuse semantic classes while still producing segmentations likely under the diffusion prior.

B.4 TSP

The TSP extension optimizes a latent adjacency representation with the diffusion prior, extracts tours greedily, and studies generalization to larger image encodings and problem sizes.

The TSP diffusion model is trained on 64 × 64 images encoding ground-truth solutions, using a U-Net and T = 1000 diffusion steps.Each model is trained for 8 epochs with batch size 16.
Tour extraction ranks edges by Aij/∥vi − vj∥, inserts them while avoiding short cycles, and then applies 2-opt uncrossing refinement.The 2-opt baseline uses the same procedure with a uniform adjacency matrix.
The 200-vertex 128 × 128 example has a tour 2.12% longer than the optimal one.
3.77% optimality gap on 200-node problems is comparable to 3.81% for 2-opt, suggesting near-saturated generalization at this size.The dense 64 × 64 encoding makes many tour edges invisible.
Proposed future directions include larger encodings, coarse-to-fine local search using crops, and improved 2-opt move selection.The same denoising objective could be applied at multiple scales to refine tours locally.
2.59% optimality gap is obtained on 128 × 128 inference images, with an average of 81 uncrossings, for a model trained on 64 × 64 representations.The larger encoding is presented as a way to increase vertex count without changing the model.