Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Improving GFlowNets for Text-to-Image Diffusion Alignment

Dinghuai Zhang, Yizhe Zhang, Jiatao Gu, Ruixiang Zhang, Josh Susskind, Navdeep Jaitly, Shuangfei Zhai

TL;DR

This work reframes text-to-image diffusion alignment as a GFlowNet problem to sample outputs with probability proportional to a non-negative reward, addressing RL drawbacks like slow credit assignment and mode collapse. It introduces DAG, a family of algorithms that repurpose the denoising structure of diffusion models into GFlowNet transitions, with two variants: DAG-DB (DB-based) and DAG-KL (KL-based with REINFORCE-like gradients). A forward-looking diffusion-specific technique and a KL-based objective stabilize training and improve sample efficiency, demonstrating superior reward-diversity trade-offs on Stable Diffusion across aesthetic, image-quality, and prompt-alignment rewards. The approach enables effective post-training alignment with black-box rewards, offering a practical path for controllable generation in large-scale diffusion models while maintaining sample diversity.

Abstract

Diffusion models have become the de-facto approach for generating visual data, which are trained to match the distribution of the training dataset. In addition, we also want to control generation to fulfill desired properties such as alignment to a text description, which can be specified with a black-box reward function. Prior works fine-tune pretrained diffusion models to achieve this goal through reinforcement learning-based algorithms. Nonetheless, they suffer from issues including slow credit assignment as well as low quality in their generated samples. In this work, we explore techniques that do not directly maximize the reward but rather generate high-reward images with relatively high probability -- a natural scenario for the framework of generative flow networks (GFlowNets). To this end, we propose the Diffusion Alignment with GFlowNet (DAG) algorithm to post-train diffusion models with black-box property functions. Extensive experiments on Stable Diffusion and various reward specifications corroborate that our method could effectively align large-scale text-to-image diffusion models with given reward information.

Improving GFlowNets for Text-to-Image Diffusion Alignment

TL;DR

This work reframes text-to-image diffusion alignment as a GFlowNet problem to sample outputs with probability proportional to a non-negative reward, addressing RL drawbacks like slow credit assignment and mode collapse. It introduces DAG, a family of algorithms that repurpose the denoising structure of diffusion models into GFlowNet transitions, with two variants: DAG-DB (DB-based) and DAG-KL (KL-based with REINFORCE-like gradients). A forward-looking diffusion-specific technique and a KL-based objective stabilize training and improve sample efficiency, demonstrating superior reward-diversity trade-offs on Stable Diffusion across aesthetic, image-quality, and prompt-alignment rewards. The approach enables effective post-training alignment with black-box rewards, offering a practical path for controllable generation in large-scale diffusion models while maintaining sample diversity.

Abstract

Diffusion models have become the de-facto approach for generating visual data, which are trained to match the distribution of the training dataset. In addition, we also want to control generation to fulfill desired properties such as alignment to a text description, which can be specified with a black-box reward function. Prior works fine-tune pretrained diffusion models to achieve this goal through reinforcement learning-based algorithms. Nonetheless, they suffer from issues including slow credit assignment as well as low quality in their generated samples. In this work, we explore techniques that do not directly maximize the reward but rather generate high-reward images with relatively high probability -- a natural scenario for the framework of generative flow networks (GFlowNets). To this end, we propose the Diffusion Alignment with GFlowNet (DAG) algorithm to post-train diffusion models with black-box property functions. Extensive experiments on Stable Diffusion and various reward specifications corroborate that our method could effectively align large-scale text-to-image diffusion models with given reward information.
Paper Structure (25 sections, 1 theorem, 15 equations, 11 figures, 1 table, 1 algorithm)

This paper contains 25 sections, 1 theorem, 15 equations, 11 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Diffusion models
  4. GFlowNets
  5. Detailed balance (DB)
  6. Methodology
  7. Denoising Markov decision process
  8. Diffusion alignment with GFlowNets
  9. Incorporating denoising diffusion-specific structure
  10. A KL-based GFlowNet algorithm with REINFORCE gradient
  11. Related Works
  12. Diffusion alignment
  13. GFlowNets
  14. Experiments
  15. Experimental setups
  16. ...and 10 more sections

Key Result

Proposition 4

The KL term in Equation eq:db_kl has the same expected gradient with $b(\mathbf{x}_{t}, \mathbf{x}_{t-1})\log p_{\boldsymbol{\theta}}(\mathbf{x}_{t-1}|\mathbf{x}_t)$:

Figures (11)

  • Figure 1: Generated samples before (top) and after (bottom) the proposed training with Aesthetic reward.
  • Figure 2: Top: samples from the original Stable Diffusion model. Middle: the proposed method trained with compressibility reward; these images have very smooth texture. Down: the proposed method trained with incompressibility reward; the texture part of images contains high frequency noise.
  • Figure 3: Sample efficiency results of our proposed methods and our RL baseline (DDPO). The number of training steps is proportional to the number of sampled trajectories. The experiments are conducted on reward functions including aesthetic score, ImageReward, and HPSv2.
  • Figure 4: Sample efficiency results of our proposed methods and our RL baseline (DDPO) on learning from compressibility and incompressibility rewards.
  • Figure 5: Text-image alignment results. We display four prompts and the corresponding generation visualization from the original Stable Diffusion (1st row), DDPO (2nd row), DAG-DB (3rd row), and DAG-KL (4th row) models to compare their alignment abilities. See Figure \ref{['fig:vis_hpdphoto_painting_more']} for more results.
  • ...and 6 more figures

Theorems & Definitions (7)

  • Remark 1: diffusion model as GFlowNet
  • Remark 2: GPU memory and the choice of GFlowNet objectives
  • Remark 3
  • Proposition 4
  • Remark 5: gradient equivalence to detailed balance
  • Remark 6: analysis of $b(\mathbf{x}_{t}, \mathbf{x}_{t-1})$
  • proof