Efficient Diversity-Preserving Diffusion Alignment via Gradient-Informed GFlowNets

While one commonly trains large diffusion models by collecting datasets on target downstream tasks, it is often desired to finetune pretrained diffusion models on some reward functions that are either designed by experts or learned from small-scale datasets. Existing methods for finetuning diffusion models typically suffer from lack of diversity in generated samples, lack of prior preservation, and/or slow convergence in finetuning. Inspired by recent successes in generative flow networks (GFlowNets), a class of probabilistic models that sample with the unnormalized density of a reward function, we propose a novel GFlowNet method dubbed Nabla-GFlowNet (abbreviated as -GFlowNet), the first GFlowNet method that leverages the rich signal in reward gradients, together with an objective called -DB plus its variant residual -DB designed for priorpreserving diffusion alignment. We show that our proposed method achieves fast yet diversity-and prior-preserving alignment of Stable Diffusion, a large-scale text-conditioned image diffusion model, on different realistic reward functions. Diffusion models [16, 56, 47] are a powerful class of generative models that model highly complex data distributions as the results of a sequence of multi-scale denoising steps. State-of-the-arts diffusion models for downstream applications are typically large in network size and demand a significant amount of data to train.