Diffusion models as plug-and-play priors

In plain words

Diffusion models are generators that learn to make images by undoing noise in many small steps. Once trained, they capture what realistic images look like. This paper asks whether that knowledge can be reused as a plug-in part. The generator stays as it is, and any other rule that can be written as a smoothly changing score is added on top. The rule might be a classifier that says whether a face is smiling, or a check that a land-cover map agrees with colour clusters and rough labels.

The search looks for an input that both looks realistic to the generator and satisfies the rule. It takes small gradient steps on a candidate. At each step the candidate is noised a little, the generator reports how wrong its guess of that noise was, and the rule adds its own penalty. The noise level starts high and falls over the run. Early steps shape the overall layout, and late steps add detail. The search itself needs no training, which is the main appeal once a generator exists.

The authors try the idea on handwritten digits, face attributes, land-cover maps, and a travelling-salesman problem written as a picture of a route. The price is compute: every sample needs hundreds to thousands of passes through the generator.

TL;DR

• Takes a denoising diffusion model trained on its own data as the prior p(x), and combines it with a differentiable constraint c(x, y) at inference time, with no further training of the diffusion model §1 §2.1.
• Inference (Algorithm 1) minimises the denoising error minus log c over a point estimate, drawing fresh noise at each step, with the noise level t following a schedule annealed from high to low §2.2 alg. 1.
• Author-reported: conditional MNIST and FFHQ samples, judged visually fig. 1 fig. 2 fig. 3; land-cover transfer comparable to Rolf et al. tab. 1; TSP tours with a 1.28% gap at N=50 without beam search tab. 2.
• Costs: many denoiser passes per sample; for TSP, a greedy tour extraction and a naive 2-opt pass; one run per experiment, with no error bars.

Problem

The goal is the posterior p(x|y) proportional to p(x) c(x, y), where p(x) is a fixed prior and c is a constraint on x given extra information y §2.1. For a DDPM prior, the marginal p(x) requires an intractable integral over the latent chain x_T ... x_1 eq. 5. DDPMs are trained with a forward noising process eq. 6, so the marginal cannot be plugged in directly §2.2. Without this method, the usual route is to train a conditional model or to fine-tune with paired data §1. The method assumes the constraint is differentiable, or is a differentiable function of some latent y that deterministically produces x §2.2.

Method

Objective. Fix y and introduce an approximate posterior q(x). The free energy is eq. 1:

F=-\mathbb{E}_{q(x)} [\log p(x) + \log c(x,y) - \log q(x)] eq. 1

Minimising F gives the true posterior when q is expressive enough, and otherwise a mode-seeking approximation §2.1.

Point estimate. Take q(x) as a Dirac delta at η eq. 7. The DDPM term becomes a sum of KL divergences between Gaussian conditionals eq. 9. With the standard reparametrisation, that sum becomes a sum of expected denoising errors, with weights w_t that equal 1 in most pretrained diffusion models §2.2 eq. 10:

\sum_t w_t(\beta) \mathbb{E}_{\epsilon}\big[\|\epsilon - \epsilon_\theta(x_t,t)\|_2^2\big], \quad x_t = \sqrt{\bar\alpha_t}\,\eta + \sqrt{1-\bar\alpha_t}\,\epsilon eq. 10

With w_t = 1, the objective to minimise over η is eq. 12:

F = \sum_t \mathbb{E}_{\epsilon}\big[\|\epsilon - \epsilon_\theta(x_t,t)\|_2^2\big] - \log c(\eta, y) eq. 12

The paper minimises F over η by sampling ε inside each summand over t §2.2. Each run is stochastic and can return a different η §2.2.

Annealing. Optimising all t at once makes it hard to reach a mode for images. The paper therefore anneals t from high to low: the first iterations explore the search space coarsely, and later iterations reduce the noise and settle into a nearby local maximum of p(x|y) §2.2. If η starts near a good mode, the sum can stop at t_max < T §2.2.

Algorithm 1 (point estimate under a DDPM prior and constraint) alg. 1:

# Point-estimate inference under a DDPM prior  [alg. 1]
# Input:  denoiser eps_theta, data y, constraint c, time schedule, rate lr
# Output: point estimate eta

def infer_point_estimate(denoiser, y, constraint, schedule, lr):
    # Start from pure Gaussian noise  [alg. 1]
    x = random_normal(shape)
    for t_i in schedule:
        # Draw fresh noise and noise the current estimate  [alg. 1]
        eps = random_normal(shape)
        x_t = sqrt(alpha_bar[t_i]) * x + sqrt(1 - alpha_bar[t_i]) * eps
        # Gradient step on the denoising error minus the log-constraint  [alg. 1]
        denoise_error = sum_of_squares(eps - denoiser(x_t, t_i))
        loss = denoise_error - log(constraint(x, y))
        x = x - lr * gradient_wrt_x(loss)
    # The final iterate is the estimate  [alg. 1]
    return x

The gradient is taken with respect to x through the denoiser alg. 1. The paper notes that many variants of this loop are possible §2.2.

Latent parametrisation. If x = f(y) is deterministic, the constraint is a hard one. Steps 4 and 5 of Algorithm 1 are then replaced by gradient steps on y for the objective eq. 13. If the prior on the latent is pushed forward by f, the objective must be weighted by the Jacobian of f §2.2:

\|\epsilon - \epsilon_\theta(\sqrt{\bar\alpha_{t_i}} f(y) + \sqrt{1-\bar\alpha_{t_i}}\epsilon,\ t_i)\|_2^2 eq. 13

Gaussian posterior (appendix only). Appendix A derives the same objective for a scalar-covariance Gaussian q(x) = N(sqrt(1-ψ) η, ψ I) eq. 16, with the reparametrisation trick for the constraint term app. A eq. 21. This variant is derived but not used in the experiments app. A.

Instantiations. Each experiment picks its own constraint c:

• MNIST (conditional digits). Hand-crafted symmetry terms and a separately trained digit classifier, which estimates the likelihood of the target digit §3.1.

• FFHQ faces. The constraint is the sum of log-likelihoods of the chosen attributes under a ResNet-18 attribute classifier eq. 14 §3.2. The denoiser runs from t = 200 at the end to remove classifier artefacts §3.2.

  \log c(x,y) = \sum_{y \in \mathbf{y}} \log p(y \mid x) eq. 14

• Land-cover segmentation. The constraint combines the mutual information between predicted labels and a local colour clustering, and the KL divergence between predicted and weak-label distributions in blocks §4 app. B.3:

  \log c(x, z, \ell_{weak}) = \mathrm{MI}(x, z) - \mathrm{KL}(x \,\|\, \ell_{weak}) eq. 15

  The sample is initialised with the weak labels, which lets inference start from a smaller t §4.

• Travelling salesman. A tour is a symmetric matrix A = S + S^T, where S is row-wise softmax of a parameter matrix with zero diagonal §5. A DDPM trained on 64×64 pictures of solved tours serves as the prior on the picture f(A) of A §5. There is no tour-length term. After optimisation, A is rounded to a tour: edges are inserted greedily in decreasing order of A_ij / ||v_i − v_j||, and then a naive 2-opt pass runs §5 app. B.4.

Hyperparameters

Values marked paper are stated in the paper. Values marked code appear only in the official notebooks or scripts (see Paper vs code).

Results

Author-reported. Setup: land-cover patches from Pittsburgh (training) and three other EnviroAtlas cities; accuracy and class-mean IoU in % §4 tab. 1.

Source: tab. 1. The caption says the second-to-last row (Ours, fine-tuned) was pretrained on Pittsburgh labels tab. 1.

Author-reported. Setup: Euclidean TSP test sets of 1280 graphs each, N = 50 and N = 100 §5. Objective is mean tour length; gap is relative to Concorde tab. 2.

Source: tab. 2 (left). The paper compares with baselines that do not use beam search; it notes that the baselines improve with large beams, at growing cost §5.

Number of denoising steps for the 50-node model, with the mean count of 2-opt steps per instance tab. 2 (right):

Appendix results for 200-node instances, with the 20–50 model: optimality gap 3.77% (219 uncrossing moves on average) against 3.81% for 2-opt (115 moves) app. B.4. With 128×128 encodings at inference, the gap is 2.59% (81 moves) app. B.4.

MNIST and FFHQ have no numbers in the paper. Their evidence is visual: sample grids for conditional digits fig. 1, face attributes before and after the final denoise fig. 2, and the intermediate steps of one face fig. 3. The paper does not report error bars: each main experiment was run once § checklist.

Paper vs code

Official repository at 80e4ad8, MIT licence, belonging to the first author (the LICENSE copyright line matches). The code is given as notebooks for MNIST, FFHQ, land cover, and scripts for TSP.

• TSP learning rate. Paper: linear decay from 1 to 0.1 app. B.4. Code: LinearLR with end_factor=0.1 and total_iters=1000, but only 256 steps run, so the rate ends near 0.77 (1 − 0.9 × 256/1000) inference.py inference.py.
• TSP time schedule. Paper: the MNIST schedule, linearly resampled to the step count app. B.4. Code: (steps − i) + (steps − i) // 3 · cos(i/50) inference.py. The MNIST notebook uses // 4 inference.ipynb.
• TSP encoding. Paper: superimposed line segments with intensity A_ij and small black dots at the vertices §5. Code: vertices are drawn at the maximum value (point_color=1), and lines at 0.5 before rescaling to [-1, 1]; the encoded pixels at vertices are then overwritten with 1 at every step inference.py inference.py. The overwrite is not described in the paper.
• TSP softmax. Paper: S is a softmax of a parameter matrix over rows §5, with no scale given. Code: logits are −0.1 × latent before the row softmax, and the diagonal is masked with −1e5 inference.py.
• TSP 2-opt acceptance. Paper: a 2-opt move is made when two tour edges cross app. B.4. Code: a reversal is accepted when the full tour length decreases tsp_utils.py.
• Denoising loss scale. Paper: squared L2 norm of the noise residual, summed over pixels eq. 10 eq. 12. Code: F.mse_loss, the mean over elements inference.py train.py inference.ipynb, inference cell. The MNIST and land-cover constraint weights (0.01, 1) were set against this mean-scaled loss, so their effect differs from the paper's sum-scaled objective.
• FFHQ learning rate. Paper: linearly decreasing from 1 to 0.5 app. B.2. Code: LinearLR(start_factor=1, end_factor=1), so the rate stays at 1; the comment says "Optional" infer_with_attributes.ipynb, inference cell.
• FFHQ gradient clipping and updates. Paper: clip the auxiliary gradient to ½ of the denoising gradient norm app. B.2, in one combined update alg. 1. Code: clips to 0.1 × an exponential moving average (α = 0.5) of the denoising gradient norm, and runs two separate Adamax steps per iteration, one for each loss infer_with_attributes.ipynb, inference cell.
• FFHQ default constraint. Paper: the example in Figure 3 is blonde + smiling + not male §3.2 fig. 3. Code: the default mask maximises the logits of Young and Goatee; the comment beside the Young target reads "0 corresponds to young, 1 corresponds to old", which is the opposite of the attribute name, so check the sign before reusing it infer_with_attributes.ipynb, target-attributes cell.
• FFHQ classifier. Paper: ResNet-18 trained on CelebA §3.2. Code: loads models/Resnet18.pth, which is not in the repository; the README says any pretrained face-attribute classifier can be used attr_classifier.py README.md.
• FFHQ denoiser output. The denoiser has six output channels; only the first three are used as the noise prediction diffusion.py. The paper does not say this.
• MNIST learned constraint. Paper: the log-likelihood of the digit under the classifier §3.1. Code: the sigmoid output is penalised directly, so the term is linear in probability and not in log-probability inference.ipynb, learned-condition cell.
• MNIST constraints. Paper: panels (a) and (b) use mean intensity to make digits thin or thick §3.1 fig. 1. Code: no intensity term appears in the notebook; the symmetry terms are the only hand-crafted ones inference.ipynb, hand-crafted cell.
• Land-cover weak-label KL. Paper: KL between block-averaged predicted labels and weak labels, in non-overlapping 31×31 blocks app. B.3. Code: the predicted map is average-pooled with stride 1 (overlapping); the weak labels are not pooled; the non-overlapping pooling is commented out inference.ipynb, NLCDLoss cell.
• Land-cover KL direction. Paper: KL(x ‖ ℓ_weak) eq. 15. Code: F.kl_div(log(predicted), weak) computes KL(ℓ_weak ‖ predicted) inference.ipynb, NLCDLoss cell.
• Land-cover mutual information. Paper: MI(x, z), unnormalised eq. 15. Code: divides by H(predicted) + H(clustering), with the clustering entropy detached inference.ipynb, MILoss cell.
• Land-cover classes. Paper: 5-class land cover §4. Code: six channels, indices 0–5, where 0 is "Unclassified" utils.py.
• Land-cover inference settings. Paper: 200 steps, Adam lr 5e-3 → 5e-6, betas (0, 0.999) app. B.3. The dataset-generation cell matches this. The single-patch inference cell uses 500 steps, lr 1e-3 → 1e-4, and default Adam betas inference.ipynb, run-inference cells.
• Land-cover denoise pass. Paper: no final denoising pass is described for segmentation. Code: runs diffusion.inverse from t = 200 (single patch) or t = 100 (dataset) inference.ipynb.
• Land-cover clustering. Paper: fixed number of components, unspecified app. B.3. Code: 5 components, applied to NAIP tiles pooled with a 3×3 kernel at stride 4 (256 → 64 px), 30 EM iterations inference.ipynb, overlap-clustering cell overlap_cluster.py.
• Land-cover DDPM training. The training notebook passes image_size=256 to the U-Net while training on 64×64 labels; the inference notebook uses 64 train_diffusion.ipynb inference.ipynb.
• Gaussian variant. Appendix A has no matching code; the repository has no ψ-parametrised posterior app. A.

Pitfalls

• Noise-level index. The code indexes the cumulative schedule at t − 1 with t from 1 to T diffusion.py. Off-by-one errors are likely if the schedule is reimplemented from the paper.
• Typo in ᾱ. Section 2.2 writes ᾱ_t as a product with the index on α_t instead of α_i §2.2; the appendix repeats the same form app. A. The intended product runs over the index i, so use ᾱ_t = ∏ α_i. The appendix also sets β_0 = ψ, which applies only to the Gaussian variant app. A.
• Algorithm is plain gradient descent; the experiments use Adam. Algorithm 1 is written with a single rate λ alg. 1, but the appendix and code use Adam (or Adamax) with per-coordinate scaling app. B.1 app. B.2. Constraint weights such as 0.01 were tuned in that setting and do not carry over to plain gradient descent.
• Schedules are the heavy lifting, and they are not in the paper. The exact t schedule formula is only in code, and each notebook uses different constants: MNIST k // 4, TSP k // 3 (cos(i/50)), and FFHQ (k/1.5 + (k/3) cos(i/10)) / steps · 800 + 200, where k = steps − i. The paper describes the shape only.
• Random jitter on t (FFHQ). Each step adds a random integer in [−50, 50] to t, then clips to [1, T] infer_with_attributes.ipynb, inference cell. Not described in the paper.
• Attribute-classifier augmentation (FFHQ). Eight random views per step, with horizontal flip, brightness ±20%, and Gaussian blur, each applied with probability ½, and the image resized to 224 px with nearest-neighbour interpolation infer_with_attributes.ipynb, inference cell. Not in the paper.
• Probability mapping (segmentation). Labels are converted to probabilities by the normal CDF of N(η, σ) over [0.5, 1.5], then normalised eq. 15 inference.ipynb, encode_prob. As σ falls to 0.02, the map becomes nearly a step function, and gradients concentrate at the 0.5 and 1.5 boundaries.
• Mutual-information window. A window schedule is written in code but overridden to 31 in the same cell inference.ipynb, run-inference cell. The paper's fixed 31×31 is what runs.
• TSP hidden cost. Rounding the optimised matrix does not always give a tour; the paper reports disjoint cycles as a common failure §5. The naive 2-opt in code recomputes the full tour length for each candidate pair, so each pass costs O(N³) for N cities tsp_utils.py. The paper argues that cost grows sublinearly with graph size because most inference time is spent on the denoiser, which runs on a fixed-size image §5. That argument does not cover the extraction and 2-opt steps.
• TSP fixed-point overwrite. The node pixels are overwritten at every step inference.py. Without it, the optimised picture drifts away from the input points.
• Data and weights outside the repo. The TSP training data must be downloaded separately from the repository linked in the README. The 100-node test set is not referenced, and the 50-node model is the only TSP checkpoint linked README.md. The FFHQ classifier and DDPM come from other repositories.
• Single runs. Results are one run each, and inference is stochastic §2.2 § checklist. Compare with several seeds before reading small gaps as real.
• Cost. MNIST inference uses 1000 steps, each with a denoiser forward and backward pass app. B.1 inference.ipynb. FFHQ uses 200 steps with eight classifier passes each app. B.2 infer_with_attributes.ipynb.
• Padding and input range (MNIST). The code pads 28 px digits to 32 px and maps them to [-1, 1] train_diffusion.ipynb, Load Data cell. The paper does not state this.
• Classifier for MNIST. A three-layer convolutional classifier with stride-2 convolutions, trained on "digit ≠ 3" labels with binary cross-entropy, for 5 epochs at Adam lr 1e-4 inference.ipynb, classifier cell. The paper says only that it was trained separately §3.1.

Lineage

• Builds on: denoising diffusion probabilistic models (Sohl-Dickstein et al., 2015; Ho et al., 2020) §2.2, and the U-Net of Dhariwal et al. (2021) §2.2 app. B.1. None of these has a card in the wiki yet.
• Uses a pretrained FFHQ DDPM from Baranchuk et al. (2022) §3.2.
• Related: score-based inverse-problem solvers (Song et al., 2021; Kawar et al., 2021; Kadkhodaie et al., 2021) §1; diffusion-based purification (Nie et al., 2022) §1.
• Compared against: Rolf et al. (2022) for land-cover transfer tab. 1; Kool et al. (2019), Joshi et al. (2019) and Bresson et al. (2021) for TSP tab. 2.