Improving GFlowNets for Text-to-Image Diffusion Alignment

2024-06-02

arXiv.org

AI Summary

To address alignment difficulties, low sample quality, and inefficient credit assignment in text-to-image diffusion models under black-box reward functions, this paper proposes Diffusion Alignment with GFlowNet (DAG), the first post-training alignment framework for diffusion models based on Generative Flow Networks (GFlowNets). DAG abandons policy-gradient-based reinforcement learning, instead directly modeling high-reward generation trajectories via flow-matching objectives—enabling gradient-free, efficient, and scalable black-box alignment. Compatible with Stable Diffusion, DAG significantly outperforms RL-based fine-tuning baselines in both CLIP-Score and human preference evaluations, simultaneously improving alignment fidelity and image quality while accelerating convergence by 2–3×. Its core contribution lies in the first application of GFlowNets to diffusion model alignment, establishing a novel paradigm for black-box reward-driven generative alignment.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Neural Architectures and Foundation Models
• Multiagent Systems: Adversarial Agents

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Web Mining and Content Analysis: Large pretrained models with web data

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.

Problem

• Image-Text Matching
• Image Quality Improvement
• Complex Scoring Criteria

Innovation

• DAG Algorithm
• Enhanced Diffusion Model
• Accelerated Learning Process