Nabla-GFlowNet:兼顾扩散模型微调的多样性与效率

香港中文大学(深圳)等机构提出Nabla-GFlowNet,一种高效且能保留多样性的扩散模型微调方法,速度与质量兼得。

原文标题:扩散模型奖励微调新突破:Nabla-GFlowNet让多样性与效率兼得

原文作者:机器之心

冷月清谈:

本文介绍了一种名为Nabla-GFlowNet的新方法,用于解决扩散模型在奖励函数微调中面临的挑战。传统方法如强化学习收敛慢,直接最大化奖励函数则容易过拟合和丢失多样性。Nabla-GFlowNet基于生成流网络(GFlowNet),将扩散过程视为水流系统,并推导出新的流平衡条件。该方法通过对数流梯度参数化设计和残差梯度估计,实现了更高效且能保留生成样本多样性的微调。实验结果表明,在美学指标和指令跟随能力等奖励函数上,Nabla-GFlowNet相较于其他方法,能在快速提升奖励的同时,有效避免过拟合,并保持生成图像的多样性,从而在速度和质量之间取得平衡。

怜星夜思:

1、Nabla-GFlowNet 通过将扩散模型视为水流系统来解决微调问题,这个思路很有意思!除了视觉生成领域,这种“流”的视角还能应用到哪些其他领域?
2、文章中提到 Nabla-GFlowNet 在 Aesthetic Score 上进行了微调,并且能生成更自然的图片。除了 Aesthetic Score,还有哪些其他的奖励函数可以用来提升图像生成的质量和用户满意度?
3、文章提到 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 这一常用的文生图扩散模型上显示出相较于其他方法的优势。


© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:[email protected]

从“流”的角度看问题确实能打开新思路!我也来补充几个例子:

* 气候模拟: 可以把热量、水汽等在地球系统中的流动看作一个整体,用 GFlowNet 优化气候模型的参数,提高预测的准确性。
* 供应链管理: 产品和信息在供应链上的流动,可以借助 GFlowNet 优化库存管理、物流配送等环节,提高效率、降低成本。
* 能源网络: 电力、天然气等能源在网络中的流动,GFlowNet 可以用来优化能源分配,提高能源利用率。

我来抖个机灵,如果想生成更有“梗”的图,是不是可以把点赞数、转发数作为奖励函数? 这样训练出来的模型,岂不是能源源不断地生产表情包?:dog_face:

我从另一个角度来补充一下:

多样性可以看作是模型泛化能力的一种体现。一个泛化能力强的模型,不仅能够记住训练数据,还能理解数据的本质规律。因此,它能够生成更多样化的、符合规律的样本。

过拟合的模型,实际上是缺乏对数据本质规律的理解。它只是记住了训练数据中的一些特例,因此无法生成多样化的样本。

这个问题问得好!Aesthetic Score 只是一个起步,图像生成领域有很多其他有潜力的奖励函数:

1. CLIP Score: 利用 CLIP 模型 (Contrastive Language-Image Pre-training) 来衡量生成图像与给定文本描述的匹配程度。可以用来提升文生图模型的质量。
2. 用户点击率 (Click-Through Rate, CTR): 如果模型用于生成广告素材,可以直接用用户点击率作为奖励函数。这种方法能够更直接地优化用户参与度。
3. 图像质量评估指标 (如 PSNR, SSIM, LPIPS): 虽然这些指标更多用于图像复原任务,但也可以作为奖励函数,提升生成图像的清晰度、真实感。
4. 人脸识别准确率: 如果模型用于生成人脸图像,可以用人脸识别模型的准确率作为奖励函数,提高生成人脸的逼真度。
5. 主观评价 (Human Evaluation): 最直接的方法是让人来评价生成图像的质量。虽然成本较高,但可以得到最真实的反馈。

多样性在图像生成领域绝对是核心指标之一!原因很简单:

1. 提升用户体验: 如果模型只能生成几种相似的图片,用户很快就会感到厌倦。
2. 覆盖更广的应用场景: 多样性强的模型能够适应不同的输入和需求。
3. 缓解模型偏见: 如果训练数据不够全面,模型可能会学到一些错误的关联。多样性强的模型可以减轻这种偏见的影响。

过拟合和多样性缺失之间存在着密切的关联:

过拟合通常意味着模型过度关注训练数据中的噪声和细节,而忽略了数据的整体分布。这会导致模型在训练集上表现很好,但在测试集上表现很差。同时也会导致模型生成的结果缺乏多样性,因为模型只会生成训练集中出现过的样本。
简单来说,就是模型“背书”背得太死了,不会“举一反三”了。

楼上各位说的都很有道理!我再从一个更抽象的角度来补充一下:

GFlowNet 的核心思想是学习一个生成策略,使得生成结果的概率与奖励函数成正比。这种思想可以应用到任何需要优化生成过程的场景。关键在于如何定义合适的“流”和“奖励函数”。

例如,在软件工程中,可以把代码生成看作一个“流”,奖励函数可以是代码的质量、性能等指标。 GFlowNet 或许能帮助自动生成高质量的代码。

除了上面这些,可以考虑结合一些更高级的奖励函数:

* 对抗损失 (Adversarial Loss): 利用判别器来判断生成图像的真假。可以提升生成图像的真实感。
* 特征匹配损失 (Feature Matching Loss): 提取真实图像和生成图像的中间层特征,让它们尽可能相似。可以提升生成图像的细节和结构。
* 多样性奖励 (Diversity Reward): 鼓励模型生成更多样化的图像,避免坍塌到少数几种模式。

如果把图像生成比作画画,那多样性就像是画家的风格。一个只会模仿别人作品的画家,肯定没有自己的风格。只有那些理解了绘画本质,并且勇于创新的画家,才能画出独具特色的作品。

所以,模型也一样,只有具备了足够的“理解力”和“创新精神”,才能生成多样化的图像。

这个问题很有意思!把扩散模型看作水流系统确实是个巧妙的类比。我觉得除了视觉生成,这种“流”的视角在很多领域都有潜力:

1. 自然语言处理 (NLP): 可以将文本生成视为信息在句子中的流动。GFlowNet 或许能帮助生成更连贯、信息量更丰富的文本,尤其是在长文本生成或者对话系统中。
2. 分子生成/药物发现: 在药物设计中,可以把分子结构的构建看作原子或基团的流动过程。GFlowNet或许能帮助生成具有特定性质的分子结构,加速新药的研发。
3. 推荐系统: 用户的行为可以看作是在商品或内容之间的流动。GFlowNet 可以用来优化推荐策略,提升用户满意度。
4. 金融建模: 资金在市场中的流动也可以用类似的模型来分析,帮助预测市场趋势和风险。
5. 交通流量优化: 城市交通网络中的车辆流动,GFlowNet 可以用于实时交通优化,减少拥堵。