PaperlayerView exact revision on arXiv

Source-linked AI summary

Training Diffusion Models with Reinforcement Learning

Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, Sergey Levine

arXiv:2305.13301v4cs.LGcs.AIcs.CV

TL;DR

Diffusion models are usually trained toward likelihood objectives, although many applications instead require directly optimizing downstream goals. This paper introduces DDPO, which frames denoising as multi-step decision-making and uses policy gradients to adapt diffusion models to objectives such as compressibility, aesthetic quality, and prompt-image alignment using VLM feedback.

  • Problem

    Many diffusion-model applications target downstream objectives rather than likelihoods, but directly optimizing those objectives is challenging because exact full-process likelihoods are intractable.

  • Method

    DDPO frames iterative denoising as a multi-step decision-making problem and applies policy-gradient reinforcement learning, including VLM-derived rewards for prompt-image alignment.

  • Results

    DDPO effectively adapts text-to-image diffusion models to objectives difficult to specify through prompting or evaluate programmatically, and improves alignment without additional human annotation.

  • Takeaways & Limitations

    The framework enables diffusion models to optimize user-specified objectives, including image properties and semantic alignment, through direct reinforcement learning and automated VLM feedback.

  • Takeaways & Limitations

    Reward-weighted regression is not theoretically justified for diffusion training because the denoising loss is a variational bound rather than an exact log-likelihood.

Abstract

from arXiv · show

Diffusion models are a class of flexible generative models trained with an approximation to the log-likelihood objective. However, most use cases of diffusion models are not concerned with likelihoods, but instead with downstream objectives such as human-perceived image quality or drug effectiveness. In this paper, we investigate reinforcement learning methods for directly optimizing diffusion models for such objectives. We describe how posing denoising as a multi-step decision-making problem enables a class of policy gradient algorithms, which we refer to as denoising diffusion policy optimization (DDPO), that are more effective than alternative reward-weighted likelihood approaches. Empirically, DDPO is able to adapt text-to-image diffusion models to objectives that are difficult to express via prompting, such as image compressibility, and those derived from human feedback, such as aesthetic quality. Finally, we show that DDPO can improve prompt-image alignment using feedback from a vision-language model without the need for additional data collection or human annotation. The project's website can be found at http://rl-diffusion.github.io .

1 INTRODUCTION

The paper proposes DDPO, a policy-gradient method that treats diffusion denoising as a multi-step decision process to optimize downstream objectives directly. It applies this approach to text-to-image adaptation, including objectives difficult to specify through prompting or evaluate without human feedback.

  • Problem and approach: DDPO frames denoising as a multi-step decision-making problem and uses exact per-step likelihoods to optimize black-box downstream rewards.This avoids relying on the intractable likelihood of the complete denoising process.
  • Applications: The method targets objectives such as compressibility, aesthetic quality, and prompt-image alignment determined by vision-language models.
  • Applications: DDPO is applied to finetune large text-to-image diffusion models for objectives that are difficult to specify through prompting or derive from ordinary data.
  • Applications: Vision-language-model feedback replaces additional human annotations for improving prompt-image alignment on unusual subject-setting compositions.
  • Contributions: The paper compares DDPO with reward-weighted likelihood methods and demonstrates generalization to unseen prompts.

2 RELATED WORK

The related work spans diffusion modeling, controllable generation, reinforcement learning from human feedback, and sequential decision-making formulations for image generation. The paper differs by targeting downstream objectives rather than primarily matching a data distribution.

  • Diffusion models: Diffusion models have been used for images, videos, 3D shapes, and robotic trajectories, with denoising objectives conventionally derived as approximations to likelihood.
  • Controllable generation: Controllable-generation research includes finetuning on user data, text-embedding optimization, model composition, adapters, and inference-time techniques.
  • Reinforcement learning from human feedback: Reinforcement learning from human feedback has been studied across robotics, games, translation, retrieval, question answering, summarization, and instruction following.
  • Sequential decision-making: Prior sequential-decision approaches to diffusion training include methods for image generation and policy gradients, but a concurrent method focused on data-distribution matching rather than downstream objectives.

3 PRELIMINARIES

This section introduces conditional diffusion models, their iterative forward-noising and reverse-denoising processes, and the Markov decision-process formulation used to describe sequential reinforcement learning.

  • Diffusion models: Conditional diffusion models represent p(x_0|c) and learn a reverse process for a Markovian forward process that progressively adds noise.
  • Diffusion models: The reverse process is trained with a neural network μ_θ(x_t,c,t) using an objective justified as a variational lower bound on data log-likelihood.
  • Diffusion models: Sampling starts from x_T drawn from N(0,I) and follows reverse transitions to produce a trajectory ending at x_0.
  • Diffusion models: Popular samplers use isotropic Gaussian reverse transitions with fixed timestep-dependent variance.
  • MDPs and reinforcement learning: An MDP consists of states, actions, an initial-state distribution, transition kernel, and reward function governing sequential decisions.
  • MDPs and reinforcement learning: RL policies generate trajectories of states and actions, with the objective of maximizing expected cumulative reward.

4 REINFORCEMENT LEARNING TRAINING OF DIFFUSION MODELS

The paper frames diffusion denoising as reinforcement-learning decision making and introduces DDPO to optimize sample rewards directly. This multi-step formulation enables exact policy-gradient estimation, whereas reward-weighted regression only approximately optimizes the objective.

  • Reward-weighted regression: Reward-weighted regression reuses the denoising loss with model-generated samples and reward-dependent weights.Exponentiated rewards provide one weighting scheme, while binary weights retain samples above a reward threshold.
  • Reward-weighted regression: RWR is only an approximate optimization method because diffusion denoising loss is a variational bound rather than an exact log-likelihood.Its one-step formulation also ignores the sequential nature of denoising and uses only final samples.
  • Denoising diffusion policy optimization: DDPO reframes denoising as a multi-step MDP, making the policy an isotropic Gaussian whose exact log-likelihoods and gradients are tractable.This avoids the arbitrarily complicated final-sample distribution used in the RWR formulation.
  • Denoising as a multi-step MDP: Diffusion training can be treated as an MDP in which denoising trajectories receive the final sample reward r(x0, c).The policy maps each denoising state to a transition distribution, and maximizing the diffusion reward objective is equivalent to maximizing the corresponding RL objective.
  • Policy gradient estimation: DDPO alternates denoising-trajectory collection with parameter updates using policy-gradient estimators, including the score-function variant DDPOSF.Importance sampling supports multiple updates from trajectories generated by older parameters, with clipping used as a trust region.

5 REWARD FUNCTIONS FOR TEXT-TO-IMAGE DIFFUSION

The paper evaluates reward functions for text-to-image diffusion that are difficult to express through prompting, including compressibility, aesthetic quality, and prompt-image alignment. A vision-language model supplies alignment feedback without requiring additional human annotation.

  • Reward-function spectrum: Text-to-image diffusion is used to study reward functions ranging from straightforward computations to complex workflows involving large vision-language models.The setting provides pretrained models and diverse reward functions for evaluating reinforcement-learning methods.
  • Compressibility and incompressibility: File-size rewards test objectives unavailable through prompting because images are rarely captioned with their file size.With samples fixed at 512x512, JPEG file size depends solely on image compressibility; the paper studies both minimizing and maximizing it.
  • Aesthetic quality: Aesthetic quality is measured with the LAION aesthetics predictor trained on 176,000 human image ratings.The predictor is a linear model over CLIP embeddings, with annotations ranging from 1 to 10.
  • Automated prompt alignment: Prompt-image alignment is optimized using an existing vision-language model instead of additional large-scale human labeling.LLaVA describes each generated image, and the reward is semantic similarity between that description and the prompt measured by BERTScore recall.
  • Automated prompt alignment: The VLM interface can use questions to specify complex or hard-to-define reward functions, with reward complexity limited by the vision-language models’ capabilities.The paper illustrates this with the question “what is happening in this image?” and notes that questions could be generated automatically.

6 EXPERIMENTAL EVALUATION

The experiments test DDPO against reward-weighted likelihood methods across several objectives, including compressibility, aesthetic quality, prompt alignment, and generalization. DDPO improves targeted rewards, supports VLM-based alignment without additional human labels, and transfers beyond the finetuning prompts.

  • Experimental setup: The evaluation compares DDPO variants with RWR on compressibility, incompressibility, and aesthetic quality, then examines VLM-based alignment and prompt generalization.Stable Diffusion v1.4 provides the base model, with reward functions and evaluation questions spanning algorithm comparisons, automated rewards, and transfer.
  • Reward optimization: DDPO transforms naturalistic images into artistic drawings for aesthetic quality, removes backgrounds and smooths foregrounds for compressibility, and finds compression-resistant artifacts for incompressibility.These reward-optimizing strategies emerge without further data curation.
  • Algorithm comparisons: DDPO attains higher rewards than RWR across all three tasks, while importance sampling slightly outperforms the score-function estimator within DDPO.RWR weighting schemes perform comparably, making sparse weighting preferable here because of its simplicity and lower resource requirements.
  • Automated prompt alignment: VLM feedback improves prompt-image alignment without additional human labels, including prompts for which the pretrained model initially has zero success rate.Transfer across prompts supplies reward signals that allow difficult compositions such as “a dolphin riding a bike” to improve.
  • Prompt alignment: Samples become more cartoon-like or artistic during alignment finetuning, although this behavior was not directly optimized and is attributed hypothetically to pretraining or reward-function effects.The caption notes that thick lines show average activity scores and faint lines show selected individual-prompt scores.
  • Generalization: Finetuning effects generalize from 45 animals and 3 activities to unseen animals, non-animal objects, and novel activities such as “taking an exam”.Figure 6 illustrates this transfer with prompts including “a capybara washing dishes” and “a duck taking an exam”.

7 DISCUSSION AND LIMITATIONS

The discussion presents DDPO as a framework for directly optimizing diffusion models with diverse rewards, while identifying restricted prompt diversity and unexamined overoptimization as limitations. The broader-impact discussion notes both beneficial uses and potential misuse.

  • Discussion: DDPO frames iterative denoising as multi-step decision-making and provides policy-gradient algorithms for directly optimizing diverse reward functions.The framework addresses objectives difficult to specify through prompts or evaluate programmatically, including compressibility and semantic prompt alignment.
  • Discussion: VLM feedback offers an automated way to derive rewards for image quality without relying on additional human annotation.The paper presents this as a route toward more targeted training of large generative models.
  • Limitations: The experiments cover a constrained range of images, largely involving animals performing activities, and do not study overoptimization or divergence from the original distribution.The authors identify broader VLM questions and more diverse prompt distributions as future directions.
  • Broader impacts: Generative-model adaptation may support educational applications but may also enable disinformation, impersonation, or phishing, making reliable synthetic-content detection important.These risks are stated as broader impacts rather than as evaluated outcomes of the experiments.

APPENDIX A OVEROPTIMIZATION

The appendix examines reward overoptimization and compares DDPO with DPOK and universal guidance. It finds that optimizing learned or proxy rewards can degrade semantic quality, although DDPO outperforms DPOK in the reported comparison.

  • Overoptimization: Reward overoptimization can move diffusion models away from useful outputs, producing high reward while drifting from the pretraining distribution.The paper frames this as a general limitation of finetuning on reward functions, especially learned ones.
  • Overoptimization: Optimizing incompressibility eventually produces high-frequency noise, while optimizing animal-count prompts can trigger typographic attacks instead of correct animal counts.The VLM may interpret generated text as the requested number rather than assessing the generated animals.
  • Mitigation: There is no general-purpose method for preventing overoptimization, so the authors manually select the last checkpoint before deterioration for qualitative results.They note that KL regularization may be empirically equivalent to early stopping and identify prevention as future work.
  • Comparison to Classifier Guidance: Universal guidance improves aesthetic score significantly, but the improvement is small compared with DDPO and generation is much slower.The universal-guidance results average 50 samples for one prompt; one image takes almost 2 minutes on an NVIDIA A100 GPU.
  • Comparison to DPOK: DDPOIS outperforms DPOK across the reported comparison without KL regularization.The comparison uses ImageReward for training and evaluates with both ImageReward and the LAION aesthetic score.
  • Comparison to DPOK: Significant overoptimization appears within 25k reward queries for one count prompt, but the authors describe it as neither severe nor unreasonably fast.At 20k reward queries, qualitative samples remain high quality.

E.1 CFG TRAINING

The appendix explains CFG training and interleaving choices for reinforcement-learning finetuning of diffusion models. CFG training becomes essential across repeated RWR rounds, while DDPO ultimately exceeds RWR across interleaving levels.

  • CFG training: Classifier-free guidance jointly trains conditional and unconditional objectives, then mixes their predictions during sampling with a guidance weight.The guided epsilon prediction is used to compute the next denoised sample.
  • CFG training: Training only the conditional objective causes rapid performance deterioration after the first finetuning round.The authors attribute this to guidance-weight miscalibration as the model is updated.
  • CFG training: CFG training uses a fixed guidance weight and the guided epsilon prediction during both training and sampling.The procedure has no effect after one round but becomes essential for subsequent RWR rounds.
  • Interleaving: DDPO uses a policy-gradient objective and more on-policy data, collecting 256 samples per iteration compared with 10,000 for RWR.RWR is designed for off-policy data, whereas policy-gradient methods specifically require on-policy data.
  • Interleaving: More frequent RWR interleaving helps up to a point, after which performance degrades, and RWR still cannot match DDPO asymptotically.The ablation varies sampling and training rounds to isolate the effect of data distribution.

APPENDIX F QUANTITATIVE RESULTS FOR GENERALIZATION

The appendix quantitatively evaluates whether aesthetic-quality finetuning generalizes beyond its training prompts. Generalization is strong for unseen animals and weaker for ordinary objects.

  • Quantitative generalization: The analysis measures average aesthetic reward throughout training across several prompt distributions.The finetuning prompts contain 45 common animals, while evaluation includes 38 unseen animals and 50 ordinary objects.
  • Quantitative generalization: The aesthetic-quality objective generalizes very well to unseen animals.The unseen-animal set contains animal prompts not used during finetuning.
  • Quantitative generalization: Generalization to everyday objects is positive but weaker than generalization to unseen animals.The ordinary-object set includes items such as toasters, chairs, and coffee cups.
  • Quantitative generalization: The appendix supplements prior qualitative evidence that both aesthetic-quality and image-prompt-alignment models generalize to unseen prompts.Additional figures show reward curves and qualitative samples for the evaluated prompt distributions.

APPENDIX G MORE SAMPLES

This appendix provides additional qualitative samples for baseline RWR, DDPO image-prompt alignment, and generalization of aesthetic and alignment objectives. The examples span seen and unseen prompts with unusual subjects, activities, and settings.

  • Baseline RWR: Figure 13 shows qualitative samples generated by the baseline RWR method.
  • Image-prompt alignment: Figure 14 shows additional DDPO samples on seen prompts using the image-prompt alignment reward.
  • Aesthetic quality generalization: Figure 15 shows aesthetic-quality generalization to unseen animals and everyday objects.Examples include a capybara washing dishes and a parrot driving a car.
  • Image-prompt alignment generalization: Figure 16 shows image-prompt-alignment generalization to unseen subjects and activities.Examples include a duck taking an exam, a horse typing on a keyboard, and a tree riding a bike.
Loading 2305.13301v4…