PaperlayerView exact revision on arXiv

Source-linked AI summary

Directly Fine-Tuning Diffusion Models on Differentiable Rewards

Kevin Clark, Paul Vicol, Kevin Swersky, David J Fleet

arXiv:2309.17400v2cs.CVcs.LG

TL;DR

Diffusion models often need fine-tuning to optimize desired rewards, while existing reward-based methods can be sample inefficient. DRaFT backpropagates differentiable reward gradients through diffusion sampling and introduces truncated and lower-variance variants. Across diverse rewards and prompts, the methods substantially improve efficiency and image quality, though reward hacking can reduce diversity.

  • Problem

    Existing reward fine-tuning methods improve alignment or aesthetics on narrow domains but are sample inefficient and have not broadly improved generation style across diverse prompts.

  • Method

    DRaFT differentiates through the diffusion sampling process, with DRaFT-K truncating backpropagation to the last K steps and DRaFT-LV reducing gradient variance for K = 1.

  • Results

    >200× faster than the RL algorithms from Black et al. (2023) for maximizing LAION Aesthetics Classifier scores, while DRaFT-LV learns roughly 2× faster than ReFL.

  • Takeaways & Limitations

    DRaFT works across diverse differentiable rewards and substantially improves Stable Diffusion 1.4 on human preference rewards and aesthetic quality.

  • Takeaways & Limitations

    Reward hacking can cause the fine-tuned model to lose diversity and collapse toward a certain high-reward image.

Abstract

from arXiv · show

We present Direct Reward Fine-Tuning (DRaFT), a simple and effective method for fine-tuning diffusion models to maximize differentiable reward functions, such as scores from human preference models. We first show that it is possible to backpropagate the reward function gradient through the full sampling procedure, and that doing so achieves strong performance on a variety of rewards, outperforming reinforcement learning-based approaches. We then propose more efficient variants of DRaFT: DRaFT-K, which truncates backpropagation to only the last K steps of sampling, and DRaFT-LV, which obtains lower-variance gradient estimates for the case when K=1. We show that our methods work well for a variety of reward functions and can be used to substantially improve the aesthetic quality of images generated by Stable Diffusion 1.4. Finally, we draw connections between our approach and prior work, providing a unifying perspective on the design space of gradient-based fine-tuning algorithms.

1 INTRODUCTION

Diffusion models can generate high-quality data but may not match desired behaviors such as aesthetic quality, while existing reward fine-tuning methods are sample inefficient. DRaFT addresses this by backpropagating differentiable rewards through sampling, with truncated and lower-variance variants improving efficiency.

  • Diffusion models model training distributions well, but those distributions may not align with desired behaviors such as aesthetically pleasing image generation.Classifier-free guidance and curated datasets are commonly used to improve alignment and image quality.
  • Existing supervised and reinforcement learning approaches improve alignment or aesthetics on narrow domains but remain sample inefficient and do not broadly improve style across diverse prompts.
  • DRaFT backpropagates reward gradients through the full diffusion sampling chain, using gradient checkpointing and LoRA weights to reduce memory and compute costs.
  • DRaFT-K truncates backpropagation to the last K sampling steps, while DRaFT-LV averages multiple noise samples to obtain lower-variance gradients for K = 1.
  • >200× faster than the RL algorithms from Black et al. (2023) for maximizing LAION Aesthetics Classifier scores, while DRaFT-LV learns roughly 2× faster than ReFL.The methods also improve Stable Diffusion on PickScore and Human Preference Score v2 and support multiple reward functions.

2 RELATED WORK

Prior work aligns diffusion models using guidance, supervised reward fine-tuning, reinforcement learning, or reward gradients. DRaFT differs by directly optimizing final-image rewards through the sampling process and reports greater efficiency than ReFL.

  • Learning human preferences: Human preference learning trains reward models from judgments and then learns policies that maximize those rewards.
  • Guidance: Guidance steers diffusion sampling toward objectives using auxiliary signals from models such as classifiers, facial recognition systems, and object detectors.
  • Backpropagation through diffusion sampling: Some methods backpropagate through diffusion sampling to optimize sampler hyperparameters, accelerate sampling, or train diffusion policies for reinforcement learning.
  • Reward Fine-Tuning with Supervised Learning: Supervised reward fine-tuning weights or filters images generated by pretrained models, whereas online variants train on examples generated by the current policy.
  • Reward Fine-Tuning with Reinforcement Learning: Reinforcement learning methods formulate denoising as multi-step decision making and use policy gradients to optimize diffusion models for black-box objectives.
  • Reward Feedback Learning (ReFL): ReFL evaluates rewards on one-step predicted clean images from randomly selected denoising steps, while DRaFT evaluates the final image; DRaFT-LV trains approximately 2× faster.

3 BACKGROUND ON DIFFUSION MODELS

Diffusion models generate samples by iteratively denoising pure noise with a learned noise-prediction function. Conditional models incorporate context such as text prompts, while classifier-free guidance combines conditional and unconditional estimates.

  • Diffusion Models: Diffusion models generate samples by repeatedly applying a learned denoising function to pure noise over T timesteps.
  • Diffusion Models: The denoiser is trained to predict noise added to clean data, and image models typically parameterize it with a UNet.
  • Diffusion Models: Conditional diffusion models pass context c, such as a text prompt, to the denoising function during sampling.
  • Classifier-Free Guidance: Classifier-free guidance combines conditional and unconditional score estimates using a guidance weight w(t), with empty conditioning represented by ∅.The unconditional score function is learned by replacing the context with ∅ during part of training.

4 METHOD

DRaFT fine-tunes diffusion models by backpropagating differentiable reward gradients through sampling into adaptable model weights. DRaFT-K and DRaFT-LV reduce the cost or variance of this gradient computation.

  • DRaFT: DRaFT fine-tunes pre-trained diffusion models by maximizing a differentiable reward over images generated by the sampling process.The method computes reward gradients through multiple diffusion-model calls and applies gradient ascent.
  • DRaFT: LoRA freezes pre-trained weights and adds low-rank matrices whose contributions are summed with the original model outputs.DRaFT uses LoRA instead of fine-tuning the full parameter set, reducing the number of trainable parameters.
  • DRaFT: Gradient checkpointing reduces memory use by storing selected activations and recomputing other UNet activations during backpropagation.The implementation stores the input latent for each denoising step and re-materializes UNet activations as needed.
  • DRaFT-K: DRaFT-K truncates backpropagation to the last K sampling steps, reducing UNet backward passes while improving optimization speed and overall performance.For small K, such as K = 1, unrolling has low memory cost and does not require gradient checkpointing.
  • DRaFT-LV: DRaFT-LV reduces DRaFT-1 gradient variance by adding examples obtained by noising the generated image and summing their reward gradients.Using n = 2 is around 2× more efficient than DRaFT-1 while adding around 10% compute overhead for the evaluated reward functions.
  • Unified framework: Algorithm 1 unifies DRaFT, DRaFT-K, DRaFT-LV, and ReFL according to how each method computes gradients through diffusion sampling.The sampling process runs from t = T to 0 with context c, starting from xT sampled from a standard normal distribution.

5 EXPERIMENTS

Experiments show that DRaFT methods efficiently optimize diverse differentiable rewards for Stable Diffusion 1.4, with truncated variants often improving speed and performance. The methods generalize across prompts and rewards, support controllable reward mixing through LoRA weights, and reveal optimization and reward-overfitting trade-offs.

  • Experimental setup: DRaFT is evaluated with Stable Diffusion 1.4 across differentiable rewards, using DDIM with 50 sampling steps and backpropagation through both latent sampling and image decoding.DRaFT-50 denotes full-chain backpropagation, with classifier-free guidance weight 7.5.
  • Aesthetic quality: DRaFT outperforms reinforcement-learning baselines in aesthetic fine-tuning, while single-step truncation and lower-variance gradients further improve sample efficiency.The LAION Aesthetic Classifier rates images from 1 through 10; training used 45 animal-name prompts, and DRaFT-LV improved efficiency beyond DRaFT-1.
  • Human preference rewards: DRaFT-LV achieves the best HPSv2 reward and learns approximately 2× faster than DRaFT-1, while DRaFT-1 slightly outperforms ReFL with simpler implementation.The comparison uses HPDv2 test categories and includes DRaFT-K values for K ∈{1, 5, 10, 30, 50}; smaller K improves training time and final performance.
  • Generalization and reward behavior: DRaFT transfers well across prompt sets and fairly well across reward functions, while different rewards produce distinct generation styles.HPSv2 generally encourages more colorful but less photorealistic generations; a model trained on 45 animal prompts preserves style on other subjects.
  • LoRA control: Interpolating or combining LoRA weights enables smooth movement between pretrained and fine-tuned outputs and mixes PickScore and HPSv2 styles without additional training.Scaling LoRA parameters acts as a control over fine-tuning strength, analogous to using guidance strength to control generations.
  • Ablations and limitations: Truncated backpropagation improves compute and sample efficiency, but gradients grow with K and aesthetic performance degrades for K > 10; LoRA scaling works best among tested overfitting controls.Gradient clipping alleviates but does not eliminate the large-K gap, while KL regularization and early stopping were less effective than LoRA scaling.

A.1 ADDITIONAL EXPERIMENTAL DETAILS AND HYPERPARAMETERS

The experiments use shared optimization settings for DRaFT and ReFL, with LoRA-based UNet adaptation, 50-step DDIM sampling, and different compute scales for small- and large-scale runs.

  • Implementation: ReFL and DRaFT variants use the same hyperparameters for comparison.The appendix states that the same hyperparameters were used for ReFL and the DRaFT variants.
  • Implementation: LoRA adapts the UNet feedforward and cross-attention layers, excluding the CLIP text encoder and VAE decoder.The excluded components did not improve results in initial experiments.
  • Sampling: 50-step DDIM sampling uses classifier-free guidance with weight 7.5 during training and testing.Ancestral/DDPM sampling produced similar results.
  • Compute: Small-scale runs take about 1.5 hours on 4 TPUv4s, while large-scale runs take about 8 hours on 16 TPUv4s.Large-scale runs target human-preference rewards, whereas small-scale runs cover the other experiments.
  • Reward models: The LAION aesthetic predictor scores images from 1 through 10 using a feedforward network on a CLIP image encoder.HPSv2 and PickScore instead use prompt-aware or preference-trained OpenCLIP-H reward models.
  • Baselines: DDPO applies REINFORCE-style policy gradients to increase the likelihood of high-reward diffusion sampling actions.The method adds importance-weight clipping, reward normalization, and classifier-free-guidance-aware training.
  • Baselines: Prompt engineering searches candidate prompts for high aesthetic scores, while Best of 16 reranks sixteen pretrained generations per caption.DOODL instead optimizes the initial noise with backpropagation through sampling.

B.2 FULL QUANTITATIVE RESULTS

The appendix reports quantitative-comparison setup and additional applications of differentiable reward fine-tuning, including compressibility, adversarial classification, LoRA interpolation, and object detection.

  • Quantitative results: Table 2 compares models fine-tuned on HPDv2 training data across Animation, Concept Art, Paintings, Photos, and their aggregate.Results above the separator come from Wu et al. (2023a), while results below come from the authors’ implementations and experiments.
  • JPEG rewards: JPEG compressibility optimization rewards small reconstruction error after differentiable compression and decompression.The reward is the negative Euclidean distance between the original image and its JPEG reconstruction.
  • JPEG rewards: Negating the JPEG compressibility reward produces difficult-to-compress images with high-frequency foreground and background information.Over-optimizing incompressibility can also reduce semantic information, such as generating an octopus with more than eight arms.
  • Adversarial examples: Diffusion adversarial examples condition on one text prompt while targeting a different class in a pretrained ImageNet classifier.The experiments use ResNet-50 and negative cross-entropy to a fixed target class.
  • LoRA interpolation: Scaling LoRA parameters controls reward-adaptation strength for HPSv2 and PickScore, producing semantic interpolations with the original model.Linear combinations of LoRA weights trained for different rewards yield smooth transitions between styles.
  • Object detection: Maximizing OWL-ViT scores for “strawberry” produces fruit bowls containing progressively more strawberries during fine-tuning.The object-detection setup experiments with rewards based on summed detection scores or bounding-box areas.

B.7 UNDERSTANDING THE IMPACT OF K

The ablations show that DRaFT-K adaptation affects more than the differentiated suffix of sampling, benefits from early LoRA application, and has optimization behavior linked to gradient clipping and reward overfitting.

  • LoRA timing: LoRA parameters must be applied for at least 10–20 sampling steps to produce substantial image changes, despite truncated backpropagation through only the last K steps.Applying LoRA early in the sampling chain is also important.
  • LoRA timing: Applying LoRA only during early sampling steps can be beneficial, demonstrating that adaptation is not confined to the final sampling steps.The experiments compare LoRA use in the first or last M steps across M values from 2 to 50.
  • Gradient behavior: DRaFT-1 and DRaFT-K gradients become nearly orthogonal for K > 30.The comparison measures angles between the DRaFT-1 gradient and gradients from different K values.
  • Gradient clipping: For DRaFT-50, c = 0.001 improved optimization substantially, whereas c = 1000 impeded training; DRaFT-1 was nearly unchanged across c = 0.001 to c = 1000.0.The results support exploding-gradient concerns for long diffusion sampling chains.
  • Reward overfitting: LoRA scaling was more effective than early stopping or KL regularization for reducing cross-reward overfitting.The comparison trains on HPSv2 and evaluates generalization to PickScore.
  • Diversity: Large dropout rates such as 0.95 did not prevent diversity collapse, while a strong LPIPS diversity term increased diversity at the expense of aesthetic score.The experiments tested dropout and minibatch dissimilarity rewards as diversity interventions.
  • CLIP reward: Optimizing CLIP similarity improved CLIP scores but degraded image quality after many training steps.The result illustrates a mismatch between optimizing the reward and preserving broader image quality.
  • Safety: The reported safety experiment changed outputs from mostly swimsuit images before fine-tuning to lay-flat swimsuit images afterward.The passage places this result within multiple possible safety stages in the text-to-image pipeline.

B.12 LEARNING SAMPLER HYPERPARAMETERS

DRaFT can optimize sampler hyperparameters as well as model weights by differentiating through the denoising process, including guidance schedules and negative prompt embeddings.

  • Sampler optimization: End-to-end DRaFT gradients enable optimization of sampler hyperparameters alongside model weights.The paper explores guidance weights and negative prompt embeddings as sampler inputs.
  • Guidance weights: DRaFT learns a separate guidance weight w(t) for each DDIM step and linearly interpolates the schedule for other step counts.Guidance training uses a 50× larger learning rate than LoRA because the weights have larger magnitude.
  • Negative prompts: The learned negative prompt is a prompt-dependent embedding parameterized as MLP(Multi-Headed-Self-Attention(e)).Here e denotes text embeddings produced by the diffusion model’s text encoder.
  • Negative prompts: Negative-prompt optimization did not improve results much at large scale.This is the reported practical outcome for the explored negative-prompt parameterization.

C METHOD EXTENSIONS

The paper extends reward fine-tuning through deterministic-policy-gradient formulations and compares these methods with direct backpropagation and prior approaches. DPG can optimize the full sampling process, but its return estimates are poor at large timesteps and it underperforms the more efficient DRaFT variants.

  • Gradient-Based Extensions: DRaFT-LV requires additional reward-gradient computations, and a variant reusing one input-pixel gradient was unstable and did not learn effectively.The unstable variant may rely on the assumption that gradients at the generated and predicted-clean images are approximately equal.
  • Deterministic Policy Gradient: DPG models DDIM sampling as a deterministic policy in a deterministic Markov decision process, with the denoiser as policy and the final clean-image reward as the only nonzero reward.States include the current latent, timestep, and prompt; transitions perform DDIM sampling steps.
  • Deterministic Policy Gradient: The DPG critic is trained from one-step-denoised images, using LoRA-adapted reward models to estimate expected returns more efficiently than evaluating noisy inputs.The method executes an action, performs one-step denoising, and applies the adapted reward model; gradients through denoising are not propagated into the policy.
  • Results: DPG outperforms vanilla DRaFT but not the more-efficient DRaFT variants, while poor critic estimates at large timesteps remain a central challenge.A ReFL-style policy-training variant that rarely samples high timesteps is suggested as a possible mitigation.
  • Connections to Prior Work: Unlike ReFL and the DRaFT variants, vanilla DRaFT and DPG optimize the full sampling process in an unbiased way.The paper identifies improving DPG for reward fine-tuning as an open direction because many DPG variants remain unexplored.
  • Connections to Prior Work: DRaFT differs from prior work by optimizing arbitrary differentiable rewards through diffusion sampling, while related methods target sampler speed, supervised reward weighting, or reinforcement-learning objectives.DRaFT retains standard diffusion sampling cost after fine-tuning and supports compositionality and interpolation through different LoRA parameter sets.

E UNCURATED SAMPLES

The paper presents uncurated Stable Diffusion samples fine-tuned with DRaFT for HPSv2, PickScore, and a combined reward. The single-reward models overfit their rewards, while LoRA scaling mitigates overfitting for the combined reward.

  • HPSv2: DRaFT-1 samples are shown for Stable Diffusion fine-tuned on the HPSv2 reward.The samples are uncurated.
  • PickScore: DRaFT-1 samples are shown for Stable Diffusion fine-tuned on the PickScore reward.The samples are uncurated.
  • Reward Overfitting: The HPSv2- and PickScore-fine-tuned models overfit their reward functions, decreasing diversity and photorealism.The paper contrasts these single-reward examples with a combined-reward model.
  • Combined Reward: For the combined reward, LoRA weights are scaled down by 0.75 to reduce reward overfitting.The combined reward uses PickScore = 10, HPSv2 = 2, and Aesthetic = 0.05, with prompts from the HPSv2 benchmark.
Loading 2309.17400v2…