扩散模型奖励微调新突破：Nabla-GFlowNet让多样性与效率兼得

本文作者刘圳是香港中文大学（深圳）数据科学学院的助理教授，肖镇中是德国马克思普朗克-智能系统研究所和图宾根大学的博士生，刘威杨是德国马克思普朗克-智能系统研究所的研究员，Yoshua Bengio 是蒙特利尔大学和加拿大 Mila 研究所的教授，张鼎怀是微软研究院的研究员。此论文已收录于 ICLR 2025。

在视觉生成领域，扩散模型（Diffusion Models）已经成为生成高质量图像、视频甚至文本的利器。然而，生成结果往往离我们所偏好的不一致：结果不美观，图文不符，等等。

虽然我们可以像大语言模型中的 RLHF（基于人类反馈的强化学习）一样直接用传统强化学习来微调扩散模型，但收敛速度往往慢；而基于可微计算图直接最大化奖励函数的方法又往往陷入过拟合和多样性缺失的问题。

有没有一种方法，既能保留生成样本的多样性，又能快速完成微调？我们基于生成流网络（Generative Flow Network，GFlowNet）提出的 Nabla-GFlowNet 实现了这一速度和质量间的平衡。

• 论文标题：Efficient Diversity-Preserving Diffusion Alignment via Gradient-Informed GFlowNets

• 论文地址：https://arxiv.org/abs/2412.07775

• 代码地址：https://github.com/lzzcd001/nabla-gfn

                                            利用 Nabla-GFlowNet 在 Aesthetic Score 奖励函数（一个美学指标）上高效微调 Stable Diffusion 模型。

扩散过程的流平衡视角

生成流网络 GFlowNet 示意图。初始节点中的「流」通过向下游的转移概率流经不同节点，最后汇聚到终端节点。每个终端节点所对应的流应匹配该终端节点对应的奖励。

在生成流网络（Generative Flow Network, GFlowNet）的框架下，扩散模型的生成过程可以视为一个「水流从源头流向终点」的动态系统：

• 从标准高斯分布采样的噪声图像作为初始状态，其「流量」为；

• 去噪过程 如同分配水流的管道网络，把每一个节点的水流分配给下游每一个节点；

• 而加噪过程则可以回溯每一个的水流来自哪里；

• 最终生成的图像将累积总流量。

                                          扩散模型示意图

流梯度平衡条件

在 GFlowNet 框架下，前后向水流需要满足一定的平衡条件。我们通过推导提出我们称为 Nabla-DB 的平衡条件：

其中是残差去噪过程，和分别是微调模型和预训练模型的去噪过程。是这个残差过程对应的对数流函数。

这个残差去噪过程应该满足，其中 β 控制微调模型在奖励函数和预训练模型之间的平衡。如果 β 为零，那么残差过程为零，也就是微调网络等于预训练网络。

稍作变换，就可以得到我们提出的 Nabla-GFlowNet 对应的损失函数 Residual Nabla-DB（其中 sg 为 stop-gradient 操作）：

前向匹配损失：

后向匹配损失：

终端匹配损失：

对数流梯度参数化设计

上述损失函数需要用一个额外的网络估计，但我们观察到：如果我们对 做单步预测得到不准确的去噪结果，那么是一个很好的估计，其中是一个权重常数。因此，我们提出如下参数化：

其中是用 U-Net 参数化的残差梯度，而单步去噪估计为（是扩散模型ε-预测参数化的网络）。

                                           方法示意图。每条采样路径的每个转移对中每张图的上下两个「力」需要相互平衡。

直观解释

如果我们只计算这一转移对的前向匹配损失对扩散模型参数的梯度，我们有：

其中第一项是基于内积的匹配度函数（残差扩散模型与奖励梯度估计之间的匹配），第二项是让微调模型趋近于预训练模型的正则化。

伪代码实现

实验结果

我们分别用以下奖励函数微调 Stable Diffusion 网络：

• Aesthetic Score，一个在 Laion Aesthetic 数据集上训练的美学评估奖励函数；

• HPSv2 和 ImageReward，衡量指令跟随能力的奖励函数。

定性实验结果表明，通过 Nabla-GFlowNet 微调，我们可以快速得到奖励更高但避免过拟合的生成图像。

ReFL 和 DRaFT 等直接奖励优化的方法虽然收敛速度快，但很快会陷入过拟合；而 DDPO 这一基于传统策略梯度的强化学习微调方法由于没有理由梯度信息，微调速度显著劣于其他方法。

同时，我们的定量实验表明，我们的 Nabla-GFlowNet 可以更好保持生成样本的多样性。

Aesthetic Score 奖励函数上的微调结果（微调 200 步，取图片质量不坍塌的最好模型）。Nabla-GFlowNet（对应 Residual Nabla-DB 损失函数）方法微调的网络可以生成平均奖励更高且不失自然的生成图片。

                                      相较于 ReFL，DRaFT 等直接奖励优化的方法，Nabla-GFlowNet 更难陷入过拟合。

Aesthetic Score 奖励函数上的定量结果 Nabla-GFlowNet 在奖励收敛快的同时，保持更高的 DreamSim 多样性指标（越高代表多样性越好）和更低的 FID 分数（越低代表越符合预训练模型的先验）。

                                        HPSv2 奖励函数上的微调结果

                                       ImageReward 奖励函数上的微调结果

结语

我们利用生成流网络（GFlowNet）的框架，严谨地得到一个可以更好保持多样性和先验的高效的扩散模型奖励微调方法，并且在 Stable Diffusion 这一常用的文生图扩散模型上显示出相较于其他方法的优势。