Reduce, Reuse, Recycle: Compositional Generation with Energy-Based Diffusion Models and MCMC

Yilun Du; Conor Durkan; Robin Strudel; Joshua B. Tenenbaum; Sander Dieleman; Rob Fergus; Jascha Sohl-Dickstein; Arnaud Doucet; Will Grathwohl

Source-linked AI summary

Reduce, Reuse, Recycle: Compositional Generation with Energy-Based Diffusion Models and MCMC

Yilun Du, Conor Durkan, Robin Strudel, Joshua B. Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Sohl-Dickstein, Arnaud Doucet, Will Grathwohl

arXiv:2302.11552v6Feb 22, 2023cs.LGcs.AIcs.CVstat.ML

TL;DR

The paper addresses how to reuse and compose diffusion models when existing composition techniques fail or cannot express certain operators. It introduces MCMC-based samplers and an energy-based parameterization, finding notable improvements across diverse compositional generation settings. The work also identifies higher sampling and model-computation costs as important drawbacks.

Problem
Existing diffusion-model composition methods do not fully support compositional operators such as mixtures, and some compositions fail despite using pretrained models.
Method
The paper composes pretrained diffusion models without retraining, uses MCMC-derived samplers, and parameterizes diffusion models with explicit energy functions.
Results
The proposed approach yields notable improvements across domains, scales, and compositional operators, including 2D densities, cube composition, classifier-guided modeling, and text-to-image generation.
Takeaways & Limitations
MCMC sampling and energy-parameterized diffusion models provide supported ways to improve diffusion-model composition and enable additional compositional operators.
Takeaways & Limitations
More sophisticated MCMC samplers can take 5-times longer than typical diffusion sampling, while energy-parameterized models require double the memory and compute cost of score-parameterized models.

Abstract

from arXiv · show

Since their introduction, diffusion models have quickly become the prevailing approach to generative modeling in many domains. They can be interpreted as learning the gradients of a time-varying sequence of log-probability density functions. This interpretation has motivated classifier-based and classifier-free guidance as methods for post-hoc control of diffusion models. In this work, we build upon these ideas using the score-based interpretation of diffusion models, and explore alternative ways to condition, modify, and reuse diffusion models for tasks involving compositional generation and guidance. In particular, we investigate why certain types of composition fail using current techniques and present a number of solutions. We conclude that the sampler (not the model) is responsible for this failure and propose new samplers, inspired by MCMC, which enable successful compositional generation. Further, we propose an energy-based parameterization of diffusion models which enables the use of new compositional operators and more sophisticated, Metropolis-corrected samplers. Intriguingly we find these samplers lead to notable improvements in compositional generation across a wide set of problems such as classifier-guided ImageNet modeling and compositional text-to-image generation.

1 Introduction

Diffusion models provide reusable generative priors and support post-training conditioning, but their scale makes retraining costly. This work develops composition methods that reuse pretrained models without retraining.

Large generative models can require 10’s or even 100’s of billions of parameters and internet-scale datasets, making training or finetuning costly.
Diffusion models combine scalability, sample quality, and log-likelihood with the ability to guide sampling after training.
Classifier-based and classifier-free guidance reuse pretrained diffusion models for downstream conditional generation.
Model composition combines probabilistic models to create new models, building on approaches such as Mixtures-Of-Experts and Product-Of-Experts.
The paper introduces composition methods without retraining, MCMC-derived sampling remedies, energy-based parameterization, and evaluations from 2D data to high-resolution text-to-image generation.

2 Background

Diffusion models learn to reverse a progressively noised data distribution, while energy-based models represent distributions through unnormalized energies and require approximate sampling. Their score-based relationship supports classifier guidance, classifier-free guidance, and an energy-based parameterization with broader compositional flexibility.

Diffusion models: Diffusion models progressively add Gaussian noise until x_T approximately follows N(0, I), then learn a reverse distribution for ancestral denoising.The forward transition scales x_{t−1} by √(1−β_t) and adds Gaussian noise with variance β_t.
Diffusion models: Training uses analytically tractable noisy marginals and a noise-prediction objective to recover scores of the Gaussian-smoothed data distribution.The model estimates ∇_x log q_σt(x) through the learned noise predictor.
Energy-based models: Energy-based models omit the normalizing constant, gaining flexibility but losing efficient likelihood computation and direct sampling.Denoising score matching can learn the score of a Gaussian-smoothed data distribution, after which MCMC or numerical ODE methods provide approximate sampling.
Energy-based models: MCMC methods provide approximate EBM sampling, including ULA, Metropolis-corrected MALA, and momentum-augmented HMC variants.Metropolis correction accepts or rejects proposed updates to address discretization error from larger step sizes.
Diffusion–EBM connection: The energy-based diffusion parameterization defines the noise predictor as the negative gradient of a scalar potential, enabling more sophisticated samplers and additional composition forms.This parameterization explicitly models a noise-conditional potential function f_θ(x,t).
Guidance: Classifier guidance combines unconditional data scores with predictive-model scores, whereas classifier-free guidance derives predictive gradients from unconditional and conditional diffusion models.Classifier guidance can repurpose one unconditional generative model with different predictive models, while classifier-free guidance has reduced modularity.

3 Compositional Generation Beyond Guidance

Beyond classifier-based guidance, distributions can be composed through products, mixtures, and negation. Standard score-only diffusion parameterization supports products and negation in theory but not mixtures, motivating explicit probability or energy-based modeling for broader composition.

Products: Product composition re-normalizes multiple distributions and behaves like an intersection, assigning high probability where all component distributions do.Guidance can be viewed as a two-expert product combining data density with a conditional property model.
Mixtures: Mixture composition represents a union of concepts, but score functions alone cannot define mixtures because generation requires ratios of component normalizers.The experiments assume the ratio of normalizers is 1; another unknown ratio would produce a weighted mixture.
Negation: Negation inverts one density relative to another, assigning high likelihood to points favored by p_0(x) but not by p_1(x).The parameter α controls the degree of inversion, and the experiments use α = 0.5.
Nested composition: Negation can be nested with products and mixtures to construct complex combinations of distributions.These operators extend compositional generation beyond the two-expert product structure used in common guidance methods.
Limitation: Because standard diffusion models learn log-density gradients rather than explicit log-likelihoods, they cannot directly support mixture composition, unlike product and negation composition in theory.This limitation distinguishes score composition from probability-level composition.

4 Scaling Compositional generation with Diffusion Models

Standard reverse diffusion does not correctly sample composed distributions because the composed score at intermediate noise levels differs from the sum of component scores. Annealed MCMC sampling addresses this mismatch, while energy-based parameterization enables mixture composition and Metropolis-corrected samplers.

Improving Sampling with MCMC: Composing diffusion models by adding their score estimates produces incorrect samples under standard reverse diffusion.The required score is that of the diffused target product distribution, which is not generally the sum of the component scores for t > 0.
Improving Sampling with MCMC: Annealed MCMC applies transition kernels across intermediate distributions, beginning from a tractable distribution and ending at the composed target.The proposed procedure uses Langevin Dynamics and HMC-based kernels.
Improving Sampling with MCMC: MCMC sampling faithfully generates samples from composed distributions without modifying the underlying diffusion models.Quantitative results further suggest that sampler choice may explain prior failures in compositional generation.
Energy-Based Parameterization: Energy-based parameterization recovers an explicit unnormalized log-likelihood, enabling mixture composition in addition to other compositional forms.The parameterization defines ϵθ(x, t) = −∇xfθ(x, t), while the standard score parameterization does not explicitly model log-likelihood.
Energy-Based Parameterization: Metropolis-corrected MALA and HMC produce the best compositional generation results and support automated hyper-parameter tuning through acceptance rates.The energy-based parameterization is what makes these Metropolis corrections available.
Energy-Based Parameterization: On CLEVR, MCMC improves the accuracy of simultaneously controlling multiple cube positions, with further gains from the Metropolis adjustment.Training conditions on one object location per image, whereas composition imposes multiple positional constraints.

5 Experiments

The experiments evaluate MCMC-based samplers and energy-parameterized diffusion models across compositional density, shape, ImageNet, and text-to-image tasks. MCMC methods improve generation, while Metropolis corrections and energy parameterization provide further gains in several settings.

2D densities: MCMC methods improve sample quality and likelihood over reverse diffusion on illustrative compositional datasets.Metropolis-adjusted methods perform best under matched score-function evaluation costs.
Composing Cubes: More MCMC steps improve the accuracy of placing all five cubes in their correct locations, with larger gains from Metropolis adjustment.Accuracy is measured as the fraction of samples whose object locations are all correct.
Classifier conditioning: MCMC sampling improves classifier-guided ImageNet performance over reverse sampling, with further improvements from Metropolis corrections.The evaluation uses classifier accuracy, FID, and Inception Score on 128x128 ImageNet.
Text-2-Image: Composing text-conditioned diffusion models enables more faithful generations of scenes with multiple concepts.The approach trains an energy-parameterized model for 64x64 text-conditioned image generation and presents upsampled results.
Image tapestries: Spatially composing overlapping text-to-image models supports image tapestries with specified content at different locations and scales.The experiments illustrate controllable compositions across multiple scales.

6 Discussion

The proposed compositional diffusion approach improves generation across domains, scales, and compositional operators, but its advanced samplers impose substantial computational costs. The authors identify lower-cost methods as an important direction for future work.

Limitations: Advanced MCMC samplers can take 5-times longer to generate samples than typical diffusion sampling.The authors describe the higher sampling cost as a major drawback.
Limitations: Energy-parameterized models require a second backward pass, doubling memory and compute cost relative to score-parameterized models.The additional pass computes the derivative implicitly.
Future work: Distillation and easier-to-differentiate neural networks are proposed as possible ways to reduce the cost of the sampling procedures.These are identified as future directions rather than evaluated solutions in the paper.
Conclusion: The approach yields notable improvements across a variety of domains, scales, and compositional operators.The conclusion frames the work as demonstrating the feasibility of composing pretrained diffusion models to model new distributions.

A Detailed Derivation of Diffusion Models

Diffusion models define a forward Markov process with tractable Gaussian marginals and learn a stepwise reverse process for sampling. MCMC methods provide alternative sampling procedures, including ULA and HMC variants, whose step sizes and momentum handling affect implementation.

Forward diffusion: Diffusion models define latent variables through a forward Markov process that progressively transforms data toward Gaussian noise.The process has analytically computable Gaussian time marginals, and positive noise values make the limiting distribution N(0, I).
Reverse diffusion: A parameterized reverse model learns to invert the forward process step by step, with p(xT) set to N(0, I) for ancestral sampling.Training maximizes a variational bound, focusing on the parameter-dependent middle terms because the first term lacks parameters and the final term is approximately zero.
Reverse diffusion: The reverse model commonly predicts the noise added to xt rather than directly predicting the mean of xt−1.This noise-prediction parameterization determines the reverse transition used during sampling.
MCMC sampling: HMC augments the target with momentum, integrates Hamiltonian dynamics using leapfrog steps, and can retain part of the momentum between updates.U-HMC omits Metropolis acceptance because the parameterized models do not provide an explicit likelihood.
MCMC sampling: A diffusion reverse step at fixed noise level is equivalent to ULA sampling, enabling MCMC composition by repeating reverse updates before changing noise levels.Sampler step sizes are initialized proportional to βt; adjusted MALA and HMC variants target average acceptance rates of approximately 60% and 70%, respectively.

B.4 MCMC Implementation Details

The implementation shows that diffusion composition must account for both sampler behavior and the mismatch between diffusing modified distributions and modifying diffused distributions. MCMC methods exploit the resulting annealed sequence, but image sampling requires avoiding intermediate clipping.

Image-domain implementation: Intermediate clipping of diffusion samples to [-1, 1] caused uniform textures during image-domain MCMC sampling.Removing this clipping was important for proper MCMC sampling.
Image-domain implementation: Image-domain mixing improved when one reverse-process step initialized MCMC before multiple MCMC steps at each timestep.The initialization contracts a sample between noise levels and is close in form to Langevin sampling.
Composition mismatch: Products of diffused distributions are not generally the diffusion of the product distribution, so standard reverse diffusion does not sample the desired product model.The same issue applies to tempering and predictive model composition.
Composition mismatch: Scaling a data-distribution score at time zero does not remain valid for positive diffusion times, invalidating a direct tempered-score sampling interpretation.The equality holds at t=0 but not for t>0.
Annealed MCMC: Products, tempering, and guidance applied to diffused distributions form an annealed sequence from N(0, I) to the composed target, which supports annealed MCMC sampling.This permits sampling composed models directly with available score estimates.

D Energy-Based Parameterizations

The paper introduces energy-based parameterizations that expose an explicit scalar potential for diffusion models. This enables Metropolis-corrected sampling and additional compositional operators, though parameterization choices involve computational and modeling trade-offs.

Benefits and limitations: Unadjusted samplers such as ULA and U-HMC can work with score estimates but do not produce exact samples at non-zero step sizes.Metropolis corrections can make approximate samplers exact, but their acceptance probabilities require evaluating log pθ(xt).
Energy parameterization: The proposed energy parameterization uses a scalar neural network fθ(x, t) and defines the diffusion noise prediction as its input gradient.The resulting score estimate remains approximately −ϵθ(x,t), as in standard diffusion models.
Benefits and limitations: An explicit energy function enables Metropolis corrections and more compositional operators than are available from score estimates alone.With only score estimates, composition is limited to products, while Metropolis corrections require evaluating likelihood-related quantities.
Parameterization comparison: The paper compares standard, negative-L2, and inner-product energy parameterizations for unconditional ImageNet generation using FID.L2 and inner-product perform best among the alternatives, but both are outperformed by the standard parameterization; L2 was more stable for compositional sampling.
Parameterization comparison: The negative-L2 parameterization was considered more stable for compositional sampling because its energy function is bounded above.This stability observation is distinct from unconditional ImageNet FID performance.

E Synthetic Distribution Compositions

Synthetic experiments show that HMC improves composition for mixtures and negations of diffusion models. The approach nevertheless fails when the component distributions are disjoint.

Mixture: HMC sampling enables more accurate mixtures of different synthetic distributions.Additional two-dimensional mixture examples are provided in Figure A1.
Negation: HMC sampling enables accurate negations of different synthetic distributions.Additional two-dimensional negation examples are provided in Figure A3.
Failure cases: The proposed composition approach fails to generate the product of two distributions when those distributions are disjoint.This is presented as a failure case in Figure A3.

F Experimental Details

The experiments span synthetic distributions, CLEVR, ImageNet, and text-to-image settings, using score- and energy-based diffusion models with MCMC sampling. Evaluation combines distributional metrics, classifier-based accuracy, and qualitative failure-case analysis.

Experimental Scope: Experiments cover 2D synthetic datasets, CLEVR, ImageNet, and text-to-image settings with detailed metrics, training details, and architectures.The reported settings include synthetic distributions, cube-position accuracy, ImageNet classifier accuracy, Inception Score, and FID.
Failure Cases: Reverse diffusion fails to sample accurately from negated distributions, while the approach fails for products of distributions with no overlap.These failure cases are illustrated for component distributions, reverse diffusion, and HMC sampling.
Synthetic Datasets: Synthetic experiments use score- and energy-based models trained with residual MLPs, 100 timesteps, and fixed-step-size MALA and HMC samplers.The synthetic models use four residual blocks and hidden dimension 128; MALA uses 10 steps per timestep and HMC uses three steps with three leapfrog steps each.
CLEVR: CLEVR experiments generate 64 × 64 images containing 1 to 5 cubes and evaluate whether classifiers correctly identify cubes at specified positions.A cube is counted as correctly generated when the binary classifier confidence exceeds 0.5.
ImageNet: ImageNet experiments train unconditional 128 × 128 diffusion models and report classifier accuracy, Inception Score, and FID on 50,000 generated samples.Sampling is initialized with one reverse-process step before MCMC; the implementation uses 1,000 timesteps and a cosine beta schedule.
Text-to-Image: Text-to-image experiments train an energy-based model on 400 million images and use HMC sampling, while image tapestries combine overlapping models across a canvas.The text-to-image setup uses 1,000 timesteps; tapestry models are placed every 32 pixels with 50% overlap and seam-hiding edge tapering.

G Text-to-Image Results

The text-to-image results show that composing energy-parameterized diffusion models improves control over image details, object colors, confounding factors, and object counts. Negating one composed model also supports unusual configurations such as a sandy beach without coastal water.

Compositional Rendering: Composing energy-parameterized diffusion models enables more accurate rendering of detailed caption information, object colors, and scene object counts.Examples include details such as Central Park or snow, correct colors, and accurate numbers of objects.
Negation: Negating one energy-parameterized diffusion model while composing another enables images with uncommon configurations, such as a sandy beach without coastal water.The composition removes a commonly occurring confounding factor.
Prompt Examples: The examples evaluate compositions involving a red car and pink forest, a lake with purple trees, and multiple oranges.The displayed prompts include conjunctions such as “AND pink forest” and “Two Oranges AND Two of the”.