PaperlayerView exact revision on arXiv

Source-linked AI summary

Contrastive Energy Prediction for Exact Energy-Guided Diffusion Sampling in Offline Reinforcement Learning

Cheng Lu, Huayu Chen, Jianfei Chen, Hang Su, Chongxuan Li, Jun Zhu

arXiv:2304.12824v2cs.LG

TL;DR

Diffusion guidance with unnormalized energy functions requires intermediate guidance that is difficult to estimate and whose exactness previous methods do not establish. The paper derives the exact guidance and learns it with contrastive energy prediction, which is theoretically exact under unlimited capacity and data. Applications show improved offline-RL results and image-synthesis performance roughly on par with classifier guidance.

  • Problem

    Intermediate guidance for diffusion sampling from an energy-modified, generally unnormalized target distribution is unknown and difficult to estimate, while prior methods do not guarantee the desired final distribution.

  • Method

    CEP learns exact intermediate energy guidance by contrasting energy values across noise-perturbed samples, using self-normalized labels for improved numerical stability.

  • Results

    CEP outperforms existing guidance methods in offline RL and achieves image-synthesis results roughly on par with classic classifier guidance.

  • Takeaways & Limitations

    The method provides an exact energy-guided diffusion approach with demonstrated effectiveness across offline RL and image synthesis.

  • Takeaways & Limitations

    Previous MSE-based energy guidance is exact only at t = 0 and is inexact for t > 0.

Abstract

from arXiv · show

Guided sampling is a vital approach for applying diffusion models in real-world tasks that embeds human-defined guidance during the sampling procedure. This paper considers a general setting where the guidance is defined by an (unnormalized) energy function. The main challenge for this setting is that the intermediate guidance during the diffusion sampling procedure, which is jointly defined by the sampling distribution and the energy function, is unknown and is hard to estimate. To address this challenge, we propose an exact formulation of the intermediate guidance as well as a novel training objective named contrastive energy prediction (CEP) to learn the exact guidance. Our method is guaranteed to converge to the exact guidance under unlimited model capacity and data samples, while previous methods can not. We demonstrate the effectiveness of our method by applying it to offline reinforcement learning (RL). Extensive experiments on D4RL benchmarks demonstrate that our method outperforms existing state-of-the-art algorithms. We also provide some examples of applying CEP for image synthesis to demonstrate the scalability of CEP on high-dimensional data.

1. Introduction

The paper studies diffusion sampling guided by flexible scalar energy functions, where the correct intermediate guidance is unknown. It derives exact guidance and introduces CEP, with applications showing strong offline-RL performance and image-synthesis results comparable to classifier guidance.

  • Scalar energy functions let diffusion models encode preferences such as rewards, Q-functions, human feedback, feature similarity, or frame distances.
  • The target distribution combines the data distribution q(x) with an energy factor e^-βE(x), while remaining generally unnormalized.
  • Previous energy-guided samplers define guidance manually or arbitrarily across diffusion time, without establishing that final samples follow the desired distribution.
  • CEP estimates exact intermediate guidance by comparing energy values across noise-perturbed samples, and is theoretically exact with unlimited capacity and data.
  • Experiments show CEP outperforms existing algorithms on most D4RL tasks, especially AntMaze, while image-synthesis quality is roughly on par with classifier guidance.

2. Background

The background connects diffusion models’ reverse-time score-based sampling with offline RL’s constrained policy optimization. This connection motivates treating optimal-policy sampling as energy-guided diffusion from a behavior model.

  • 2.1. Diffusion (Probabilistic) Models: Diffusion models gradually add Gaussian noise to data, then recover clean samples through a reverse process initialized from an approximately Gaussian terminal distribution.
  • 2.1. Diffusion (Probabilistic) Models: The reverse diffusion process depends on score functions, which neural networks estimate from noisy data during training.
  • 2.2. Constrained Policy Optimization in Offline Reinforcement Learning: Offline RL uses a fixed dataset of state-action histories from a behavior policy to optimize a new policy without further environment interaction.
  • 2.2. Constrained Policy Optimization in Offline Reinforcement Learning: Constrained policy optimization combines policy improvement through Q-values with a policy constraint relative to behavior data.
  • 2.2. Constrained Policy Optimization in Offline Reinforcement Learning: The optimal constrained policy belongs to the energy-guided distribution family, motivating sampling it with a pretrained diffusion behavior model.

3. Exact Energy-Guided Sampling

The paper derives an exact intermediate-energy representation for diffusion sampling under an energy-modified target distribution. CEP learns the corresponding guidance from noisy sample comparisons, with theoretical exactness and stabilization through self-normalized labels.

  • 3.1. Exact Formulation of Intermediate Energy Guidance: The desired target is represented as a product of the diffused data distribution and an exponential intermediate-energy term.
  • 3.1. Exact Formulation of Intermediate Energy Guidance: Intermediate energy is determined by the original data-space energy, while both target and data marginals approach the same Gaussian terminal distribution.
  • 3.1. Exact Formulation of Intermediate Energy Guidance: Exact guided sampling requires the pretrained data score together with the gradient of intermediate energy, whose direct estimation is generally difficult.
  • 3.2. Learning Energy Guidance by Contrastive Energy Prediction: CEP trains an energy model by contrasting K noise-perturbed samples and matching predicted labels to soft labels derived from their energies.
  • 3.2. Learning Energy Guidance by Contrastive Energy Prediction: With unlimited capacity and data, the optimal CEP model has a gradient equal to the exact intermediate-energy guidance.
  • 3.2. Learning Energy Guidance by Contrastive Energy Prediction: Self-normalized energy labels improve numerical stability, and their bias decreases as K increases, becoming equivalent to the original objective as K approaches infinity.

4. Comparison with Previous Methods for Guided Sampling

The paper finds that prior energy-guided samplers are inexact, while CEP provides a contrastive alternative to classifier guidance for conditional sampling and extends to general energy guidance.

  • Previous Energy-Guided Samplers: Previous energy-guided samplers are inexact and do not guarantee convergence to the desired distribution.
  • MSE for Predicting Energy: MSE-trained energy models match the true energy only at t = 0 and are inexact for all t > 0.
  • CEP and Classifier Guidance: Classifier guidance normalizes across conditions, whereas CEP compares data samples for a specified condition and normalizes across samples.
  • CEP and Classifier Guidance: CEP and classifier guidance can both guarantee exact guidance with unlimited data and model capacity.

5. Q-Guided Policy Optimization for Offline Reinforcement Learning

QGPO formulates offline constrained policy optimization as energy-guided diffusion sampling, using CEP to estimate Q-dependent intermediate guidance. It trains support-action-based objectives and outperforms referenced baselines in most D4RL tasks, especially difficult AntMaze tasks.

  • 5.1. Problem Formulation: QGPO expresses constrained policy optimization as energy-guided sampling, with the desired guidance determined by an action evaluation model.The method estimates the targeted policy score using a behavior diffusion model, a Q-function, and an energy-guidance model.
  • 5.2. In-Support Contrastive Energy Prediction: CEP trains the guidance model from behavior-model-generated support actions because the offline dataset lacks multiple independent actions for each continuous state.The resulting objective is called in-support CEP because it is optimized over the generated support action set rather than the true conditional dataset distribution.
  • 5.4. Results: D4RL evaluation reports normalized performance means and standard deviations across five random seeds at the end of training.Numbers within 5 percent of the maximum in each task are highlighted, and Diffusion-QL is rerun for consistent evaluation.
  • 5.3. Q-Function Training: The Q-function can use Bellman-style targets, but importance sampling with the support action set avoids repeated diffusion-model sampling during training.The paper identifies diffusion sampling inside the Bellman target as time-consuming and proposes support actions to reduce that cost.
  • 5.4. Results: QGPO outperforms referenced baselines in most D4RL tasks, particularly difficult tasks such as Antmaze-Large.The comparison includes traditional Gaussian-like-policy methods and recent diffusion-based offline RL algorithms.

6. Image Synthesis Examples with CEP

CEP supports both class-conditional and continuous energy-guided image synthesis. On ImageNet it performs roughly on par with classifier guidance, while hue-based energies control sampled image colors as guidance scales vary.

  • 6.1. Results in Class-Conditional Image Synthesis: On ImageNet, CEP achieves image sample quality roughly on par with classic classifier guidance.The comparison uses the same architecture, training pipeline, and evaluation methods without hyperparameter tuning; fixed-seed samples are almost visually identical.
  • 6.2. Energy-Guided Image Synthesis: A hue-based energy function enables continuous control of an image’s overall color appearance instead of discrete class conditioning.Hue is represented angularly, with red at 0, green at 2π/3, and blue at 4π/3.
  • 6.2. Energy-Guided Image Synthesis: Three energy-guidance models and adjustable guidance scales control whether generated images emphasize red, green, or blue.Figure 3 shows samples for the three color guidances while holding the random seed fixed.

7. Conclusion

The paper derives exact energy guidance for diffusion sampling and introduces CEP to estimate it. Applications show improved offline RL performance and image-synthesis results roughly matching classifier guidance.

  • 7. Conclusion: The paper formulates energy-edited diffusion sampling, derives exact guidance, and trains an estimator with contrastive energy prediction.CEP is described as converging to the desired guidance under unlimited model capacity and data.
  • 7. Conclusion: Experiments report that the method outperforms existing guidance methods in offline RL and is roughly on par with classifier guidance for conditional image synthesis.These applications demonstrate the method’s effectiveness and scalability within the reported settings.

A. Limitations and broader impacts

The paper identifies potential broader-impact risks from energy-guided diffusion sampling, while the supplied technical passages establish CEP’s theoretical convergence conditions and proof relation to exact guidance.

  • A. Limitations and broader impacts: Energy-guided diffusion sampling could generate harmful content such as deepfakes and amplify unwanted social bias in training data.The paper presents these as potential risks shared with other deep generative modeling methods.
  • D.4. Proof of Theorem 3.2: The proof establishes that the optimal model differs from the intermediate energy by a time-dependent constant, so their gradients agree.The constant is independent of the perturbed input, yielding ∇_x_t f_ϕ*(x_t,t) = ∇_x_t E_t(x_t).

D.5. Proof of Theorem D.1

The proof establishes that the learned intermediate energy model recovers the exact energy guidance under unlimited data and model capacity. The argument shows the optimal model differs from the intermediate energy by an input-independent constant, so their gradients coincide.

  • The proof treats the intermediate time variables and noisy samples jointly to characterize the optimal energy model.
  • For every fixed time sequence and sample sequence, the optimization equality holds pointwise under the stated support conditions.
  • With unlimited data and model capacity, the optimal model satisfies the derived characterization.
  • The optimal model differs from the intermediate energy by a constant independent of the noisy input.
  • Because the constant has zero input gradient, the learned model’s gradient equals the exact intermediate energy-guidance gradient.

E. Comparison with Existing Energy-Guided Sampling Algorithms

The section compares exact CEP guidance with MSE- and DPS-based approximations, relating their errors to nonlinear expectation exchanges and finite-time assumptions. It also connects CEP to contrastive learning and discusses guidance-scale tuning and E-MSE.

  • Comparison with Existing Energy-Guided Sampling Algorithms: Exchanging nonlinear functions with expectations introduces approximation errors; the paper associates larger errors with DPS because E0 may be complex.
  • Comparison with Existing Energy-Guided Sampling Algorithms: MSE guidance is a first-order approximation that assumes the diffused and original energies are similar, an assumption most appropriate near t = 0.
  • Comparison with Existing Energy-Guided Sampling Algorithms: DPS introduces another approximation by replacing the original sample with its conditional expectation, which is also mainly justified near t = 0.
  • Comparison with Existing Energy-Guided Sampling Algorithms: Across additional 2-D and offline-RL experiments, CEP performs significantly better than MSE, which performs significantly better than DPS.
  • Relationship with Contrastive Learning: For β = 1 and E(x0) = −log q0(c|x0), CEP is equivalent to a traditional contrastive InfoNCE objective when summed over conditions.
  • Relationship with Contrastive Learning: CEP generalizes the contrastive objective because it uses q0(x0) samples when general energy-guided sampling cannot generally provide samples from the target distribution.
  • Relationship between Inverse Temperature and Guidance Scale: Changing guidance scale s linearly changes guidance strength, whereas changing β alters guidance through the exponential energy term.
  • Relationship between Inverse Temperature and Guidance Scale: In complex image and reinforcement-learning tasks, increasing s can be necessary because varying β alone may not provide good performance.

I.2. Experiment Details of QGPO

QGPO’s offline-RL implementation trains separate diffusion, Q-value, and energy-guidance networks, then evaluates guided policies with DPM-Solver. Experiments use repeated trials, ablations, and comparisons with diffusion and resampling baselines.

  • Model Training: QGPO trains three neural networks per task: a diffusion behavior model, an action evaluation model, and an energy guidance model.
  • Model Training: The behavior model follows prior diffusion-policy architecture and is trained for 600k gradient steps with Adam at learning rate 1e-4.
  • Model Training: The action evaluation model is a 3-layer, 256-hidden-unit ReLU MLP trained for 500k steps; βQ and K differ between MuJoCo and AntMaze tasks.
  • Model Training: The energy guidance model is a 4-layer, 256-hidden-unit SiLU MLP trained for 1M steps, with K task-dependent and β fixed at 3.
  • Evaluation: Evaluation averages five independent trials and multiple test seeds, using 15 diffusion steps with second-order DPM-Solver.
  • Evaluation: Ten evaluation diffusion steps perform equally well as 15, while five steps slightly underperform 15; support actions remain generated with 15 steps.
  • Ablations: Ablations examine gradient scales, diffusion-step counts, guidance methods, and resampling-based alternatives.
  • Additional Experiments: The low-dimensional experiments train one generative diffusion model and one energy guidance model on 1M two-dimensional datapoints.

N. More Results for Energy-Guided Image Synthesis

The image-synthesis results vary color guidance strength for conditional and unconditional diffusion priors. The figures show samples across guidance scales from 0.0 to 10.0, with fixed seeds in the conditional comparison.

  • Conditional-prior color guidance is visualized across guidance scales from 0.0 to 10.0.
  • Unconditional-prior color guidance is likewise visualized across guidance scales from 0.0 to 10.0.
  • Conditional-prior samples are compared across different guidance methods with random seeds fixed across experiments.The evaluated ImageNet classes include animals, objects, food, and scenes identified by class number.
Loading 2304.12824v2…