扩散模型奖励微调新突破：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 的平衡条件：

前向匹配损失：

后向匹配损失：

终端匹配损失：

对数流梯度参数化设计

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

伪代码实现

实验结果

我们分别用以下奖励函数微调 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 这一常用的文生图扩散模型上显示出相较于其他方法的优势。

文章来自微信公众号 “ 机器之心 ”