Efficient Diversity-Preserving Diffusion Alignment via Gradient-Informed GFlowNets

Authors: Zhen Liu, Tim Z. Xiao, Weiyang Liu, Yoshua Bengio, Dinghuai Zhang
Primary area: generative models
Keywords: Reward Finetuning, Diffusion Models, Alignment, GFlowNet, Generative Models

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 $\nabla$-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.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730467043352

thread_type:
official_review

<thread_object>
creator: reviewer_5qst

<content>
summary:
This paper introduces Nabla-GFlowNet with a new objective, $ \nabla$-DB, and its variant, residual $ \nabla$-DB, for finetuning pretrained diffusion models. The method aims to enhance sample diversity and finetuning efficiency by utilizing reward gradients. Experiments on two reward functions, Aesthetic Score and Human Preference Score (HPSv2), demonstrate improved performance.

strengths:
- The paper has a clear motivation, addressing the challenge of preserving diversity and improving efficiency in finetuning diffusion models, which is crucial for real-world applications.
- The paper introduces a unique application of GFlowNet principles to diffusion model finetuning, specifically focusing on the reward gradient to preserve sample diversity.

weaknesses:
- In the quantitative evaluation, the proposed method performs well on the smaller dataset of Aesthetic Score, nearly doubling the baseline in the DreamSim metric. However, on the larger dataset of HPSv2, its performance is similar to baselines, showing no clear advantage. The effectiveness of the method remains inconclusive due to the differences in dataset size.
- While the authors mention “fast and efficient finetuning” as a contribution, only Figure 4 shows comparable convergence speed to the baseline. It would be helpful to include details on training resource consumption, such as GPU usage and computational cost, to substantiate this claim.
- Figure 3 only compares the results of the pretrained model and the proposed method, lacking visual comparisons with other baselines.
- Figure 5 lacks sufficient information in the title and annotations to clarify what each point represents (e.g., training iteration). To improve clarity, consider adding an explanation in the figure caption or in Section 4.4 to specify the meaning of each point.
- The conclusion section is missing, which may limit the clarity of the paper’s overall findings and contributions.

questions:
Please see weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730593399068

thread_type:
official_review

<thread_object>
creator: reviewer_5n8z

<content>
summary:
The authors propose a method, called $\nabla$-GFlowNet which which is modification of GFLowNets. They define a new objective $\nabla$-DB, which is a gradient informed version the Detailed Balance objective. The paper then goes on to propose a residual version of this loss which at optimality samples proportionally to the argumented distribution $r(x_T)p^\sharp (x_T)$, which maintains diversity of generations. The paper then presents experiments which shows a good pareto frontier of reward vs diversity on two tasks.

strengths:
- This paper offers a good way at maintaining diversity in generations while aligning generations to a reward model.
- The paper is theoretically founded and shows that their new training objective maintains the validity of GFlowNets while taking into account gradient information.
- The work has a number of good ablations and experiments to show diversity and reward tradeoff.

weaknesses:
- The problem of bias is well known in optimizations for diffusion models (elaborated in question section). Some treatment of this problem problem would be desirable since it seems that this method may be optimizing for the biased distribution

questions:
- It is known that optimizing the KL constrained optimization problem, with the closed form solution of the augmented distribution, can lead to a biased result [0, 1]. Does this method have this bias problem? If not, how do you get around it?

[0] https://arxiv.org/abs/2409.08861
[1] https://arxiv.org/abs/2402.15194
</content>
</thread_object>

</thread>
<thread>
creation date:
1730617055921

thread_type:
official_review

<thread_object>
creator: reviewer_y62f

<content>
summary:
The authors propose Nabla-GFlowNet (∇-GFlowNet) to efficiently finetune pretrained diffusion models. This approach addresses issues of limited sample diversity and slow convergence in existing methods by leveraging reward gradients through ∇-DB and its variant, residual ∇-DB. Empirical results show that residual ∇-DB enables fast, diversity-preserving finetuning of models like StableDiffusion on various realistic reward functions.

strengths:
- The paper offers a comprehensive theoretical deduction of the proposed method, thoroughly explaining how the objectives nabla-DB and residual nabla-DB are derived. 
- By introducing residual ∇-DB, the authors extend the applicability of their work to pretrained large-scale models, which is crucial.
- The paper enhances the quantitative evaluation of diversity in generated samples. By employing a broader range of metrics and more extensive comparisons.

weaknesses:
- The current experimental setting appears somewhat outdated. To enhance the study's relevance, please consider using more recent schedulers and pre-trained models instead of DDPM or Stable Diffusion 1.5.
- The qualitative results shown in Figure 2 are confusing. Additional explanation is needed to clearly demonstrate the superiority of ∇-DB, as DDPO and DAG-DB also exhibit strong performance.
- A user study would be helpful for evaluating diversity.

questions:
Please see weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730721336765

thread_type:
official_review

<thread_object>
creator: reviewer_beky

<content>
summary:
This paper introduces an interesting method to address the challenges of fine-tuning multistep sampling in diffusion models. It employs GFlowNets to incorporate a middle term, $F(x_t)$ or $g_\phi(x_t)$, which allows the reward score to effectively influence different timesteps.  This method successfully eliminates the need to train a reward model that accepts noisy input. This paper implements their idea in both theoretical and practical contexts. Section 3.1 covers the theoretical aspect, while sections 3.2 and 3.3 address the practical application. Experiments also show that this method can enhance reward tuning in diffusion models.

strengths:
1. This paper presents a new method for addressing the challenges of fine-tuning multistep sampling in diffusion models using GFlowNets. This method effectively eliminates the need to train a reward model that processes noisy input.
2. This paper implements their idea in both theoretical and practical contexts. Section 3.1 covers the theoretical aspect, while sections 3.2 and 3.3 address the practical application.

weaknesses:
The main weakness is in the experiment part.
1. The function $g_\phi(x_t)$ is an interesting and reasonable choice for achieving the fitness task; however, it results in approximately zero vectors, with a terminal constraint of $g_\phi(x_T) = 0$. It remains unclear whether Unet is a suitable option for this purpose.
2. The regularization term appears significant, with $\lambda=1000$ in the Aesthetic Score experiments and $\lambda=100$ in the HPSv2 experiments. However, Section 3.2 states that it "may eventually over-optimize the reward and thus neglect the pretrained prior." Is Section 3.2 more helpful for preventing over-optimization than the regularization term?
3. The method is interesting and useful, but the paper's claim of "Diversity-Preserving" in the title creates a gap. Does this imply that the method theoretically ensures better diversity, or is it based solely on observational results? From my perspective, this method prioritizes improving information backpropagation for fine-tuning multistep sampling rather than enhancing diversity.

If the authors can resolve my issue, I will contemplate raising my score.

questions:
See the weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730962572904

thread_type:
official_review

<thread_object>
creator: reviewer_zkle

<content>
summary:
The paper addresses the problem of fine-tuning the pretrained diffusion models on a target reward while aiming for 1) preserving diversity of generated images and 2) fast convergence. It proposes Nabla-GFlowNet to do so, inspired by generative Flow Nets (GFlow-Nets) that sample with unnormalized density of the reward function. Experiments on different benchmarks show that the proposed method generally achieves the best diversity vs reward trade-off compared to baselines.

strengths:
- The proposed idea is based on the generative flow nets, which makes it intuitive and straightforward. 
- The Nabla-GFlowNet can leverage the first order information of the reward function (gradient) while the baselines only use the zero-order information. 
- The experimental results show that the proposed method can generally achieve the best diversity vs. reward trade-off frontiers.

weaknesses:
- I think the "predicted reward" estimation in Eq. 15 can be severely unreliable, especially for the high-noise time-steps of the diffusion model. The predicted clean image will be noisy, and if the reward function is calculated by a model that has been trained on not noisy images, the predicted reward will be inaccurate.

- The parameter \lambda and the output regularization described in Page 7 seems to be crucial to the model's performance, but they are not the paper's contribution.

- The qualitative samples that are compared with the baselines are only for the aesthetic score. I think a qualitative comparison on the HPSv2 can be more valuable and insightful about the model's performance.

questions:
- I suggest that the authors include the qualitative comparative results for the HPS-v2 reward in the paper.

If the authors address my concerns, I am willing to increase my score.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732189724813

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response to the reviewers

comment:
We sincerely thank the reviewers for their time and effort in reviewing our paper and making valuable suggestions. We have uploaded a revised draft of our paper. The major differences are:

- We have 1) fixed the numerical issue during training, 2) introduced **time-dependent weighting on the predicted rewards** of the residual $\nabla$-DB loss. As a result, our proposed method is now able to achieves a **much better trade-off between reward, diversity and prior-following on all reward models**, shown in Table 1 and Figure 6, 7, 12 & 13 for the main experiments in the revised draft. We have therefore also revised our Figure 2 on the results using Aesthetic Score.
- To further demonstrate the effectiveness of our methods and in response to Reviewer zWHc and 5qSt, we include more qualitative results in Figure 3, 4, 5 as well as the last three pages in the appendix.
- We now introduce a **new hyperparameter to control the strength of the pretrained prior**, for which the ablation results are shown in Figure Figure 16 and 17.
- We now include a **new metric of prior preservation** by measuring the FID score (in general, lower the better) between samples generated by the finetuned model and the pretrained model (as our objective is $P_\text{pretrained}(x)^\eta * R(x)^\beta$). For a fair comparison, we draw the Pareto frontier to show the FID scores of models at roughly the same reward. Figure 7 shows that our method is significantly better in this dimension.
- We now include **another reward function ImageReward** [1], a general-purpose text-to-image human preference reward model.
- In response to Reviewer y26f and 5qSt, we include some more details in the experiment section to avoid confusion.
- In Figure 2, we used to show the final results (trained with 200 update steps) of all methods. While it is indeed the case that our method is more robust, it might be a bit unfair as one may perform early stopping when finetuning models. Therefore, we now visualize the samples from the best model checkpoint (throughout the finetuning process) for each method with an average reward of the visualized set of images, while we qualitatively show the training collapse issue of the baseline methods in Figure 3.
- Due to the limited space in the paper, we moved some figures like the one for ablation study on trajectory subsampling rate to the appendix.
- In response to Reviewer y62f, we have experimented with a different scheduler. Specifically, we pick **SDE-DPM-Solver++** and construct an MDP. In Figure 18 and 19, we show that our method is still able to perform well.
- For every experiment we now show the standard deviation of metrics with 3 random seeds.
- We polished the figures, the captions the related work section for better flow.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732426207688

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Following-up on reviewers' concerns

comment:
Dear Reviewers and AC,

We thank again for your time and efforts in reviewing our paper and going through the possibly long responses to your questions. As it is coming to the end of the rebuttal period, we would like to learn if we have addressed your concerns or the reviewers have any other questions. If any, we are more than happy to answer them.

Many Thanks,

Paper 8475 Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1732504603794

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of our contributions

comment:
We would like to copy and paste parts of our supplementary response to Reviewer ZKLE, to give a better summarization and explanation on our contributions (especially from the perspective of GFlowNets):

**From the perspective of reward finetuning.**

Our proposed method is a *practical, scalable and provable* method that

- works and scales well on **large diffusion models**;
- has **theoretical guarantees** of unbiased optimal solution to the target distribution; and,
- achieves **Pareto improvements** on reward convergence and diversity & prior preservation on **all reward functions** we experimented with.

**From the perspective of GFlowNets.**

- We propose the **first** GFlowNet objective that leverages **gradient signals** in training GFlowNets, while all other GFlowNet methods only use zeroth-order signals.
- Our method demonstrates scalability to large diffusion models, such as StableDiffusion with a substantial number of sampling steps—something that few existing GFlowNet approaches have achieved.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a novel GFlowNet method dubbed Nabla-GFlowNet (abbreviated as $\nabla$-GFlowNet), together with an objective called $\nabla$-DB, plus its variant residual $\nabla$-DB for finetuning pretrained diffusion models, to achieve fast yet diverse sampling.

All 5 reviewers give a final rating of 6. I am on board with them and recommend accept.