PaperlayerView exact revision on arXiv

Source-linked AI summary

Large-scale Reinforcement Learning for Diffusion Models

Yinan Zhang, Eric Tzeng, Yilun Du, Dmitry Kislyuk

arXiv:2401.12244v1cs.CVcs.AIcs.LG

TL;DR

Text-to-image models trained on web-scale text-image pairs can miss prompt semantics and human aesthetic preferences. This paper scales reinforcement-learning fine-tuning across objectives, reporting improvements in human preference, fairness, and composition, as well as stronger aesthetic scores and human preference on PartiPrompts than baseline models.

  • Problem

    Web-scale text-image pretraining can leave diffusion models weak at faithfully depicting prompt semantics and relationships and at matching human aesthetic preferences.

    Source: §1 · p. 1

  • Method

    The paper presents a scalable RL framework that fine-tunes diffusion models across millions of prompts and multiple objectives, including a distribution-based reward for diversity.

    Source: Yilun Du MIT CSAIL · p. 1

  • Results

    The approach improves the model across human preference, fairness, and composition objectives, and achieves the highest Aesthetic score and human preference rate among baseline models on PartiPrompts.

    Source: Yilun Du MIT CSAIL · p. 1, §4.1 · p. 6

  • Takeaways & Limitations

    Joint training mitigates the alignment-tax issue common in RLHF while improving multiple objectives.

    Source: §5 · p. 8

    ImageReward alone is not a robust performance indicator because the model was directly trained against it, and single-reward optimization can overoptimize that metric.

    Source: §4.1 · p. 6

Abstract

from arXiv · show

Text-to-image diffusion models are a class of deep generative models that have demonstrated an impressive capacity for high-quality image generation. However, these models are susceptible to implicit biases that arise from web-scale text-image training pairs and may inaccurately model aspects of images we care about. This can result in suboptimal samples, model bias, and images that do not align with human ethics and preferences. In this paper, we present an effective scalable algorithm to improve diffusion models using Reinforcement Learning (RL) across a diverse set of reward functions, such as human preference, compositionality, and fairness over millions of images. We illustrate how our approach substantially outperforms existing methods for aligning diffusion models with human preferences. We further illustrate how this substantially improves pretrained Stable Diffusion (SD) models, generating samples that are preferred by humans 80.3% of the time over those from the base SD model while simultaneously improving both the composition and diversity of generated samples.

1. Introduction

Text-to-image models trained on web-scale text-image pairs can struggle to depict prompt semantics and relationships accurately, and may produce images that do not match human aesthetic preferences.

1. Introduction

  • Web-scale text-image pretraining exposes models to diverse prompts, but downstream applications still observe weaknesses in fidelity, controllability, and aesthetic alignment.The passage links these weaknesses to dependence on the pretraining-data distribution.Source: §1 · p. 1
  • Fidelity and controllability failures include inaccurate prompt semantics, such as incorrect object composition and relationships.Source: §1 · p. 1
  • Human aesthetic mismatch occurs when generated outputs are not perceived as aesthetically pleasing.Source: §1 · p. 1

Yilun Du MIT CSAIL

The paper proposes scalable reinforcement-learning fine-tuning for diffusion models across millions of prompts and diverse objectives, including human preference, fairness, and composition.

Yilun Du MIT CSAIL

  • The proposed generic RL framework fine-tunes diffusion models across millions of prompts and arbitrary combinations of objectives.Source: Yilun Du MIT CSAIL · p. 1
  • A distribution-based reward function is proposed to improve output diversity.Source: Yilun Du MIT CSAIL · p. 1
  • Multi-objective RL training targets human aesthetic preference, fairness, and object composition together.Source: Yilun Du MIT CSAIL · p. 1
  • The experiments compare the approach with existing reward-optimization methods across a suite of tasks.Source: Yilun Du MIT CSAIL · p. 1
  • Figure 1 tracks sample progression across training objectives, with base SDv2 samples in the leftmost columns.Source: Figure 1 · p. 2

2. Related Work

Prior diffusion-model reward fine-tuning uses supervised reward-weighted data, reward gradients, or reinforcement learning; the paper positions its approach as scaling RL and addressing composition and inclusion without inference-time guidance.

2. Related Work

  • Prior diffusion reward fine-tuning uses reward-weighted data, reward-function backpropagation, or reinforcement learning, while earlier RL work used narrow prompt sets and objectives.The paper contrasts those settings with RL across millions of prompts and different objectives.Source: §2 · p. 2
  • Prior compositional-generation methods modify inference or add layout guidance, whereas this method targets fidelity without additional layout guidance.Source: §2 · p. 2
  • Prior inclusive-generation approaches use balanced synthetic data, textual inference guidance, or reference images; this method requires neither synthetic-data collection nor inference-time intervention.Source: §2 · p. 2

3. Method

The method fine-tunes an existing diffusion model to maximize rewards on generated images, framing denoising as a multi-step decision process. It also addresses distribution-level objectives and jointly trains across multiple reward tasks.

3.1. Policy Gradient with Multi-step MDP

  • The method frames iterative denoising as a multi-step Markov decision process, with the diffusion model as policy and reward assigned only to the final sample.Policy-gradient estimates use the likelihood-ratio method.Source: §3.1 · p. 3
  • For millions of prompts, rewards are normalized per training batch; importance sampling and a clipped trust region support updates, while the original diffusion objective helps prevent reward over-optimization.The clipped interval is controlled by a hyperparameter.Source: §3.1 · p. 3

3.2. Distribution-based Reward Functions

  • Distribution-level rewards address goals such as diversity that cannot be evaluated from isolated generations; the method approximates them using empirical minibatch samples and backpropagates minibatch rewards across those samples.The paper applies this approach to fairness and diversity objectives.Source: §3.2 · p. 3

3.3. Multi-task Joint Training

  • Joint training optimizes one model across human preference, skintone diversity, and object composition by sampling prompts from each task and applying its reward model and gradient step sequentially.Each task uses its own prompt distribution.Source: §3.3 · pp. 3–4
  • Figure 2 evaluates human preference through head-to-head image comparisons on 400 prompts, using matched prompts and random seeds and five raters per query.Raters judge image quality and image-text alignment; preference rates reflect positive votes.Source: Figure 2 · p. 4

4. Reward Functions and Experiments

The section evaluates RL fine-tuning on SDv2 for human preference, skintone fairness and diversity, and object composition, then jointly optimizes the three reward objectives.

4.1. Learning from human preference

  • On DiffusionDB and out-of-domain PartiPrompts, our method achieves the highest Aesthetic score and human preference rate among baselines, and converges in about 1,000 steps versus DRaFT's about 4,000.ImageReward alone is not a robust performance indicator when the model was directly trained against it; RAFT diverged as iterations increased.Source: Table 1 · p. 5, §4.1 · pp. 5–6, Table 2 · p. 5

4.2. Optimizing Fairness and Diversity

  • A four-category skintone classifier rewards uniform minibatch distributions; the fine-tuned model reduces SDv2's skintone bias, especially for occupations with stronger social stereotypes.The test sets include 100 randomly sampled occupations and 200 HRSBench prompts, both out of domain relative to the BLIP-generated training prompts.Source: §4.2 · p. 7

4.3. Optimizing Compositionality

  • Training pairs 450 of 532 object classes across five spatial relations, and detector-based evaluation finds better caption-constraint adherence than SDv2, including on unseen objects.The model's learned compositional abilities generalize to unseen objects.Source: Table 4 · p. 7, §4.3 · p. 7

4.4. Multi-reward Joint Optimization

  • Over 80% relative performance of individually fine-tuned models is retained simultaneously across all three metrics, while the joint model significantly outperforms base SDv2 on every task.The best score for each metric still comes from its task-specific model; single-objective fine-tuning can regress on other tasks.Source: Table 5 · p. 8, §4.4 · p. 8

5. Conclusion

The conclusion presents scalable RL fine-tuning for diverse objectives, including distribution-based rewards, and reports multi-task improvements and reduced alignment tax. The appendix outlines supporting experimental details, results, and analyses.

5. Conclusion

  • The framework fine-tunes diffusion models at scale for arbitrary downstream objectives, including distribution-based rewards.The authors report joint fine-tuning for human preference, fairness, and object composition, with mitigation of the alignment tax common in RLHF.Source: §5 · p. 8
  • The authors report generality and robustness from evaluating the trained model against baselines on diverse out-of-domain test sets.They identify more complex compositional relationships and additional social dimensions of bias as future topics.Source: §5 · p. 8

Appendix

  • The appendix covers experimental setup, additional qualitative comparisons, human-rating guidelines, diversity and compositionality evaluations, training curves, and reward-hacking examples.Source: Appendix · p. 11
  • It also includes a pretraining-dataset ablation and complete occupation and object lists for the diversity and compositionality evaluations.Source: Appendix · p. 11

A. Experiment Details and Hyperparameters

Experiments use shared implementation settings and common training data for baselines; the appendix also provides qualitative results for single-reward and jointly trained models.

A. Experiment Details and Hyperparameters

  • Experiments, including baseline training, used 128 80GB A100 GPUs, and all fine-tuning methods used the same aesthetic-filtered 12M LAION-5B subset when pretraining data was required.Source: A. Experiment Details and… · p. 11
  • For RL fine-tuning, each iteration collected 16x128 samples, used 50 DDIM sampling steps, and performed five gradient updates across sampled timesteps.The experiments used a 1e−4 clip range; baseline hyperparameters followed the original implementations, and models were trained on the same training set until convergence.Source: A. Experiment Details and… · p. 11

B. Additional Qualitative Results

  • Additional qualitative results cover ImageReward, compositionality, and skintone-diversity single-reward models, as well as a model jointly trained with all three rewards.Source: B. Additional Qualitative Results · p. 11

B.1. Results on Human Preference Fine-tuning

Qualitative comparisons report more aesthetically aligned outputs from the fine-tuned model, including on out-of-domain prompts, and show its comparisons with reward-tuning baselines.

B.1. Results on Human Preference Fine-tuning

  • The fine-tuned model's real-user-prompt outputs favor finer details, focused composition, vivid colors, and high contrast in line with human aesthetic preferences.Source: Figure 7 · p. 12, §B.1 · p. 12
  • Figure 8 compares the fine-tuned model with other reward-optimization methods on same-seed real-user prompts from DiffusionDB.Source: Figure 8 · p. 13
  • On PartiPrompts, the model generates samples with higher aesthetic quality and better image-text alignment than the other baseline models.The authors also report generalization to unseen prompts from a different distribution than the training prompts.Source: §B.1 · p. 12, Figure 9 · p. 14

B.2. Results on Optimizing Diversity

Single-reward fine-tuning improves skintone diversity and object composition, while joint optimization improves human preference, skintone diversity, and composition. The accompanying evaluation materials define rating criteria and comparison procedures.

B.2. Results on Optimizing Diversity

  • Skintone-diversity fine-tuning mitigates inherent bias and stereotypes in base SDv2, increasing skintone diversity in generated human samples.The comparison includes RAFT, which also improves skintone diversity but can overfit and produce over-saturated samples with decreased realism.Source: §B.2 · p. 15, Figure 10 · p. 15

B.3. Results on Optimizing Compositionality

B.4. Results on Multi-reward Joint Optimization

  • Joint optimization of all three rewards significantly improves human preference, skintone diversity, and object composition over base SDv2.For occupation portraits, the joint model also shows an additional aesthetic-quality increase over single-reward skintone-diversity training.Source: §B.4 · p. 16, Figure 12 · p. 17, Figure 13 · p. 17

C.1. Evaluation Criteria and Guidelines

C.1.4 Overall Preference Ranking

  • For most tasks, fidelity and aesthetic quality outweigh text alignment, except when stronger alignment accompanies only slight fidelity loss and no significant discomfort-causing artifacts.Source: §C.1.4 · pp. 20–21

C.2. Evaluation Interface

  • The head-to-head interface presents same-prompt, same-seed generations and asks evaluators to choose based on fidelity, aesthetic quality, and image-text relevance.Source: §C.2 · p. 22, Figure 23 · p. 23

D.1. Additional Results

These sections describe human evaluations of compositionality and skintone diversity, including the rating criteria, evaluator guidance, and comparison procedure.

D.1. Additional Results

  • Compositionality samples were rated for image–text relevance, while skintone-diversity samples were judged by their overall diversity.Figure 24 reports these two human evaluations.Source: §D.1 · p. 23, Figure 24 · p. 23

D.2.1 Guidelines

  • Skintone diversity is operationalized as a balanced distribution of light and dark skintones, rather than only one or the other.The guidelines illustrate this with portraits of a police officer, chemist, and builder.Source: §D.2.1 · p. 24
  • For each prompt, evaluators choose the more diverse of two image groups, each containing six images.The guidelines include examples in which a group with a balanced distribution of light and dark skintones is considered more diverse.Source: §D.2.1 · p. 24, Figure 25 · p. 24

D.2.2 Interface

  • Skintone-diversity comparisons use two groups generated from the same prompt and random seed, and evaluators are trained with the provided guidelines.The evaluation interface presents the two groups for comparison.Source: §D.2.2 · p. 24, Figure 26 · p. 25

E. Training Curve

The training-curve analysis reports rapid convergence for the proposed approach and reward hacking in ReFL, with examples of artifacts associated with that issue.

E. Training Curve

  • Except for RAFT, which diverged, online-learning methods show steadily increasing sample rewards until they saturate at convergence.The authors consider models converged once rewards reach a maximum level.Source: E. Training Curve · p. 25
  • 1,000 steps versus DRaFT’s 4,000: our approach converges faster, while DRaFT optimizes only differentiable rewards.The training curves plot average reward per batch against training iteration.Source: Figure 27 · p. 25, E. Training Curve · p. 25

F. Reward Hacking

  • ReFL can exploit imperfections in a human-annotation-trained reward model to pursue high scores, resulting in reward hacking.The paper identifies this as a known RLHF issue.Source: F. Reward Hacking · p. 26
  • ReFL’s ImageReward values rise quickly, but examples show over-detailed images with high-frequency noise.The paper describes these generations as reward hacking.Source: Figure 28 · p. 26

G. Effect of Pretraining Dataset

The appendix reports that omitting the pretraining denoising loss is associated with grainy artifacts and reduced realism, and lists the occupations and objects used in evaluation and training.

G. Effect of Pretraining Dataset

  • Without pretraining loss, samples are more prone to grainy artifacts and decreased realism, compared with samples from the same number of iterations and random seeds.The loss is incorporated to stabilize training and prevent reward over-optimization.Source: G. Effect of Pretraining Dataset · p. 26, Figure 29 · p. 27

H. Full List of Occupations and Objects

Loading 2401.12244v1…