Efficient Diversity-Preserving Diffusion Alignment via Gradient-Informed GFlowNets

2024/12/10 by Zhen Liu, Tim Z. Xiao, Liu, Zhen +7 · 12 citations
Biochemistry, Genetics and Molecular Biology · Computer Science · #Cancer-related molecular mechanisms research #Computer Vision and Pattern Recognition (cs.CV) #Domain Adaptation and Few-Shot Learning #FOS: Computer and information sciences #Generative Adversarial Networks and Image Synthesis #Machine Learning (cs.LG)

paper · pdf · doi:10.48550/arxiv.2412.07775

openalex publication_date 2024/12/10 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

While one commonly trains large diffusion models by collecting datasets on target downstream tasks, it is often desired to align and finetune pretrained diffusion models with some reward functions that are either designed by experts or learned from small-scale datasets. Existing post-training methods for reward finetuning of diffusion models typically suffer from lack of diversity in generated samples, lack of prior preservation, and/or slow convergence in finetuning. In response to this challenge, we take inspiration from recent successes in generative flow networks (GFlowNets) and propose a reinforcement learning method for diffusion model finetuning, dubbed Nabla-GFlowNet (abbreviated as ∇-GFlowNet), that leverages the rich signal in reward gradients for probabilistic diffusion finetuning. We show that our proposed method achieves fast yet diversity- and prior-preserving finetuning of Stable Diffusion, a large-scale text-conditioned image diffusion model, on different realistic reward functions.

Cited by

• Tilt Matching for Scalable Sampling and Fine-Tuning
• RMLer: Synthesizing Novel Objects across Diverse Categories via Reinforcement Mixing Learning
• Value Gradient Guidance for Flow Matching Alignment
• EraseFlow: Learning Concept Erasure Policies via GFlowNet-Driven Alignment
• Diffusion Alignment as Variational Expectation-Maximization
• CREPE: Controlling Diffusion with Replica Exchange
• Data from Discrete Flow-Based Generative Models for Measurement Optimization in Quantum Computing
• FlowRL: Matching Reward Distributions for LLM Reasoning
• Mitigating Hallucinations in LM-Based TTS Models via Distribution Alignment Using GFlowNets
• Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference
• Concept-TRAK: Understanding how diffusion models learn concepts through concept-level attribution
• Nabla-R2D3: Effective and Efficient 3D Diffusion Alignment with 2D Rewards

Related

• Improving GFlowNets for Text-to-Image Diffusion Alignment
• Large-scale Reinforcement Learning for Diffusion Models
• Towards Better Alignment: Training Diffusion Models with Reinforcement Learning Against Sparse Rewards
• Aligning Text-to-Image Diffusion Models with Reward Backpropagation
• Self-Play Fine-Tuning of Diffusion Models for Text-to-Image Generation