我自己的原文哦~ https://blog.51cto.com/whaosoft/12839828
#生成扩散策略算法归纳整理:优势与挑战
记得两年前刚投稿ICLR时,扩散策略(Diffusion Policy)还基本是一个纯理论的概念,全网只有寥寥两三篇arxiv,而现在它俨然已成为RL和领域的“显学”了。组里目前也在探索大规模扩散通用智能体的构建(RDT-1B)。最近闲了点,打算梳理一下近两年领域的理论进展,也算总结下自己研究的心路历程,做个宣传。
要回答的问题:
扩散策略究竟“好”在哪?
扩散模型引入RL,带来的本质挑战和核心难点(坑)是什么?
总结Diffusion Policy的经典RL优化算法。以我自己的工作为主线,尽量用统一的理论体系归纳相关工作。侧重high-level的理解。(可能更新后续系列细致讲解某一类算法)
(我内心中)各个算法、特别是自己提出的算法缺陷到底是什么?目前Diffusion+RL还有哪些Open problem?
文章前置知识:
- 有一定ML数学基础,比如熟悉KL散度的定义,大概知道forward/reverse KL的定义和区别,知道minimize KL和maximize likelihood是等价的。
- 了解Diffusion modeling,知道它本质上是在估计score function,能看懂Diffusion Loss function (不需要会推导),了解Diffusion Guidance (CG&CFG)。
- 有一丢丢RL基础,知道s,a,r,s',a' 这五个物理量的意思,大概知道Q函数,Advantage函数和V函数是什么。如果熟悉Offline RL背景会更好。
PS:对于领域的读者,需要说明本文研究的Diffusion Policy并非单纯模仿学习训练,这只是第一步,更加侧重如何给定Task Reward或Human Preference后进一步优化扩散策略。
扩散策略
一、扩散策略的优势/研究动机
目前“扩散策略优于高斯/确定性策略”似乎已经被当做某种常识了,大家无外乎有这么几点论据:
(1)Policy Diversity。如果一个任务有多个最优解,确定性高斯策略最多只能学会一个,扩散策略可能全学会。
RDT-1B[1] 双臂抓取任务中最优解多样性
(2)刷榜。在D4RL这种标准榜单上性能领先Gaussian Policy一个身位(DiffusionQL[2]),在领域尽管没有统一benchmark,但众多文章的ablation study也显示扩散优于高斯(DP[3])。
(3)不会“左右为难”。在模仿学习中如果既有向左走也有向右走的数据,确定性策略可能卡死不动,而扩散策略会学到“多峰”动作分布。(SfBC[4]; DP[3])
扩散策略不会"左右为难"
(4)Expressive Modeling Ability of Human Behaviors。扩散策略模拟人类复杂多样的行为数据,其拟合和建模能力肯定远优于高斯,上限更高。(Diffusion BC[5])
(5)兼容性。因为视频模型已经是基于Diffusion了,扩散策略方便与Diffusion World (video) Model无缝对接。目前也只有扩散模型适合做world model中一次性predict future sequence而不仅仅是predict one step的任务(Diffuser[6], Decision-Diffuser[7])
乍一看都非常有道理,Diffusion Policy 牛逼! 但这些论据合理性其实还是值得推敲的。
对(1),扩散策略天然具有多样性,但是传统RL任务以及任务中,“策略多样性”真的是我们的top priority吗?我们是希望模型稳定,可预测地完成任务,最大化奖励,还是希望它搞出各种花里胡哨的解决方案?再者,真的有很多任务存在多个“严格一样好”的控制策略吗?
对(2),D4RL榜单是典型只比较evaluation reward,不看policy diversity的榜单。难道只有扩散策略能够刷到很高的分数,传统高斯策略就不行吗?22/23年那会儿,我对(2)完全是持否定态度的,其实挺多传统算法后来也刷到了比较高的分数,但是没啥人关心罢了。这里面其实有很强的幸存者偏差:“扩散策略看起来有novelty,所以我要用;因为我要用,所以他就强”。(不过最近搞的同学不断给我反馈扩散策略就是明确强,我一直在琢磨是不是还有什么本质优势我没想清楚)
对(3),“左右为难”其实并不是高斯策略的问题,而是部分forward KL/Regression算法导致的mode-covering 的问题,换用有mode-seeking性质的reverse KL算法可以基本解决这个问题。典型反例就是SRPO[8]作为高斯策略,性能和Diffusion Policy持平,动作采样又远远快于扩散策略。后续文章会详细讨论。
SRPO算法 Forward/Reverse KL 对比
对(4),这个没啥说的。特别是在游戏AI上,我们更希望人机表现更“像人”,而不是完虐人类。可惜绝大部分游戏的action输入都是离散变量,不适合扩散建模。鼠标/轮盘移动除外,[5]这篇代表作也是在cs go游戏上搞得(鼠标)。
对(5),也没啥反驳的。不过这种结合world model + planning 的方法一般计算量都超级大,我本人涉及不多(后面讨论也较少有MBRL方法,尽管这方面也有很多有意思的工作)。
Diffuser (目前已知最早diffusion policy研究工作,学术思想非常超前)
当然无可否认的事,扩散策略能够短时间内收到如此广泛的关注,其本质优势(哪怕是incremental的),也是明确而不可忽视的。上面五条单独看某一条都可以辩论下,但放在一起就是强有力的论据了。审慎辩证地看待扩散策略优势是我们找出其问题所在,开展更进一步研究的前提。
二、扩散策略带来的本质挑战是什么?
有人说扩散策略是用生成式视角处理RL问题。但是对于做强化学习的人来讲,生成模型(如GAN, VAE, EBM)用于RL训练也不是一天两天的事儿了。在强化学习中把VAE换成扩散模型,是不是简单套用前人研究,做模块化替换,A+B的组合就可以呢?虽然有些这么做的工作,但这样其实很难触及和把握领域的前沿问题/方向。
- 似然估计问题
在扩散策略之前,整个强化学习的理论系统都是围绕着Explicit Likelihood Modeling 设计的,这在连续控制上反应为高斯模型,反应在离散动作(如LLM)上是分类模型(Categorical models)。以PPO为例
PPO loss 定义
不仅仅是RL算法,由于相同原因,LLM Alignment中经典的DPO算法以及诸多Preference Learning方法都无法直接迁移到扩散策略上。以DPO为例:
DPO 算法也需要显式计算模型Likelihood
现状:
1)上面的问题对于Offline RL并不太致命,22年左右扩散策略研究工作的核心就是在解决这个问题。其中SfBC[4]这篇文章首次明确提出了这个难点,Diffuser[6],CEP[9],Decision Diffuser[7] 三篇文章分别使用Energy Guidance和Classifier Free Guidance (CFG)手段对上面问题给出了比较漂亮的理论解析,Diffusion-QL[2]用基于ELBO的监督训练方法给出了一种实用的算法。重采样/引导采样/监督训练 三类手段基本也成为了目前Diffusion RL主流方法。
CEP[9] energy guided sampling 引导采样技术
2)然而,PPO/REINFORCE算法无法直接使用,对于Online RL影响巨大。我目前还没看到特别漂亮的解决方案。现有的方法基本上还是把Online Diffusion Policy 优化近似等于iterative offline RL去优化,也没有充分考虑online RL中特有的Exploration问题(扩散策略是可以多峰探索的!有巨大想象空间),感觉有点不本质。自己做的一些理论推导尝试也都不满意。目前我感觉还是Open Problem。
3)对于Diffusion Policy+Alignment问题。由于LLM Alignment发展的助推,最近这个方向的研究有很多,但应用并不在“扩散策略”上,而是在“图像生成模型上”,譬如DDPO[10],Diffusion-DPO[11]这类算法。然而,即使在图像上,目前也没有理论特别完美的算法,基本依赖推导中引入较多“约等于”符号,导致扩散与语言模型Alignment算法基本属于两个独立的体系。在扩散策略上基本不依赖近似,直接引入现有LLM Alignment策略的,我认为自己Nips24的工作EDA[12]算一个,但其也有本质问题,如对模型结构有特殊要求,有机会后面详细介绍。总体而言,我觉得这个领域处于active exploration状态。
- 采样效率问题
扩散模型需要多步采样,图像一般是15-100步,/RL中一般是5-50步。采样效率不仅仅是扩散策略,而是整个扩散模型领域的本质问题。但这个问题反应在控制任务中比在图像视频采样问题中解决需求更加急迫一些。
(1)控制频率问题。这个好理解,扩散模型一般参数要多一些,特别是RDT-1B这种大模型,叠加上多步采样,可能会很大程度上降低控制频率,直接影响“能用”还是“不能用”的问题。目前我了解到的解决方法有
a. 强行降低采样步数到5步,Diffusion-QL[2],RDT-1B[1]就是这样。在实践中,这些工作发现5步其实已经效果很好(足够)了,但是我对这个操作还是很担心,因为从ODE数值积分理论上推,5步采样几乎没有任何采样精度保证,所以这个系统效果有点看命,我自己的工作基本采用DPM-Solver采样,感觉最低也就是到10步,和图像上的经验一致。传统Diffusion基本就是10步图片能看,5步不太能看的水平。
b. 蒸馏。这类方法基本上就是跟着图像领域的发展走,都是先训练一个diffusion policy再进一步蒸成单步策略。我目前知道主流的方案一是Consistency model/policy,最早的工作是[13,14],二是score distillation,最早的工作是SRPO[8]。目前给我的感觉,方案一的主要问题就是CM本身的问题:基于bootstrapping的损失函数有训练有点不稳定、参数敏感。方案二的主要问题是,蒸馏出来的新模型损失掉了diversity,重新变成了Gaussian Policy。不过后续有一些新工作用Gaussian Mixture [15]或者GAN [16]的思路部分解决了这个问题。
c. 近似采样。这类工作更理论一些,基本就是用Diffusion Sampling过程中的一些数学性质,直接近似估计x_0而不走完完整的采样流程。代表工作是EDP[17],这也是我所知道最早专门针对扩散策略解决采样效率问题的工作。这类方法的问题和降低采样步数到5步方法的局限性比较类似,就是引入的“约等号”有点粗糙了,缺乏强理论保证,不知道能不能scale up。
(2)Q函数优化问题
扩散策略不仅仅拖慢inference速度,也拖慢训练速度,原因就是很多RL算法训练的时候依赖训练一个Q函数,而这个过程是依赖于实时从当前policy中采样的:
Diffusion-QL 算法中Q-Learning Loss
如果a要从扩散模型中采样,这个损失函数就会变得非常低效。更加要命的问题是,比如Diffusion-QL本身动作采样就是5步完成的,采样已经很不精确了,且扩散策略在训练初期,没收敛的模型采样出来的动作样本还有严重的数值问题(比如采出来个1e5这种离群点),现在采样不精确性造成的误差会进一步通过训练传导到Q函数中,Bootstrapping Q loss本身又不稳定。几个因素叠加,就回导致这种方法训练高度依赖调参,且在部分稀疏奖励任务上还容易炸掉。
这个问题说实在还没有太好的解决方法,我在后期的工作中都选用了IQL算法预训练一个Q函数,至少解耦policy learning和Q-learning,可以大大降低调参难度。
(插一嘴,尽管目前不得不用,我认为bootstrapping q learning的算法一定不是未来长远的rl算法,合理的q学习不应当是左脚踩右脚的耦合系统,我比较看好用类似dpo的方法直接学trajectory reward从而孤立出q function,可以参考我之前from r to q* 一篇blog,很希望在这块与有兴趣同学交流合作)
(3)模型优化问题
多步采样,也会导致扩散策略较难用DDPG这种梯度反传的方式来优化policy:
DDPG算法,带着策略模型的梯度信息生成的action,再传入Q-network,通过梯度反传优化模型
原因很简单,多步采样会使模型梯度层层累积,最终导致梯度爆炸问题。
然而不得不承认,这种方法在特定场景下也还能用,甚至有奇效。比如Diffusion-QL中,通过强行压低采样步数为5,模型又只使用3层的MLP解决了这个问题。这个方法在图像Alignment领域也是主流方法之一,然而图像至少要20步以上采样,直接使用此方案必然导致梯度爆炸,现在的主流算法是采样图片的前15步强行detach gradient,只在最后5步保留梯度信息。我反正觉得有点tricky,不够优雅。(:相关工作的作者别打我
(4)Online RL问题
Online RL不再是从静态数据集学习,要一边和环境交互一边学习。然而交互需要扩散策略产生动作,采样动作有精度误差又很慢。。。。。难受,研究不了一点儿(我想现在搞llm alignment 中online算法的研究者应该也是一样的感觉吧
三、Forward KL or Reverse KL?扩散优化算法理论体系
后续系列已经更新
陈陈:扩散策略算法归纳整理(二):前向KL优化理论
https://zhuanlan.zhihu.com/p/8615075797
相关文献
[1] ="https://arxiv.org/abs/2410.07864">RDT-1B: a Diffusion Foundation Model for Bimanual Manipulation 2024.10
#扩散模型驱动的四足机器人运动
离线学习+在线对齐
想象一下,一位经验丰富的登山者行走在蜿蜒的山路上,脚下的草地中混杂着碎石,身旁有大片茂密的灌木丛。或许连他自己都没有意识到,每一步行走时他的大脑都在飞速运转,判断下一个落脚点是否稳固,评估迈出去的步伐会不会失去平衡。对登山者而言,这不是简单的“左脚右脚”重复动作,而是像一场精密的舞蹈,融合了经验、直觉和环境的实时反馈,又稳又快地向前行走。
你是否会好奇,当四足机器人在崎岖不平的地面上以不同姿态行走时,它究竟是如何保持平稳的呢?实际上,它的控制系统就像这位登山者的“智慧大脑”一样,每一秒都在根据地面状况来精确计算每一个机械关节究竟需要“用多大的力”以“多大的速度”运转。这些计算和决策充满挑战,要求它在极短的时间内完成精准的判断,让每一步都踏得稳,跑得快,最终构成我们所看到的机械狗的流畅动作。
近期,中国电信人工智能研究院(TeleAI)研究科学家白辰甲博士与香港科技大学(广州)朱美新助理教授联合提出了一种名为“偏好对齐扩散规划器”的方法,通过将离线学习和在线偏好优化相结合,对扩散模型(Diffusion Model)驱动的四足机器人控制问题进行了深入研究,提出的两阶段学习框架结合了离线和在线学习的优势。首先,在离线阶段,扩散规划器通过学习状态-动作序列的联合分布,从其他策略收集的离线数据集中进行训练。随后,通过与环境的在线交互和偏好对齐技术,有效缓解了离线数据分布外(OOD)样本对策略影响,提升了控制算法的鲁棒性。这一算法不仅节省了传统在线强化学习的时间成本和运算资源,还显著提高了机器人的稳定性和速度跟踪能力。同时,实现了模拟与现实(sim-to-real)之间的零样本迁移(zero-shot transfer),使得机器人在不同地形下的表现更加出色。香港科技大学(广州)研究生袁欣怡和尚致违是论文的共同第一作者。
论文名称:Preference Aligned Diffusion Planner for Quadrupedal Locomotion Control
论文链接:https://arxiv.org/abs/2410.13586
项目地址:https://shangjaven.github.io/preference-aligned-diffusion-legged/
研究动机
构建具备通用性和高鲁棒性的四足机器人运动控制策略是机器人学中的一个长期挑战。随着扩散模型在大规模数据中表现出强大的生成能力,它在四足机器人控制任务中的应用变得愈发引人关注。扩散模型能够从复杂的状态-动作分布中学习,从而提供灵活的决策策略。然而,仅依赖离线数据训练的策略存在显著局限性,尤其在数据集状态覆盖不足的情况下,模型对样本外状态(OOD Samples)往往表现出较强的敏感性,难以在真实环境中有效应对多样的动态变化。此外,其他则研究通过收集在线轨迹来对原始策略进行更新,但是这类方法依赖于和专家策略进行实时交互或是对奖励进行在线采样,仍需要较高的计算资源。
本研究提出了一种创新的两阶段学习框架,巧妙地结合了离线策略学习与在线偏好优化,以提升四足机器人在不同地形下的稳健表现。在离线学习阶段,扩散规划器从其他专家策略收集的离线数据中学习状态-动作序列的联合分布,而这一过程并不依赖奖励标签。接下来的在线阶段,我们运用偏好对齐技术,来有效缓解样本外数据所带来的不良影响,进而提高策略的鲁棒性。本方法可以有效利用构建的弱偏好标签,以构建不依赖真实奖励信号的偏好数据集,从而在模型的偏好一致性上实现显著提升。最终,在在线优化阶段中,仅需少量的在线交互数据便可在多种步态下实现优异的稳定性和速度跟踪性能,并能在真实的四足机器人测试中展现出稳定性和优异的性能。
研究方法
本研究首先利用专家PPO策略在多步态策略学习框架“walk-these-ways”构建的仿真环境中收集离线数据集,涵盖了不同步态的状态和动作序列。在获得离线数据集后,通过行为克隆训练离线扩散模型的策略。在给定当前状态的前提下,扩散模型能够生成相应的状态和动作序列,以模仿离线数据集中的专家运动模式。随后,在初步构建的离线扩散规划器上进行偏好对齐,通过弱偏好标签对离线策略进行微调。偏好对齐阶段在不降低整体模型表现的情况下,使得较优轨迹的生成概率增加,以提高模型的稳定性和对分布外样本状态的适应性。通过这种策略微调,模型能够更好地适应真实环境中的变化,并提高任务的鲁棒性。最后,我们将优化后的策略部署到真实的Unitree Go1机器人上,进行真实环境中的测试。
图1:系统框架
离线策略学习
在第一阶段中,利用离线收集的四足机器人运动专家数据,通过扩散模型学习状态-动作序列的联合分布,形成初步的轨迹生成策略。在此过程中,扩散模型通过反复去噪,重构高维动作序列,从而捕捉机器人运动中的复杂动态特征。通过这种方式,扩散模型能有效建立在不同步态任务下的基础行为模式。
偏好标签构建
提出了一种“弱偏好”(weak preference)的偏好标签生成方法,用于在没有真实奖励信号或人工偏好标签的情况下,对离线扩散模型进行策略优化。首先从离线策略中生成(roll-out)新的数据集,并从中随机采样两个轨迹片段,并基于轨迹的状态-动作对与最优专家轨迹的距离计算来决定偏好,并将与专家轨迹更为接近的轨迹定义为胜出轨迹(winning trajectory)。
在线偏好优化
在偏好对齐阶段,我们使用的损失函数由偏好损失和正则化项两个主要部分构成。偏好损失致力于放大“胜出轨迹”和“失败轨迹”之间的差异。在偏好损失的指引下,策略在生成动作的时候,会更倾向于生成胜出轨迹而不是失败轨迹。优化后的策略能够更好地与得到的偏好标签一致,从而生成更多符合偏好要求的样本。
然而,一味的偏向于胜出轨迹可能会导致整体扩散策略表现的崩塌。因此损失函数中还包含一个正则化项,用于防止新策略与原始策略之间出现显著的偏离。正则化项能够极大程度缓解分布外样本带来的影响,从而确保整体策略的生成能力和鲁棒性。通过偏好损失和正则化项的结合,我们最终能够得到一个与偏好标签更一致、同时更稳定的扩散策略,新的策略能够在复杂多变的现实环境中展现出出色的性能。
实验结果
本研究在多种步态和速度条件下展现出优异的稳定性和速度跟踪能力,相比现有基线离线强化学习方法具有显著优势。研究在模拟器和真实环境中分别测试了三种步态(Trotting、Pacing、Bounding),并在慢速(0.5 m/s)和快速(1.0 m/s)两种速度下评估了模型的稳定性与横向速度控制精度。如下图所示,左图展示了三种步态在模拟器中的表现,右图为它们在真实机器人Unitree GO1上的应用效果。
在实验中,以稳定性和x轴平均速度作为评估指标。针对稳定性的评估,该研究对比了三种基线方法:CQL算法、采用U-Net为主干网络的DDPM算法(标记为DDPM-Unet)以及采用Transformer为主干网络的DDPM算法(标记为DDPM-Transformer)。实验结果如下图所示:
从图中可以看出,CQL算法在所有步态任务中表现不佳,机器人要么停滞不前,要么出现轻微抖动,难以应对复杂的控制任务。相比之下,采用偏好对齐的Diffusion算法,无论是在弱偏好还是强偏好条件下,都在所有步态任务中展现了出色的稳定性和速度跟踪能力。
针对速度追踪的效果,该研究使用x轴平均速度来进行评估,比较了不同Diffusion模型在Slow Bounding任务下速度追踪结果,如下图所示:
论文所提出的基于偏好对齐的Diffusion算法在目标速度的追踪上表现更加精准,而采用U-Net或Transformer架构的DDPM算法均难以准确实现目标速度的追踪。
更进一步的,为了分析方法的有效性,该研究还进行了消融实验,探索了偏好数据集规模、偏好标签质量及正则化对稳定性的影响。下图列出了不同条件下的消融结果:
实验结果表明,随着偏好数据集从1024增加到3072,稳定性显著提高。即便在数据集规模较小时,弱偏好标签的性能也接近于强偏好标签,这验证了弱偏好标注方法的有效性。此外,从图中还可以看出正则化方法在偏好对齐中的作用至关重要。未使用正则化项的实验均告失败,这表明正则化项在解决超出分布状态(OOD)问题时不可或缺。该论文的方法通过正则化增强了对OOD状态的适应能力,从而进一步提高了稳定性和速度控制性能。
最后,该研究在真实机器人Unitree GO1上进行部署,进一步验证了该方法在四足机器人运动控制中的潜力。
,时长00:48
研究总结
研究提出了一种结合离线学习与在线偏好对齐优化的两阶段Diffusion规划框架,旨在解决四足机器人运动控制中的复杂动作分布与分布外状态(OOD)问题。离线阶段利用专家数据学习状态-动作序列的分布,无需奖励标签,提升了数据效率;在线阶段通过基于欧氏距离的弱偏好标签进行策略优化,减少对外部反馈的依赖,并通过正则化避免策略漂移,增强了模型的稳定性与泛化能力。实验结果表明,该方法在不同步态与速度场景下的稳定性和速度控制精度显著优于基线算法,并成功实现了从模拟到现实的零样本迁移,为扩散模型驱动的四足机器人在复杂环境中的应用奠定了基础。
#Auto-Regressive Moving Diffusion Models for Time Series Forecasting
时间序列演进也是种扩散过程?基于移动自回归的时序扩散预测模型
在人工智能领域,目前有很多工作采用自回归方法来模拟或是替代扩散模型,其中视觉自回归建模(Visual AutoRegressive modeling,简称 VAR)就是其中的典型代表,该工作利用自回归取得了比传统基于噪声的扩散模型更好的生成效果,并且获得了人工智能顶级会议 NeurIPS 2024 的 best paper。然而在时间序列预测领域,当前主流的扩散方法还是传统的基于噪声的方法,未能充分利用自回归技术实现时间序列建模。
近期,来自上海交通大学和东方理工的研究团队推出了自回归移动扩散(Auto-Regressive Moving Diffusion,简称 ARMD)模型。ARMD 模型受经典的自回归移动平均(ARMA)理论启发,重新定义了时间序列的扩散过程:如图一所示,与传统的基于噪声的扩散方法不同,ARMD 将时间序列的演进看成是一个扩散过程,并采用了基于链式扩散的方式。它的前向过程由未来序列扩散到历史序列,并根据扩散步数通过滑动序列的操作生成中间状态;反向过程则通过历史序列‘采样’生成未来序列,实现了采样和预测目标的统一。
图一:已有的基于噪声的时序扩散方法和 ARMD 的时序扩散方法对比。
相比于传统基于扩散的时序模型,ARMD 并没有将时间序列预测视为从噪声起步的条件生成任务,而是充分利用时间序列数据的连续演化特性,实现了扩散机制与时间序列连续性之间的对齐,从而提升了模型性能。该项研究以《Auto-Regressive Moving Diffusion Models for Time Series Forecasting》为题,被 AAAI 2025 接收。
- 论文标题:Auto-Regressive Moving Diffusion Models for Time Series Forecasting
- 论文地址:https://arxiv.org/abs/2412.09328
- 项目地址:https://github.com/daxin007/ARMD
方法介绍
针对时序预测任务,该研究用历史序列
去预测未来序列
,其中 L 为历史序列的长度,T 为未来序列的长度。在该任务中,研究人员将历史序列的长度设置为与未来序列相同,即 L=T。
在 ARMD 中,时间序列的演进被概念化为一个扩散过程,时间序列的每一步都可以看成是扩散模型的一个状态,如图二所示。未来序列
(下标表示在序列中的位置,上标表示在扩散模型中的状态)作为前向扩散(演进)过程的初始状态,而历史序列
则是最终状态。不同于传统方法逐渐添加噪声生成中间状态,ARMD 通过对
进行滑动操作来生成中间状态(序列)
,使其逐渐接近历史序列。这一过程不仅保持了时间序列的连续性,还确保每个中间状态反映了时间序列演进的特定阶段。由初始状态到中间状态添加的内容可以定义为演进趋势,类似于传统扩散模型中的噪声。
反向去噪(逆演进)过程则利用历史序列
来迭代生成(预测)未来序列
。ARMD 采用了一个基于距离的去噪(逆演进)网络,根据中间状态
和时间步 t 得到时间序列的演进趋势。
图二: ARMD 扩散模型示意图。在前向扩散(演进)过程中,未来序列逐渐被转化为历史序列。相反,反向去噪(反演进)过程则利用历史序列迭代地生成 / 预测未来序列。
在最终的采样(预测阶段),ARMD 从历史序列
开始,迭代生成对未来序列的预测,使采样过程和最终的时间序列预测目标对齐,从而使 ARMD 成为无条件的时序扩散模型。采样过程遵循 DDIM(Denoising Diffusion Implicit Models)的方法,用预测的演进趋势替代传统扩散模型中预测的噪声,并通过逐步去噪(反演进)得到对模型的最终预测。
模型结果
经过在七个广泛使用的时序预测数据集上的大量实验,ARMD 模型的表现超越了现有的基于扩散的时间序列预测模型,并且和最先进的端到端时间序列预测模型取得了相当的效果。模型的预测结果如表一和表二所示:
表一:时间序列预测实验之与基于扩散的时间序列预测基准模型的结果比较。最佳结果以粗体突出显示。“最佳次数” 列表明模型取得最佳结果的次数。
表二:时间序列预测实验之与端对端时间序列预测基准模型的结果比较。最佳结果以粗体突出显示。“最佳次数” 列表明模型取得最佳结果的次数。
另外,由于扩散过程的优化,ARMD 模型在训练推理时间和预测的稳定性方面相对于之前的时序扩散模型也有明显的优势。ARMD 模型的预测效果示例如图三所示:ARMD 在具有明显周期性或趋势特征的时间序列上展现了卓越的预测性能,能够有效地捕捉这些时间序列中的相关特征,且相对于其他扩散模型预测更加稳定。
总结
综上所述,Auto-Regressive Moving Diffusion (ARMD) 模型通过重新诠释时间序列的扩散过程,结合时间序列的自回归特性,成功弥合了扩散机制与时间序列预测目标之间的差距。不同于传统方法从白噪声开始逐步添加噪声,ARMD 采用链式扩散策略,通过滑动操作将未来序列逐渐转化为历史序列,更准确地模拟时间序列的演化路径。该模型还引入了独特的反向去噪(逆演化)过程,利用历史序列迭代生成(预测)未来序列,有效捕捉时间序列的特征。实验结果和理论分析共同证明了 ARMD 作为一种稳健且有效的时间序列预测方案的巨大潜力。
图三: ARMD 和 Diffusion-TS 在给定相同历史序列的情况下,分别做出的 10 次不同预测的分布情况。ARMD 实现了更稳定且准确的预测。
#轻微的预训练数据扰动如何造就更优秀的扩散模型
适度的噪声(如2.5%-7.5%的条件误配对)不仅不会显著负面影响模型性能,反而能帮助生成更高质量和多样性的内容。文章通过实验验证了这一现象,并提出了一种创新的扰动方法——条件嵌入扰动(CEP),可以在不直接修改数据的情况下实现类似的效果。
在人工智能领域,基础模型(Foundation Models)的预训练数据质量一直被认为至关重要。我们最新的研究《Slight Corruption in Pre-training Data Makes Better Diffusion Models》(轻微的预训练数据扰动如何造就更优秀的扩散模型),第一次全面系统地探讨了这些扰动对模型性能的影响,并提出了一种创新的方法来将数据的“瑕疵”转化为模型性能提升的“助力”。本文延续我们之前的工作[1],发现预训练数据中适度的噪声也可以提升模型预训练和下游性能。这项研究成果又卡耐基梅隆大学,香港大学,RIKEN等共同发表于NeurIPS 2024,并获得了spotlight论文荣誉。
- 论文标题:Slight Corruption Makes Better Diffusion Models
- 论文链接:https://arxiv.org/pdf/2405.20494
- 论文模型及代码:https://huggingface.co/DiffusionNoise
研究背景
扩散模型(Diffusion Models)近年来在生成高质量图像、音频和视频方面取得了令人瞩目的成功,其背后的一大推动力是大规模的预训练数据。例如,著名的Stable Diffusion模型依赖于从互联网上爬取的数十亿对图文数据(如LAION-2B),这些数据不仅为模型提供了丰富的训练语料,还帮助其实现了高度多样化和条件化的生成能力。这种方法显著提升了扩散模型的生成效果,使其在艺术创作、内容生成等任务中表现出色。
然而,大规模数据的收集和处理并非毫无挑战。尽管这些数据集通常会经过严格的筛选和过滤,但误配对样本(condition corruption)依然不可避免,例如错误的图像标签、文本描述不匹配、甚至无意义的内容。这些误配对样本被认为是大规模数据不可避免的“副产品”,传统观点通常认为它们会削弱模型性能,增加生成的噪声和不一致性。
然而,我们的研究挑战了这一传统认知。通过系统性实验,我们发现,轻微的预训练数据扰动(例如2.5%-7.5%的条件误配对)不仅不会对模型性能造成显著负面影响,反而能够帮助扩散模型生成更高质量、更具多样性的内容。令人惊讶的是,这一发现与我们此前在CLIP和ImageNet分类器的研究结论高度一致:轻微的标签噪声在特定情况下可以帮助模型在域内任务(In-Domain Tasks)中提升性能[1]。
这一观察背后的潜在原因可能是,轻微的噪声引导模型利用更多的数据维度来学习特征分布,从而增加生成数据的多样性和质量。同时,噪声还能避免模型过度拟合数据中的偏差或冗余模式,使生成的分布更贴近真实数据分布。这一发现不仅为扩散模型的优化提供了新的视角,也对理解大规模数据预训练过程中的数据偏差提出了重要见解。
随着生成模型的广泛应用,我们的研究不仅为扩散模型的设计和训练带来了新思路,还为开发更健壮、更通用的生成系统提供了理论基础和实践方法。
实验发现
为了全面研究预训练数据扰动对扩散模型的影响,我们设计了系统性实验,涵盖预训练和下游任务两个部分。以下是实验的具体设置以及主要发现。
实验设置
模型:我们采用了多种扩散模型,包括Latent Diffusion Models (LDMs)、Diffusion Transformers (DiT-XL/2) 和 Latent Consistency Models (LCMs),分别对类条件(class-conditional)和文本条件(text-conditional)的生成任务进行实验:
- LDM-4: Latent Diffusion Model,基于VQ-VAE编码的潜空间(down-sampling factor为4),图像分辨率为256×256。
- DiT-XL/2: 使用Transformer作为主干网络的扩散模型。
- LCM-v1.5: 针对快速推理优化的Latent Consistency Models,基于Stable Diffusion v1.5
预训练数据:我们在两个广泛使用的数据集上引入了合成扰动:
- ImageNet-1K (IN-1K): 为类条件生成任务引入标签翻转扰动(随机将标签替换为其他类别标签)。
- CC3M: 为文本条件生成任务引入文本替换扰动(随机交换图文对中的文本)。
扰动比例(corruption ratio, η)分别设置为0%、2.5%、5%、7.5%、10%、15%、20%。
预训练评价指标:
- 生成质量和保真度:使用Fréchet Inception Distance (FID) 和 Inception Score (IS)。
- 多样性和覆盖率:使用Precision-Recall (PR) 曲线和样本熵值(Sample Entropy)。
- 我们在Appendix里提供了其他metric的测试结果,如sFID, Memorization Ratio, RMD (sample complexity), L2 distance.
下游任务实验设置:预训练之后,我们研究了扩散模型在两个下游个性化任务中的表现:
- ControlNet:基于空间条件(如Canny边缘图和语义分割掩码)控制生成。
- T2I-Adapter:在文本条件生成中结合额外的输入模态以提高可控性。数据集:
我们选用ImageNet-100(IN-100)作为个性化任务的数据集,并使用OpenCV Canny检测器和SegmentAnything(SAM)生成边缘和分割掩码作为输入条件。
下游任务评价指标:与预训练一致,包括FID、IS、Precision 和 Recall。
Pre-training结果
- 轻微扰动提升生成质量和多样性
- 更高的质量:通过在不同程度扰动的条件下训练模型,我们发现2.5%-7.5%的扰动比例显著提高了生成图像的质量。例如,在IN-1K数据集上的类条件生成任务中,2.5%扰动的LDM-4模型,其FID从未扰动模型的9.44降至7.00。这表明轻微扰动帮助模型更好地捕捉特征分布,生成更接近真实分布的样本。
- 更丰富的多样性:轻微扰动同样显著提升了生成样本的多样性。以文本条件生成任务(CC3M数据集)为例,7.5%扰动的模型在CLIP Score上从19.85提升至17.93,同时熵值(Sample Entropy)和相对马氏距离(Relative Mahalanobis Distance, RMD)也有所增加,反映出生成分布的覆盖范围更广。
- 扰动程度的非线性效应
- 提升-退化:随着扰动比例从0%逐步增加到20%,生成质量和多样性呈现出“先提升后下降”的非线性趋势。例如,在IN-1K数据集中,FID在扰动比例达到7.5%时达到最优(7.00),而当扰动比例增加到15%时,FID略有上升至9.18。这表明,过高的扰动可能会干扰模型捕捉有用的特征。
- 偏向于低噪声的设置:与以往研究的高噪声标签学习不同,我们的结果表明低比例(2.5%-7.5%)的扰动不仅是可行的,还可以成为优化生成模型性能的有效策略。
- 跨模型一致性
- 轻微扰动的效应在多种扩散模型上得到了验证,包括LDM-4、DiT-XL/2和LCM-v1.5。这表明这种方法具有广泛适用性。即使在Transformer架构,如DiT-XL/2中,轻微扰动也能提高生成能力。
- 视觉表现的显著改进
- 我们进一步通过生成样本的视觉评估验证了轻微扰动的优越性。例如,在文本条件生成任务中,7.5%扰动的模型生成的图像更加逼真且多样化,展示了更丰富的细节和场景变化
Downstream结果
- 轻微扰动增强个性化生成性能
- ControlNet实验:在基于Canny边缘图和SAM分割掩码的个性化任务中,轻微扰动的预训练模型表现优异。这表明,预训练中的轻微扰动为下游任务提供了更通用和更具适应性的特征表示。
- T2I-Adapter实验:文本条件生成任务中,预训练扰动的优势也更加明显。
- 增强模型的多样性与保真度平衡
- 在下游任务中,我们观察到轻微扰动的模型不仅在FID指标上表现更优,还在Precision-Recall曲线上表现出更好的平衡性。例如,在IN-100的ControlNet实验中,2.5%扰动的模型在Precision和Recall上分别达到0.84和0.58,显著优于未扰动模型的0.82和0.61。这表明轻微扰动有效改善了模型生成的覆盖率,同时维持了较高的精确性。
- 下游任务的视觉评估改进
- 我们通过对生成图像的视觉质量进行对比,发现轻微扰动的模型能够生成更加自然且多样化的图像。
- 高效适配的能力
- 我们发现轻微扰动的预训练模型在进行下游任务的微调时具有更高的适应性。例如,与未扰动模型相比,轻微扰动的模型在少量标注数据的情况下,仍能实现显著的性能提升。这说明轻微扰动不仅能提升预训练阶段的表现,还能为下游任务提供更具泛化性的初始权重。
理论分析
我们通过理论分析揭示了轻微数据扰动如何在生成质量和多样性上影响扩散模型的性能。本部分分为两部分:生成多样性的提升和生成质量的改善,并结合我们的数学建模和分析,提供直观的解释和数学依据。
生成多样性:轻微扰动扩展生成分布
核心观察:
轻微扰动能有效提升生成样本的多样性,使生成分布覆盖更多的数据空间。我们通过计算生成分布的熵值(Entropy)作为多样性的度量,高熵表示生成分布更广泛地覆盖数据空间,而低熵则意味着生成分布过于集中。
理论证明:
在生成多样性方面,我们提出并证明了以下定理:
定理 1:
对于任意类别 ,在生成长度 足够大的情况下,若训练数据的经验协方差矩阵为满秩,且类别的期望范数 为常数,则有:
- 和 分别为带有扰动和无扰动条件的生成分布的熵值;
- 为扰动控制参数;
- 为数据维度。
这一结论表明,条件扰动能够增加生成分布的熵值,提升生成样本的多样性。此外,熵值的提升程度与扰动大小 和数据维度 成正比。这也与我们实验中发现的结果一致,即轻微扰动( 小且适中)能够提升多样性,而过大的扰动会导致分布失衡。
我们在不同扰动比例下观察到样本熵值的变化,如IN-1K和CC3M数据集中,2.5%-7.5%扰动的模型的熵值显著高于未扰动的模型,进一步验证了理论结论。
生成质量:轻微扰动优化分布对齐
核心观察:
轻微扰动能够改善生成分布与真实数据分布的对齐程度,进而提升生成样本的质量。我们采用2-Wasserstein距离作为衡量分布对齐程度的指标,距离越小,生成分布越接近真实数据分布,生成样本质量越高。
理论证明:
在生成质量方面,我们提出并证明了以下定理:
定理 2:
对于任意类别 ,在生成长度 足够大的情况下,若类别的期望范数 为常数,假设 ,则有:
其中:
- 为真实分布, 和 分别为无扰动和有扰动条件下的生成分布;
- 为2-Wasserstein距离;
- 为类别 的样本数量;
- 为数据维度。
这一结果表明,轻微的条件扰动能够使生成分布更接近真实分布,即生成质量有所提升。此外,当取适当的值时(如实验中的2.5%-7.5%扰动),生成分布的对齐效果最为显著,进一步验证了我们的实验结论。
在IN-1K和CC3M数据集中,我们观察到,随着轻微扰动的引入,FID显著下降,表明生成分布与真实分布的对齐程度得到了改善。
Conditional Embedding Pertubation
轻微扰动在预训练数据中展现出提升模型性能的潜力,但直接操作预训练数据可能并不总是现实可行。为此,我们提出了一种创新的扰动方法,称为Conditional Embedding Perturbation (CEP) ,通过对条件嵌入层(conditional embedding layer)直接添加随机扰动,模拟数据中的轻微噪声,从而提升模型的生成质量和多样性。
核心思想
CEP的核心思想是,在预训练和下游任务的训练过程中,对模型的条件嵌入(例如文本嵌入或类别嵌入)添加轻微扰动,而无需修改底层的数据或模型结构。具体而言,假设条件嵌入为,我们将其调整为:
其中:
- 是服从正态分布或均匀分布的随机噪声;
- 是控制扰动幅度的超参数。。
通过在条件嵌入中加入轻微扰动,CEP可以有效模拟预训练数据中的轻微噪声,同时避免了直接修改数据可能带来的额外开销或不一致性。
实验验证
我们在多个生成任务和数据集上评估了CEP的有效性,包括IN-1K的类条件生成任务和CC3M的文本条件生成任务。以下是CEP的具体表现:
- 生成质量和生成多样性的提升: 在IN-1K数据集上,使用CEP的模型(扰动幅度 (\gamma=0.05))的FID从未加扰动模型的9.44显著降低,显示了更高的生成质量。Precision-Recall曲线显示,CEP不仅提升了生成的多样性(Recall),还维持了高精确度(Precision)。
- 对不同模型的通用性: CEP在LDM、DiT和LCM等多种扩散模型中均表现出一致的性能提升,证明了该方法的通用性。
- 例如,在LDM-4模型上,CEP将FID从9.44降至7.01,而在DiT-XL/2模型上,CEP将FID从6.76降至5.51。
- 对下游任务的适用性: 在下游任务中,CEP对生成性能的提升尤为显著。例如,在基于ControlNet的个性化生成任务中,使用CEP预训练的模型,其FID从11.59进一步降低至11.14,同时多样性指标(熵值和Recall)也显著提高。类似地,在T2I-Adapter任务中,CEP的模型生成的图像更加逼真,且对额外的条件输入(如边缘图或分割掩码)的响应更精确。
未来展望
随着生成模型和预训练数据规模的不断增长,理解数据扰动对模型性能的影响变得尤为重要。我们的理论分析不仅揭示了轻微扰动的潜在优势,也为未来研究更复杂的数据偏差和噪声类型奠定了基础。此外,我们相信,这种分析方法还可以推广到其他生成模型(如GAN和自回归模型),从而为生成任务的优化提供更广泛的指导。CEP提供了一种高效且灵活的方式来模拟和利用数据中的轻微噪声。未来,我们计划将CEP推广到更多类型的生成任务中,例如对话生成、语言生成,以及多模态生成任务。同时,我们还将探索如何动态调整使CEP能够适应不同数据分布和任务需求,进一步提升模型性能和鲁棒性。
#DiC
纯卷积实现用于图像生成扩散模型?重新思考扩散模型中的 3×3 卷积
重新思考了 3×3 卷积这种最简单但是最快的模块,来实现基于纯卷积的扩散模型。
纯卷积实现且用于图像生成扩散模型。
最近的一些工作把扩散模型从传统的 U 型的 CNN-Attention 的混合结构转变为完全基于 Transformer 的架构。虽然这些模型表现出了很强的扩散性和强大的性能,但是其复杂的 Self-Attention 操作会导致推理速度变慢。
本文重新思考了 3×3 卷积这种最简单但是最快的模块,来实现基于纯卷积的扩散模型。
具体而言,作者发现 Encoder-Decoder Hourglass 的设计优于 Isotropic 的架构,但仍然表现不佳,低于原来的预期。作者引入了 Sparse Skip Connection 来减少冗余并提高模型的可扩展性。基于这个架构,本文进一步提出 Stage-Specific Embeddings,Mid-Block Condition Injection, 和 Conditional Gating 的技术,得到最终的 Diffusion CNN (DiC)。实验表明,DiC 在性能方面大大超过了现有的 Diffusion Transformer,同时保持了良好的速度优势。
图1:本文 3×3 卷积实现的扩散模型的 Roadmap。形式是 200K 训练步数的 FID 结果
1 DiC:重新思考扩散模型中的 3×3 卷积
论文名称:DiC: Rethinking Conv3x3 Designs in Diffusion Models
论文地址:http:///arxiv.org/pdf/2501.00603
项目主页:http://github.com/YuchuanTian/DiC
1.1 DiC 研究背景
Self-Attention 机制在最近的许多扩散模型中起重要的作用。早期的扩散模型的工作[1]将卷积 U-Net 结构与 Self-Attention 融合。U-ViT[2]和 DiT[3]等架构完全过渡到基于 Self-Attention 的设计,完全放弃了卷积 U-Net。这些模型展示了显著的生成能力,尤其是在扩大规模的时候。
较新的文生图或者文生视频模型,包括 Stable Diffusion 3[4]、FLUX[5]、PixArt[6]和 OpenSora[7],也采用了完全基于 Transformer 架构,实现了很高的生成质量。
不幸的是,在扩大规模时,结合 Self-Attention 的扩散模型面临着显着的计算开销和延迟。复杂的 Transformer 架构对实时、资源受限的应用程序造成了很大的挑战,其中大规模 Transformer 模型由于巨大的时间成本而变得不切实际。
所以很多工作会加速扩散模型中的 Self-Attention,比如 PixArt-Σ[6]实现了高效的 Self-Attention 机制。但是,这种方法仍局限于 Self-Attention 的范式。生成图片的速度仍然有提升的空间,未能满足实时或大规模应用的需求。
在这项工作中,本文重新思考了简单有效的完全基于卷积的扩散模型,该架构提供了显著的速度优势。而卷积作为 Self-Attention 的对应,被认为硬件友好,在各种平台上都有很好的支持。在众多的卷积操作中,规范的 3×3 Stride1 卷积因其卓越的速度而脱颖而出,这主要是由于 Winograd 加速[8]等广泛的硬件优化技术。例如,RepVGG[9]等工作已经证明了简单的卷积如何实现高效率和具有竞争力的性能。
但是,3×3 卷积有它固有的局限性,尤其是感受野受约束,可能会阻碍其在复杂生成任务中的可扩展性。作者首先进行了一些初始的实验,在现有的架构上面扩展基于 3×3 卷积的 CNN,但是发现这样得到的结果仍然不如对应的 Diffusion Transformer。
为了解决这个问题,作者对传统的 ConvNets 进行了一系列适配,专门为了缩放扩散模型进行了一系列定制化设计,形成了图 1 所示的 Roadmap。首先,作者专注于改进模型架构以更好地利用 3×3 卷积。现有的设计分为 3 类:Isotropic 架构,如 DiT;具有 skip connection 的 Isotropic 架构,如 U-ViT;以及 CNN 和 Self-Attention 混合模型中使用的 Encoder-Decoder Hourglass 结构。实验结果表明,Hourglass 结构对纯 3×3 卷积模型更有效,因为编码器的下采样和上采样扩大了感受野,弥补了 3×3 卷积感受野窄的问题。Skip connection 也至关重要,但作者认为传统的 Skip connection 的做法过于密集,无法对大 ConvNet 有效。为了解决这个问题,作者引入了 Sparse Skip Connection,减少了 skip 的数量,同时确保基本信息流从 Encoder 有效流向 Decoder。
除此之外,作者专注于改进 Condition,以更好地拟合 ConvNet。在现有模型中,一组 Condition Embedding 通常映射到不同阶段的不同 Block。然而,在 Encoder-Decoder ConvNets 中,每个 Stage 在不同的特征空间中运行。因此,作者引入了 Stage-Specific Embedding,确保每个 Stage 使用独立的、不重叠的 Condition Embedding 表示。作者还仔细检查了 Condition 注入的位置,建议在中间的 Block 中注入 Condition 以获得更好的性能。此外,作者对内部块结构进行了改进,将所有激活函数替换为 GELU。
通过上面的改进最终得到了 DiC,一种由 Conv3×3 组成的扩散模型,以上的改进使之能够在保持速度优势的同时取得出色的结果。
1.2 使用 3×3 卷积作为模型基本组成
作者选择正常的 Stride-1 full 3×3 卷积作为模型的主要成分。Stride-1 3×3 卷积的速度令人难以置信,主要是由于现代深度学习框架中广泛的硬件和算法优化。操作优化受益于 Winograd 算法[8],将常见卷积操作所需的乘法次数减少了 5/9 。对于 3×3 卷积核,该方法通过利用有效的矩阵变换显着加快了计算速度。
而且,与 Depthwise 卷积不同,3×3 卷积实现了高度的计算并行性,允许现代 GPU 充分利用它们的处理能力。这种并行性与最小的显存访问开销相结合,使得 stride-1 3×3 卷积特别有效。与更大或更复杂的卷积操作相比,3×3 卷积在计算成本和表征能力之间提供了理想的权衡,使其成为实时应用和大规模模型的实用选择。
1.3 架构探索:使用 Encoder-Decoder Hourglass 架构
已有的工作[10]也有基于 3×3 卷积的扩散模型。但是,这些模型并非是完全基于 CNN 的:模型中的 Basic Block 由两个连续的 3×3 卷积组成,然后是关注整个图像的 Self-Attention 模块。在每个卷积之前,使用 GroupNorm 对输入进行归一化,然后传递给 SiLU 的非线性激活函数。整体采用残差结构。
本文提出的 DiC 的 Roadmap 从上面提到的 Basic Block 开始。为了满足路线图的最终目标,作者在卷积后去除Self-Attention 层,形成一个由 2 个 3×3 卷积组成的 Basic Block,其他块设计暂时保持不变,如图 2 所示。
图2:本文基于 3×3 卷积的 Basic Block。3×3 卷积构建的 Basic Block 更简单,硬件支持更好,是高吞吐量的关键
目前,基于 3×3 卷积实现的有 3 种主流的架构,分别是:
图3:主流扩散模型架构。使用 3×3 卷积的 Basic Block 的现有架构 (a,b,c),用 Sparse skip connection 改进 U-Net (d)
Isotropic 架构 (以 DiT 为代表): 简单、垂直堆叠。Patchify 之后,整个模型中不会改变中间特征图的大小。它因其公认的可扩展性而被广泛采用。
具有 Skip Connection 的 Isotropic 架构 (以 U-ViT 为代表): 通过引入非相邻层之间的长跳跃连接。与传统的U-Net不同,尽管使用 Skip Connection,但该结构在整个网络中保持了一致的空间分辨率,避免了任何上采样或者下采样操作。
U-Net Hourglass 架构 (以 ADM 为代表): 这个经典的设计[11]由一个编码器和一个解码器组成,形成一个 Hourglass 的形状。Encoder 逐步对特征维度进行下采样,而 Decoder 对称地对空间维度进行上采样。除了 Hourglass 主干之外,Encoder 中的每个 Block 都通过 Skip Connection 连接到其对应的 Decoder Block。
实验分析
使用基本的 Conv 3×3 Basic Block,作者对上述架构选择进行了实验,结果总结在图 4 中。这些实验揭示了几种模式,为设计有效的基于 3×3 卷积的扩散模型提供了有价值的见解。
首先,对于仅由 3×3 卷积组成的扩散模型,Encoder-Decoder Hourglass 架构是必不可少的。这可能由于 3×3 卷积受限的感受野。在 Isotropic 的标准 Transformer 架构的设置中,每个 3×3 卷积在每个方向上仅扩展一个像素的感受野。因此,实现全局感受野将需要深度的卷积堆叠,导致计算和参数效率低下。
相比之下,Encoder-Decoder 架构显著地降低了这个限制。通过 Decoder 中的上采样和 Encoder 中的下采样操作,较高 stage 的 3×3 卷积可以有效地感知较大区域,大大增加了模型的感受野。这种分层方法允许模型有效地捕获本地和全局上下文,这使得它对生成任务更有效。
而且,Skip Connection 在基于 Conv 3×3 的架构中至关重要。一方面,在 Encoder-Decoder 结构中,它们通过将特征图直接传递给 Decoder 来减轻下采样过程中的信息损失,丰富了其表示。另一方面,Skip Connection 通过提供额外的梯度路径并保留重要特征来加速训练并提高生成质量,从而实现更高效和有效的扩散建模。
但是,缩放基于 Conv 3×3 的扩散模型需要堆叠很多层,使得密集的 Skip Connection 成为瓶颈。因为 Skip Connection 的特征需要在 channel 维度拼接起来,这样会消耗过多的计算成本。这样的低效率的做法阻碍了模型的可扩展性。为了解决这个问题,本文提出跨步 Skip Connection,其中 Skip Connection 每隔几个块用一下,而不是所有的块都使用。
图4:ImageNet 256×256 条件生成上比较不同的 Conv3×3 Scale-Up 架构。U-Net Hourglass 架构最能释放 Conv3×3 架构的潜力,稀疏 Skip Connection 进一步提升模型
如图 4 所示,实验结果表明,与密集连接的跳过设计相比,稀疏 Skip Connection 得到了更好的结果。
1.4 Conditioning 方面的改进
作者认为模型缩放知乎,Conditioning 方面也需要进一步改进。
Stage-Specific Embedding
在传统扩散模型中,以前的模型共享一个唯一的 Embedding。Conv3×3 模型依靠 Encoder-Decoder 结构来扩展感受野。Encoder-Decoder 结构的每个 stage 都使用不同的 channel 进行操作,反映了 stage 之间显著的结构和功能差异。鉴于特征维度因 stage 而异,在不同阶段的所有块上均匀地应用相同的 Condition Embedding 可能不是理想的。单个 Embedding 可能无法充分捕获每个 stage 的不同特征,因此对不同 stage 定制 Condition Embedding 可能更有效。
作者提出使用 Stage-Specific 的 Condition Embedding,对每个 stage 使用不一样的 Embedding,Embedding 的维度与这个 stage 的特征维度对齐。值得一提的是,作者分析了 Stage-Specific 的 Embedding 引入的开销。结果显示增加的开销很小。只会令模型大小增加了 14.06M 参数,仅占总模型大小的 2%。引入的计算开销仅增加了 12M FLOPs,对整体计算效率的影响很小。
Condition 注入的位置
另一个重要的点是确定注入 Condition 信息的最佳位置。扩散模型中有两种流行的条件注入策略。第一种方法在每个 Block 的一开始就注入条件,通常通过 LayerNorm,如 DiT 等模型。第二种方法在 Block 中间引入条件,如 ADM。
通过实验,作者发现对于全卷积扩散架构,将条件注入每个 Block 内的第 2 个卷积层会产生最佳性能。此放置有效地调制特征表示,在不影响其效率的情况下增强模型的生成质量。
Conditional Gating
作者还采用了 DiT 的 AdaLN 的条件门控机制。AdaLN 引入了一个门控向量,该向量沿通道维度扩展特征,为调节过程提供更动态和更细粒度的控制。这种改进增强了模型适应不同条件的能力,提高了整体生成质量。本文的实验表明,这种增强产生了显着的性能提升,进一步证明了这种调节策略在扩散模型中的有效性。
其他小改动
作者重新审视整个模型中的激活函数。受 ConvNeXt 的启发,将常用的 SiLU 激活替换为 GELU。这种变化也改进了性能。
图5:ImageNet 256×256 生成任务的改进。Conditioning 方面的改进的实验结果
1.5 实验设置
为了针对现有的 DiT 模型进行基准测试,我们设计了我们不同尺度的架构:Small (S)、Big (B) 和 Extra Large (XL),它在 FLOP 和参数计数方面与 DiT-S/2、DiT-B/2 和 DiT-XL/2 对齐,如图 6 所示。此外,作者引入了一个更大的模型 Huge (H),用于缩放性能的评估。DiC-S 提供了一种轻量级的解决方案,使用 32.8M 参数和 5.9G FLOPs,类似于 DiT-S/2,具有高效的 Encoder-Decoder 配置。DiC-H 旨在评估 DiC 模型的潜力。
除了正常的 FLOPs 计算外,作者还记录了考虑了 Winograd 之后的实际 FLOPs。Winograd 可以将 Stride-1 3×3 卷积的计算减少 5/9。可以观察到 Winograd 可以将所提出的 DiC 的理论 FLOPs 减少大约一半。
图6:不同模型尺寸的 DiC 架构配置
在实验中使用 256 的 global batch size,学习率为 1e-4,weight decay 为 0。
1.6 实验结果
如图 7 所示,DiC 模型在所有测试尺度上都优于相应的 DiT 模型。具体来说,DiC-S 的 FID 显着降低,从 67.40 下降到 58.68,同时在 IS 上也优于 DiT-S/2,从 20.44 增加到 25.82。同样,DiC-B 的性能明显优于 DiT-B/2,FID 从 42.84 降低到 32.33,并将 IS 从 33.66 增加到 48.72。
图7:DiC 与 DiT 在各种模型大小上进行比较。经过 400K 次迭代训练的模型与 DiT 的标准训练超参数进行比较。所有尺度的DiC 在相当大的范围内优于 DiT
在大模型中,DiC-XL 比 DiT-XL/2 提供了显著的改进,FID 从 20.05 下降到 13.11,IS 从 66.74 下降到 11.15。这些结果突出了我们的 DiC 模型的有效性,它在图像质量和多样性方面都始终取得卓越的性能,同时保持计算效率。
与 Diffusion Transformer 的比较
由于不同的模型使用不同的设置,包括训练超参数、采样器选择、训练迭代次数等,作者根据 U-DiT[12]采用普遍对齐的设置 (DiT 代码库上的 400K 次迭代)。如图 8 所示,DiC 模型在 ImageNet 256×256 上始终优于基线扩散架构。
图8:ImageNet 256×256 生成上的性能
除了性能之外,DiC 的速度很快。尽管扩散模型的计算量很大,但 DiC-XL 的吞吐量为 313.7,反映了 Conv3×3 高度优化的操作的使用。这种架构利用了卷积的固有速度特点,在效率和质量之间提供了很好的平衡。
DiC-H 在模型大小和速度之间取得了显着的权衡,将最先进的 FID 为 11.36,IS 为 106.52,吞吐量提高了 160.8。实际应用中,这些速度优势使 DiC 成为高质量图像合成的实用解决方案。
此外,作者在图 9 中试验了 DiC 模型的条件生成能力。DiC-XL 仅在 400K 次迭代时实现了 3.89 的 FID,大大超过了具有竞争力的 Diffusion Transformer 基线,结果来自具有 cfg=1.5 的条件生成样本。
图9:使用 CFG 的结果。DiC 在条件生成方面表现出色
图 10 报告了 DiC 模型与 DiT 在 ImageNet 512×512 上的比较结果。关键结论是 DiC 模型的计算效率。DiT-XL/2 需要 524.7G FLOPs (使用 Winograd 优化),DiC-XL 模型在 FID 和 IS 方面表现明显更好,只有 464.3G FLOPs (Winograd 优化后 228.7G)。在 DiC-H 中,改进变得更加明显,显著降低了 FID (12.89),并实现了比 DiT-XL/2 更少的 FLOPs 和更高的 IS (101.78)。
DiC 和 DiT 计算复杂度的关键区别在于它们的底层架构。DiC 使用纯卷积方法,随图像大小线性缩放。相比之下,DiT 采用 Self-Attention,会导致图像大小的二次复杂度。因此,对于较大的图像,例如 512×512,DiC 和 DiT 之间的计算差距变得更加明显。尽管如此,令人惊讶的是,DiC 模型仍然可以与 DiT 保持良好的差距。另一方面,DiC 的优势在吞吐量方面也表现出色。
图10:ImageNet 512×512 生成上的 DiC 与 DiT 进行比较。将经过 400K 次迭代训练的模型与 DiT 的标准训练超参数进行比较
1.7 缩放性能
图 11 展示了随着训练的进行,DiC-H 在 ImageNet 256×256 上的性能。随着训练的进行,DiC-H 的性能稳步提高。在 600K 步时,FID 下降到 9.73,在 7M 训练步骤中匹配 DiT-XL/2 的性能;在 800K 步,DiC-H 显示出最佳性能,FID 为 8.96。该表展示了 DiC-H 的快速收敛,其中 FID 和 IS 都随着训练的进行不断改进。
图11:DiC 快速收敛
图 12 展示了各种模型在 ImageNet 256×256 上的性能的比较,因为训练迭代被扩大。提出的关键指标是吞吐量 (TP)、Batch Size × Iteration (BS×Iter) 和 FID (Fréchet Inception Distance),这是生成模型质量的常用度量。
图12:基于 CFG 的生成性能。TP 代表吞吐量。DiC 在条件生成方面也表现出色
基线模型 DiT-XL/2 的吞吐量为 66.8,得到 2.27 的 FID 分数。U-ViT-H 模型虽然吞吐量略低 63.9,但 FID 略高达到了 2.29。相比之下,DiC-H 模型使用 256 的 Batch Size,训练 2M 次,性能显示出了显著的改进。DiC-H 可以达到 2.25 的 FID,吞吐量为 160.8。这些结果表明 CNN 可以达到不比 Diffusion Transformer 差的性能。
#BiDM
突破瓶颈!北航Ð等首次将扩散模型完全量化至1bit!52.7倍效率提升
将扩散模型量化到1比特极限,又有新SOTA了!
来自北航、ETH等机构的研究人员提出了一种名为BiDM的新方法,首次将扩散模型(DMs)的权重和激活完全二值化。
具体而言,作者们从时间和空间的角度对扩散模型进行了优化:
从时间角度来看,引入了“时间步友好二值结构”(TBS),通过可学习的激活二值化器和跨时间步特征连接来应对DMs高度时间步相关的激活特征。
从空间角度来看,提出了“空间分块蒸馏”(SPD),目标是解决二值化特征匹配的困难,特别关注图像生成任务和噪声估计网络的空间局部性。
实验结果显示,W1A1 BiDM在LSUN-Bedrooms 256×256上的LDM-4模型上取得了22.74的FID分数,远远超越了当前状态的最先进通用二值化方法的59.44分,并实现了高达28倍的存储节省和52.7倍的操作效率提升。
下面具体来看。
现有扩散模型1位量化主要集中在「权重」
目前为止,扩散模型由于其高质量和多样化的生成能力,在图像、语音和视频等多个领域引起了极大的关注和应用。它可以通过多达1000步的去噪步骤,从随机噪声生成数据。
不过,虽然一些加速采样方法能够有效减少生成任务所需的步骤数量,但每个时间步的昂贵浮点计算仍然限制了该模型在资源受限场景中的广泛应用。
因此,对扩散模型的压缩成为其更广泛应用的关键步骤,现有的压缩方法主要包括量化、蒸馏、剪枝等。这些压缩方法的目标是在保持准确性的同时减少存储和计算成本。
其中,量化被认为是一种非常有效的模型压缩技术,通过将权重和/或激活量化为低位整数或二值化,实现紧凑存储和推理中的高效计算。
所以,已有一些研究将量化应用于扩散模型,以在保持生成质量的同时实现模型的压缩和加速。
而1位量化,即二值化,能够最大限度地节省模型的存储空间,并且在卷积神经网络(CNN)等判别模型中表现良好。此外,当权重和激活同时量化为1位时,如完全二值化,使用类似XNOR和bitcount的高效位运算可以替代矩阵乘法,达到最高效的加速效果。
一些现有的工作尝试将扩散模型量化为1位,但它们的探索主要集中在权重上,离完全二值化仍有较大距离。
实际上,对于扩散模型这样的生成模型,完全二值化权重和激活的影响是灾难性的:
a) 作为生成模型,扩散模型的丰富中间表示与时间步密切相关,而高度动态的激活范围在使用二值化权重和激活时受到严重限制;
b) 像扩散模型这样的生成模型通常需要输出完整的图像,但高度离散的参数和特征空间使得在训练过程中很难与真实值匹配。离散空间中的优化难度以及与时间步动态相关的表示能力不足,导致二值化扩散模型难以收敛,甚至在优化过程中崩溃。
提出BiDM:完全二值化「权重」和「激活」
面对上述不足,作者们提出了BiDM。通过完全二值化权重和激活,将扩散模型推向极限的压缩和加速。
概括而言,BiDM目标是解决扩散模型激活特征、模型结构以及生成任务的独特需求,克服完全二值化带来的挑战。
它包含两项创新技术:
从时间角度出发,作者们观察到扩散模型的激活特征与时间步高度相关。
因此引入了“时间步友好二值结构”(TBS) ,通过可学习的激活二值量化器匹配扩散模型的动态激活范围,并设计跨时间步的特征连接,利用相邻时间步之间的特征相似性,增强二值模型的表示能力。
而从空间角度出发,作者们注意到生成任务中的空间局部性以及扩散模型使用的基于卷积的U-Net结构。于是提出了“空间分块蒸馏”(SPD),引入一个全精度模型作为监督,通过对分块的自注意模仿,专注于局部特征,更好地引导二值扩散模型的优化方向。
△具有时间步友好二值结构和空间分块蒸馏的BiDM
广泛的实验表明,与现有的SOTA完全二值化方法相比,BiDM在保持相同推理效率的同时显著提高了精度,在各种评价指标上超越了所有现有基线方法。
具体来说,在像素空间扩散模型中,BiDM是唯一一种将IS提高到5.18的方法,接近全精度模型的水平,比最佳基线方法高出0.95。
在LDM 中,BiDM将LSUN-Bedrooms的FID从SOTA方法的59.44降低到了令人印象深刻的22.74,同时节省了28.0倍的存储空间和52.7倍的OPs。
作为第一种完全二值化的扩散模型方法,大量生成的样本也证明BiDM是目前唯一一种能够生成可接受的完全二值化DM图像的方法,从而使DM在低资源场景中得到有效应用。
具体咋实现的
二值扩散模型基线
扩散模型:给定一个数据分布𝑥0~𝑞(𝑥0),前向过程生成一个随机变量序列𝑥𝑡 ∊ {𝑥1,…,𝑥𝑇},该序列通过转移核𝑞(𝑥𝑡|𝑥𝑡-1)生成,通常是高斯扰动,其表达式为:
其中 是噪声调度。高斯转移核允许作者们边缘化联合分布,因此利用 : 和 ,可以通过采样一个高斯向量 并应用变换 轻松获得 的样本。
反向过程旨在通过去除噪声生成样本,使用可学习的转移核𝑝𝜃(𝑥𝑡-1|𝑥𝑡)近似不可用的条件分布q(𝑥𝑡-1|𝑥𝑡),其表达式为:
均值 和方差 可以使用重参数化技巧得到:
其中𝜖𝜃是具有可学习参数𝜃的函数逼近,用于根据𝑥𝑡预测𝜖。对于扩散模型的训练,通常使用变分下界的简化变体作为损失函数,以提高样本质量,其表达式为:
由于 U-Net 能够融合低层次和高维特征,它已成为扩散模型的主流骨干。U-Net 的输入输出块可以表示为 和 ,其中较小 对应的块是更低层次的。跳跃连接将低层次信息从 传播到 ,因此 接收到的输入表示为:
二值化:量化通过将权重和激活离散化为低位宽,压缩并加速噪声估计模型。在二值化扩散模型的基线中,权重W被二值化为1位:
其中符号函数将W限制为+1或-1,并设0为阈值。𝜎是浮点标量,初始化为(n表示权重数量),并在训练过程中进行学习。
同时,激活通常通过简单的BNN量化器进行量化:
当权重和激活同时被量化为1位时,去噪模型的计算可以由XNOR和bitcount运算代替,从而实现显著的压缩和加速。
时间步友好二值结构
在详细介绍所提方法之前,作者们先总结对扩散模型(DMs)属性的观察:
下图中, (a) 全精度DDIM模型在CIFAR-10上第4层卷积层的激活范围随去噪步骤变化。(b) 在LSUN-Bedrooms数据集上,全精度LDM-4模型在每一步的输出特征与前一步相似。
观察1:激活范围在长期时间步中变化显著,但激活特征在短期相邻时间步中相似。
之前的研究,如TDQ和Q-DM,已经表明,DMs的激活分布在去噪过程中高度依赖于时间步,表现为相邻时间步之间的相似性,而远距离时间步之间的差异较大,如图2(a)所示。
因此,在所有时间步中应用固定的缩放因子会导致激活范围的严重失真。除了分布范围之外,Deepcache强调了连续时间步之间高维特征的显著时间一致性,如图2(b)所示。
这些现象促使作者们重新审视现有的二值化结构。
二值化,尤其是权重和激活的完全二值化,与4位等低位量化相比,会导致激活范围和精度的更大损失。这使得生成丰富激活特征变得更加困难。激活范围和输出特征的不足严重损害了像DMs这样丰富表示的生成模型。
因此,采用更灵活的激活范围的二值量化器,并通过利用其特征输出来增强模型的整体表达能力,是在完全二值化后提高其生成能力的关键策略。
作者们首先关注长期时间步之间的差异。
大多数现有的激活量化器,如BNN和Bi-Real,直接将激活量化为{+1,-1},如式(7)所示。这种方法严重扰乱了激活特征,负面影响了生成模型的表达能力。一些改进的激活二值量化器,如XNOR++,采用了可训练的缩放因子k:
其中K的形式可以是向量或多个向量的乘积,但在推理过程中保持为常量值。尽管这种方法部分恢复了激活的特征表达,但它与高度相关的时间步不匹配,可能仍会导致显著的性能损失。
作者们将注意力转向原始的XNOR,它采用动态计算的均值来构造激活二值量化器。其对二维卷积的运算可以表示为:
其中 表示二维过滤器, 和 分别表示带和不带乘法的卷积操作。这种方法自然保留了激活特征的范围,并随着不同时间步的输入范围动态调整。
然而,由于DM特征的丰富表达,局部激活在模块通过前后的范围内表现出不一致,表明预定的k值不能有效恢复激活表示。
因此,作者们使k可调整,并允许其在训练过程中进行学习,以自适应匹配激活在前后发生的范围变化。作者们提出的可学习的微小卷积k的梯度计算过程如下:
值得注意的是,使k可学习不会增加任何额外的推理负担。计算成本保持不变,允许进行高效的二值运算。
另一方面,作者们专注于相邻时间步之间的相似性。
Deepcache直接提取高维特征作为缓存,以跳过U-Net中的大量深度计算,实现了显著的推理加速。其过程表示为:
然而,这种方法不适用于二值化扩散模型,因为二值网络的每个输出的信息内容非常有限。对于本质上已经实现了显著压缩和加速但表达能力有限的二值扩散模型,作者们预计相邻时间步之间的特征相似性将增强二值表示,从而弥补表达挑战。
作者们构建了一个跨时间步的信息增强连接,以使用前一时间步的特征丰富当前时间步的表达。其过程可以表示为:
其中是一个可学习的缩放因子。正如图2(b)所示,不同块和时间步中的高维特征相似性各异。
因此,作者们设置了多个独立的𝛼值,以允许模型在训练过程中更有效地自适应学习。
总结来说,时间步友好的二值结构(TBS)包括在输入平均和跨时间步连接后应用于缩放因子的可学习微卷积。它们的组合效果能够适应扩散模型在长期时间步中的激活范围变化,并利用相邻时间步之间高维特征的相似性来增强信息表示。
从误差减少的角度来看,图3中展示了TBS的可视化。
首先,作者们将基线方法下二值扩散模型的输出抽象为向量B。
首先,作者们将基线方法下二值扩散模型的输出抽象为向量 。
缩放因子的不匹配会在它与全精度模型的输出向量 之间产生显著的长度差异。
通过使用提出的缩放因子和可学习的微卷积, 被扩展为 。 更接近 ,但仍与全精度模型存在方向上的差异。
跨时间步的连接进一步结合了前一个时间步的输出 , 和 。相邻时间步之间的高维特征相似性意味着 和 之间的差距相对较小,从而促进了 和 的结合。
最后,作者们通过应用TBS得到二值化扩散模型的输出,表达式为 ,这一输出最接近全精度模型的输出 。
空间分块蒸馏
由于生成模型的特性,扩散模型的优化过程与以往的判别模型有着不同的特征:
观察2:传统的蒸馏难以引导完全二值化的扩散模型与全精度模型对齐,而扩散模型在生成任务中的特征在空间上表现出局部性。
在以往的实践中,在量化模型的训练过程中添加蒸馏损失是常见的做法。由于二值模型的数值空间有限,直接使用简单的损失进行优化会导致调整梯度更新方向的困难,使学习变得具有挑战性。因此,向中间特征添加蒸馏损失可以更好地指导模型的局部和全局优化过程。
然而,作为生成模型,扩散模型的高度丰富的特征表示使得二值模型极难细致地模仿全精度模型。
尽管原始扩散模型训练中使用的L2损失与扩散过程中的高斯噪声对齐,但并不适用于中间特征的蒸馏匹配。在常规蒸馏过程中,常用的L2损失往往优先优化差异较大的像素,从而导致更加均匀和平滑的优化结果。
这种全局约束学习过程对于以图像生成为目标的二值模型来说是困难的,因为其有限的表示能力使得精细的蒸馏模仿难以直接调整模型以完全匹配全精度模型的方向。
与此同时,作者们注意到,使用U-Net作为骨干的扩散模型由于其基于卷积的结构和生成任务的要求,天然表现出空间局部性。
这与以往的判别模型不同,判别模型中的任务(如分类)只需整体特征提取,而不需要低层次的要求,这使得传统的蒸馏方法不适用于具有空间局部性的生成任务。此外,大多数现有的扩散模型蒸馏方法专注于减少时间步数,并未解决图像生成任务中对特征空间局部性的需求。
因此,鉴于现有损失函数在优化二值扩散模型上的困难以及扩散模型的空间局部性,作者们提出了空间分块蒸馏(SPD)。
具体来说,他们设计了一种新的损失函数,该损失函数在蒸馏之前将特征划分为小块,然后逐块计算空间自注意力。虽然传统的L2损失使二值扩散模型难以实现直接匹配,导致优化挑战,但注意力机制使得蒸馏优化可以更侧重于关键部分。
然而,对于完全二值化的扩散模型来说,这仍然是具有挑战性的,因为高度离散的二值输出信息有限,使得模型难以捕获全局信息。
因此,作者们通过将中间特征划分为多个小块,并为每个小块独立计算空间自注意力,使得二值模型在优化过程中能够更好地利用局部信息。
SPD首先将二值扩散模型和全精度扩散模型的某个块输出的中间特征 和 划分为 个小块:
然后,分别对每个小块计算自注意力:
在正则化之后,计算并累加相应位置的损失:
其中|| · ||2表示L2函数。最后,总训练损失被计算为:
其中,d表示在上采样或下采样过程中块的数量,最终包括中间块在内,共有2d+1中间特征。𝜆是用于平衡损失项的超参数系数,默认为4。
如下图所示,LDM模型在LSUN-bedroom数据集上的最后一个TimeStepBlock输出的可视化。FP32表示全精度模型的输出 。Diff表示全精度模型输出与二值化模型输出之间的差异 。Ours表示自注意力机制。
作者们对上述提到的中间特征和自注意力进行了可视化。正如图所示,作者们的空间分块蒸馏(SPD)使得模型在每个小块中更加关注局部信息。
实验
作者们在各种数据集上进行了实验,包括CIFAR-10 32×32、LSUN-Bedrooms 256×256、LSUN-Churches 256×256 和FFHQ 256×256,并在像素空间扩散模型和潜在空间扩散模型上进行了测试。
使用的评估指标包括Inception Score (IS)、Fréchet Inception Distance (FID)、Sliding Fréchet Inception Distance (sFID)、精度和召回率。截至目前,还没有研究将扩散模型压缩到如此极端的程度。
因此,作者们使用了经典的二值化算法和最新的最先进的通用二值化算法作为基线。
他们提取了扩散模型中TimestepEmbedBlocks的输出作为TBS和SPD的操作目标。且还在卷积层中采用了与ReActNet相同的快捷连接。详细的实验设置见附录A。
主要结果
像素空间扩散模型:作者们首先在CIFAR-10 32×32数据集上进行了实验。正如表1中所示,使用基线方法对扩散模型进行W1A1二值化导致了显著的性能下降。然而,BiDM在所有指标上均表现出显著改进,实现了前所未有的图像质量恢复。
具体而言,BiDM在IS指标上从4.23提升至5.18,FID指标降低了27.9%。
以下为在CIFAR-10数据集上,使用100步的DDIM二值化结果。
潜在空间扩散模型:作者们的LDM实验包括对LDM-4在LSUN-Bedrooms 256×256和FFHQ 256×256数据集上的评估,以及对LDM-8在LSUN-Churches 256×256数据集上的评估。
实验使用了200步的DDIM采样器,详细结果见表2中。
在这三个数据集中,作者们的方法相对于最好的基线方法取得了显著的改进。与其他二值化算法相比,BiDM在所有指标上都表现优异。
在LSUN-Bedrooms、LSUN-Churches和FFHQ数据集中,BiDM的FID指标分别比基线方法降低了61.7%、30.7%和51.4%。与XNOR++相比,XNOR++在去噪过程中采用了固定的激活缩放因子,导致其激活的动态范围非常有限,使得难以匹配扩散模型高度灵活的生成表示。
BiDM通过使微卷积k可学习,作用于动态计算的缩放因子,解决了这一挑战。这种优化使得各项指标取得了超过一个数量级的显著改进。
在LSUN-Bedrooms和LSUN-Churches数据集中,BiDM的FID指标分别从319.66降低至22.74和从292.48降低至29.70。
此外,与最先进的二值化方法ReSTE相比,BiDM在多个指标上实现了显著提升,特别是在LSUN-Bedrooms数据集中表现出显著的改进。
消融实验
作者们在LSUN-Bedrooms 256×256数据集上对LDM-4进行了全面的消融研究,以评估BiDM中每个提出的组件的有效性。他们评估了提出的SPD和TBS方法的有效性,结果如表3所示。
当分别将SPD或TBS方法应用于LDM时,相较于原始性能,观察到了显著的改进。当引入TBS方法时,FID和sFID分别从106.62和56.61大幅下降到35.23和25.13。
同样,当添加SPD方法时,FID和sFID分别显著下降到40.62和31.61。其他指标也显示出显著的改善。这证明了作者们的方法在训练过程中通过引入可学习的因子并加入跨时间步的连接,能够持续将二值化模型的特征逼近全精度特征的有效性。
此外,当结合这两种方法并将它们应用于LDM 时,较之单独应用每种方法,观察到额外的改进。这进一步证明了在补丁级别执行全精度和二值模型之间的蒸馏能够显著提高二值模型的性能。
作者们还进行了额外的消融实验,结果在附录B中展示。
效率分析
作者们对模型在完全二值化下的推理效率进行了分析。
表4中的结果表明,BiDM在实现与XNOR基线相同的28.0倍内存效率和52.7倍计算节省的同时,图像生成能力显著优于基线模型,其FID从106.62降低至22.74。
总结
在本文中,作者们提出了BiDM,这是一种全新的完全二值化方法,将扩散模型的压缩推向极限。
基于两个观察——不同时间步的激活特性和图像生成任务的特点,他们分别从时间和空间角度提出了时间步友好的二值结构(TBS) 和 空间分块蒸馏(SPD) 方法。
这些方法解决了完全二值化中表示能力的严重限制,以及高度离散的空间优化挑战。
作为首个完全二值化的扩散模型,BiDM在多个模型和数据集上表现出了显著优于现有最先进通用二值化方法的生成性能。
在LSUN-Bedrooms数据集上,BiDM的FID为22.74,远超最先进方法的FID为59.44,成为唯一能够生成视觉上可接受样本的二值方法,同时实现了高达28.0倍的存储节省和52.7倍的运算效率提升。
更多细节欢迎查阅原论文。
论文链接:https://arxiv.org/abs/2412.05926
项目主页:https://github.com/Xingyu-Zheng/BiDM
#Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps
扩散模型也能推理时Scaling,谢赛宁团队重磅研究可能带来文生图新范式
对于 LLM,推理时 scaling 是有效的!这一点已经被近期的许多推理大模型证明:o1、o3、DeepSeek R1、QwQ、Step Reasoner mini……
但这个说法也适用于扩散模型吗?
近日,纽约大学谢赛宁领导的一个团队对这一方向进行了探索。具体来说,他们借助通用搜索框架系统性地探索了扩散模型的推理时 scaling 情况。他们发现,推理时 scaling 对扩散模型是有效的 —— 增加推理时间计算可以显著提高扩散模型生成的样本的质量,并且由于图像的复杂性,可以针对不同的应用场景,对框架中的组件进行不同形式的组合。
这篇论文的署名部分还透露出了一个有趣的小细节:谢赛宁的所属机构标记成了谷歌并且通讯作者邮箱也使用了其谷歌域名的邮箱。但我们目前还并不清楚这是否意味着这位著名 AI 研究者已经正式加入谷歌,还是在谷歌兼职(谢赛宁此前已经是谷歌研究院访问学者)。
- 论文标题:Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps
- 论文链接:https://arxiv.org/pdf/2501.09732
有趣的是,就在前几天,来自纽约大学和哥伦比亚大学的研究者也发布了一篇类似方向的论文。看来,scaling 推理将是改进扩散模型的一个非常有前途的方向。
论文概览
OpenAI o1 等模型的问世已经证明,在推理阶段增加计算量可以让 LLM 的性能进一步提升。但对于扩散模型,如何有效 scaling 推理时间计算以进一步提升性能还不够明确。
扩散模型经过训练以去除数据中的噪声,是一类在连续数据领域占主导地位的生成模型,如图像、音频和视频。为了生成单个样本,它们的生成过程通常从纯噪声开始,需要经过训练模型的多次前向传递来去噪并获得干净的数据。这些前向传递因此被称为去噪步骤。由于去噪步骤的数量可以调整,以在样本质量和计算成本之间进行权衡,扩散模型的生成过程自然提供了在推理时分配计算预算的灵活性。
对于生成模型,这种计算预算通常用函数评估次数(NFE)来衡量,以确保与其他使用迭代采样过程但没有去噪能力的模型系列进行合理比较。
经验观察表明,仅通过将计算投入到去噪步骤中,性能提升在达到某个 NFE 后往往会趋于平稳,限制了推理过程中增加计算的收益。因此,以前关于扩散模型的工作长期以来一直专注于在保持高性能的同时尽可能减少推理时的 NFE 以提高效率。
但是,谷歌 DeepMind 和谢赛宁等人新论文的研究方向与此相反。
与 LLM 相比,扩散模型处理的是作为初始样本注入的噪声或在采样过程中注入的噪声的显式随机性。已有研究表明这些噪声并非等价,即某些噪声会带来更好的生成结果。这一观察为 scaling NFE 提供了除增加去噪步骤之外的另一个维度 —— 在采样中搜索更好的噪声。
新论文研究了在推理过程中通过搜索有效利用计算资源的方法,而不是仅将 NFE 分配给去噪步骤(这往往会很快导致性能不再增长),从而在推理时改进扩散模型的性能和可扩展性(图 1)。
作者主要考虑搜索框架中的两个设计轴:用于在搜索中提供反馈的验证器,以及用于寻找更好噪声候选项的算法,遵循 LLM 中使用的术语。
对于验证器,他们考虑了三种不同的设置,这些设置旨在模拟三种不同的使用场景:
- 掌握最终评估如何进行的特权信息的场景;
- 掌握用于指导生成的条件信息的场景;
- 没有额外信息可用的场景。
对于算法,他们研究了:
- 随机搜索,它只是从固定的候选集中选择最佳项;
- 零阶搜索,它利用验证器反馈来迭代改进噪声候选项;
- 路径搜索,它利用验证器反馈来迭代改进扩散采样轨迹。
作者首先在 ImageNet 类别条件生成这个相对简单的设置中探讨这些设计选择并展示它们的有效性,为新框架提供一个具体实例。然后他们将这些设计选择应用到更大规模的文本条件生成设置中,并评估他们提出的框架。
由于图像的复杂性质和文本条件包含的丰富信息,需要对生成质量进行更全面的评估。因此,他们采用多个验证器来进行搜索中的推理时计算 scaling。这也使他们能够探究每个验证器具有的「偏见」,以及它们与生成任务的匹配程度。为了缓解对单个验证器的过拟合,他们还尝试了验证器集成,并展示了它在不同基准测试中的良好泛化能力。
新论文的贡献总结如下:
- 提出了一个用于扩散模型推理时 scaling 的基础框架。论文表明,通过搜索去 scaling NFE 可以在各种生成任务和模型规模上带来实质性改进,超越了仅增加去噪步骤的方法。此外,作者对推理时计算预算如何影响 scaling 性能进行了全面的实证分析。
- 在提出的搜索框架中确定了两个关键设计轴:提供反馈的验证器和寻找更好噪声候选项的算法。作者研究了不同验证器 - 算法组合在各种任务中的表现,他们的发现表明没有一种配置是普遍最优的;每个任务反而需要一个独特的搜索设置才能实现最佳的 scaling 性能。
- 对验证器与不同生成任务之间的匹配度进行了广泛分析。论文的结果揭示了不同验证器中嵌入的偏见,以及在每个不同的视觉生成任务中需要专门设计验证器的必要性。
如何在推理时间 scale
本文提出了一个框架,用于扩散模型的推理时间 scaling 。本文将这一挑战表述为对采样噪声的搜索问题;具体来说,本文需要知道哪些采样噪声是好的,以及如何搜索这些噪声?
在高层次上,本文考虑了两个设计轴:
一是验证器,用于评估候选样本的质量。这些验证器通常是能够提供反馈的预训练模型;验证器函数形式如下:
这些函数接收生成的样本以及可选的对应条件,并输出一个标量值作为每个生成样本的评分。
二是算法,用于根据验证器的评分找到更好的候选样本。函数形式如下:
它接收验证器 V、预训练扩散模型 𝐷_θ ,以及 N 对生成的样本及其相应的条件,并根据噪声和样本之间的确定性映射输出最佳的初始噪声。在整个搜索过程中,f 通常会多次通过 𝐷_θ 进行前向传递。本文将这些额外的前向传递称为搜索成本,用 NFE 来衡量。
举例来说,本文使用在 ImageNet 上预训练的 SiT-XL 模型,分辨率为 256×256,并使用二阶 Heun 采样器进行采样,即除了初始噪声外,采样过程中没有其他随机性来源。此外,本文还通过去噪步骤和搜索成本的总 NFE 来衡量推理计算预算。
本文从最简单的搜索算法开始,即随机采样高斯噪声,使用 ODE 生成样本,并选择与最佳验证器评分对应的样本(图 2)。本文将这种算法称为随机搜索,它本质上是一种在所有噪声候选上应用一次的 Best-of-N 策略。
对于验证器,本文从最佳验证器开始,即 Oracle 验证器,他们假设其拥有关于选定样本最终评估的全部特权信息。对于 ImageNet,由于 FID 和 IS 通常被用作评估指标,因而直接将它们作为 Oracle 验证器。
搜索验证器
虽然通过搜索 scalingNFE 在使用 Oracle 验证器时表现出令人印象深刻的性能(如图 3 所示),但关键问题在于这种有效性是否可以推广到使用更易获取的预训练模型的监督验证器上。
为了研究这一点,本文选择了两个模型:CLIP 和 DINO 。对于 CLIP,本文遵循 Radford 等人的方法,使用通过提示工程生成的嵌入权重作为零样本分类器。对于 DINO,本文直接使用预训练的线性分类头。在搜索过程中,将样本输入分类器,并选择与生成中使用的类别标签对应的 logits 最高的样本。
如图 4 所示,与单纯通过增加去噪步骤 scaling NFE 相比(图 1),这种策略也有效提升了模型在 IS 上的性能。然而,本文注意到,由于这些分类器是逐点操作的,它们仅部分与 FID 评分的目标对齐。具体来说,它们生成的 logits 仅关注单个样本的质量,而没有考虑总体多样性,这导致样本方差显著减少,并最终表现为随着计算量增加而出现的模式崩溃。随机搜索算法也难辞其咎,由于其无约束的搜索空间,加速了搜索向验证器偏见的收敛。这种现象类似于强化学习中的奖励黑客行为,因此将其称为验证器黑客。
尽管条件信息在现实世界的生成任务中至关重要,但本文发现验证器并不一定需要条件信息来有效指导搜索。如图 5 所示,他们发现 DINO/CLIP 分类器输出的 logits 与模型在低噪声水平(σ=0.4)下的 x 预测值与最终生成的干净样本(σ=0)之间的特征空间(分别由 DINO/CLIP 提取)余弦相似度之间存在强相关性。
因此,本文使用这种相似度评分作为分类 logits 的替代,并将这类验证器称为自监督验证器,因为它们不需要额外的条件信息。图 5 再次观察到了有效的 scaling 行为。
这一结果对于条件信息不可用或难以获取的用例来说(如医学图像生成任务)是令人鼓舞的。由于这些限制在现实场景中并不常见,本文将自监督验证器的进一步研究留待未来工作。
搜索算法
前文的探索主要集中于简单的随机搜索设置,他们发现这种方法可能导致验证器黑客行为。基于这一认知,作者开始研究更细致的搜索算法,这些算法利用验证器的反馈逐步优化候选样本,每次只进行小幅调整,从而减轻过拟合风险。
具体来说,本文考虑了一种零阶搜索方法:
- 从一个随机高斯噪声 n 作为支点开始。
- 在支点的邻域中找到 N 个候选项。
- 通过 ODE 求解器运行候选项以获得样本及其对应的验证器得分。
- 找到最佳候选项,将其更新为支点,并重复步骤 1-3。
与零阶优化类似,零阶搜索不涉及昂贵的梯度计算;相反,它通过在邻域内进行多次前向函数评估来近似梯度方向。
本文还注意到,由于许多验证器是可微的,理论上可以使用真实梯度进行一阶搜索,并且在实际中已有应用。然而,这需要在整个采样过程中进行反向传播,通常在时间和空间复杂度上都极为昂贵,尤其是在扩展大模型时。在实践中,本文发现一阶搜索在 ImageNet 上并没有表现出比零阶搜索更明显的优势,尽管它的成本更高。
扩散采样过程的迭代特性为设计局部搜索算法提供了其他可能性,沿着注入噪声的采样轨迹进行搜索是可行的。本文提出了路径搜索来探索其中一种可能性。具体来说,
- 采样 N 个初始独立同分布噪声,并运行 ODE 求解器直到某个噪声水平 σ。噪声样本 x_σ 作为搜索起点。
- 为每个噪声样本采样 M 个独立同分布噪声,并模拟从 σ 到 σ+Δf 的前向加噪过程,生成大小为 M 的 {x_(σ+Δf)}。
- 在每个 x_(σ+Δf) 上运行 ODE 求解器到噪声水平 σ + Δf - Δb,并获得 x_(σ+Δf-Δb)。在这些样本上运行验证器并保留前 N 个候选样本。重复步骤 2-3,直到 ODE 求解器达到 σ=0。
- 对剩余的 N 个样本通过随机搜索并保留最佳的一个。
为了确保迭代会终止,该团队严格要求 Δb > Δf。另外,由于验证器通常不适应有噪声输入,因此他们在步骤 3 中执行一个额外的去噪步骤,并使用干净的 x 预测与验证器交互。这里,主要的 scaling 轴是步骤 2 中添加的噪声数量 M,在实践中,他们研究了不同初始噪声数量 N 的 scaling 行为。
因此该算法被命名为 Paths-N。这两种算法见图 2,从中可以看到,与随机搜索相比,零阶搜索和路径搜索都保持了很强的局部性:前者在初始噪声的邻域中运作,后者则是在采样过程的中间步骤中搜索。
这些算法的性能见图 6。由于这两种算法的局部性质,它们都在一定程度上缓解了 FID 的多样性问题,同时保持了 Inception Score 的 scaling。对于零阶搜索,可以注意到增大 N 的效果很小,N = 4 似乎已经是对局部最优的良好估计。至于路径搜索,可以看到不同的 N 值会导致不同的 scaling 行为,小 N 值在小生成预算下计算效率高,而大 N 值在增加更多计算时更具优势。
在文生图场景中的推理时 scaling
在更大规模文生图任务中,该团队研究了搜索框架的推理时 scaling 能力,并研究了验证器与特定图像生成任务之间的对齐性能。
数据集方面,该团队采用了 DrawBench 和 T2I-CompBench 以实现更全面的评估。骨干模型则是 FLUX.1-dev。验证器方面,考虑到文生图任务本质上的复杂性质,该团队还扩展了监督式验证器的选择:Aesthetic Score Predictor、CLIPScore、ImageReward。其中 Aesthetic Score Predictor 可以预测人类对合成图像的视觉质量的评分,CLIPScore 则是通过 4 亿对人工标注的图像 - 文本数据来对齐视觉和文本特征,ImageReward 则更能体现更一般化的偏好。
此外,该团队还创新地将这三个验证器组合到一起,得到了第四个验证器,称为 Verifier Ensemble。可用于进一步扩展验证器在不同评估方面的能力。
更多实验设置详情请参看原论文。
分析结果:验证器 - 任务对齐
下面比较在不同数据集上验证器和算法组合的结果。
DrawBench。这是一个高度通用的数据集,包含来自不同类别的文本提示词。图 8 展示了在该数据集上的一些结果。
如图 8 所示,根据 LLM Grader 指标,可知使用所有验证器进行搜索通常都会提高样本质量,而具体的改进行为在不同设置下有所不同。这表明,可以根据不同的应用场景选择专门的搜索设置。
从图 8 的左两列中,可以看到使用 Aesthetic 和 CLIP 验证器进行搜索会过度拟合它们固有的偏见,对彼此产生负面影响。该团队推测这两个验证器在评估上存在重大不匹配:Aesthetic Score 仅关注视觉质量,经常偏向于与文本提示词偏离的高度风格化图像,而 CLIP 优先考虑视觉 - 文本对齐,牺牲了视觉质量。因此,如果在搜索过程中使用一个验证器(例如 Aesthetic Score),其偏见便会降低由另一个验证器(例如 CLIP)评估的指标。
另外,与 ImageNet 设置类似,该团队在这里也观察到了随着搜索预算增长而表现出的 scaling 行为。
通过观察验证器在 DrawBench 和 T2I-CompBench 上的行为,可以发现:验证器的有效性取决于其标准与任务特定需求的匹配程度,某些验证器比其他验证器更适合特定任务。
算法。表 2 展示了三种搜索算法在 DrawBench 上的表现。
可以看到,这三种方法全都能有效提高采样质量,其中随机搜索在某些方面优于其他两种方法。同样,该团队认为这种行为的原因是零阶搜索和路径搜索的局部性质(图 2)。由于这里展示的所有验证器和指标都是基于每个样本进行评估的,随机搜索将大大加速向验证器偏见的收敛,而其他两种算法需要对次优候选项进行改进。
搜索与微调是兼容的
搜索和微调的目标都是将最终样本与显式奖励模型或人类偏好对齐。前者将样本模式向特定验证器的偏见移动,后者是直接修改模型的分布以与奖励对齐。于是问题来了:在模型分布被修改后,是否仍然可以根据验证器移动样本模式?
在这里研究过的所有微调方法中,作为 RLHF 方法的一个更高效和更简单的替代方案,DiffusionDPO 已被广泛用于对齐大规模文生图模型。为了回答这个问题,该团队采用论文《Diffusion model alignment using direct preference optimization》中的 DPO 微调的 Stable Diffusion XL 模型,并在 DrawBench 数据集上进行搜索。由于该模型是在 Pick-a-Pic 数据集上微调的,于是他们也用 PickScore 评估器替换 ImageReward。结果见表 3。
可以看到,搜索方法可以泛化到不同的模型,并可以改进已对齐模型的性能。看起来,这个工具很有用,可以缓解微调模型与奖励模型不一致的情况,并提高它们在其他指标上的泛化能力。
投入推理计算上不同维度的不同效果
由于扩散模型本质上的迭代采样性质,因此通过搜索来 scaling NFE 的维度有很多。这篇论文也对此展开了一些讨论。
搜索迭代的次数。直观地说,增加搜索迭代次数可以使选定的噪声相对于验证器更接近最优集,从而可以显著提高性能。
每次搜索迭代的计算量。在每次搜索迭代中,去噪步骤数量都可以调整。这里将模型采用的去噪步骤次数记为 NFEs/iter。尽管仅增加去噪步骤时模型性能很快达到平稳(图 1),但该团队观察到在搜索过程中,调整 NFEs/iter 可以揭示不同的计算最优区域,如图 10 所示。在搜索过程中,较小的 NFEs/iter 能实现高效收敛,但最终性能较低。相反,较大的 NFEs/iter 会导致收敛较慢但产生更好的性能。
最终生成的计算量。尽管可以自由调整最终生成的去噪步骤,但该团队的做法是始终使用最优设置以获得最佳的最终样本质量。在 ImageNet 中,他们将去噪预算固定为 250 NFEs,在文本到图像设置中使用 30 步采样器,因为进一步 scaling 时的性能将很快趋平。
投资计算的有效性
该团队探索了在较小的扩散模型上进行推理时计算 scaling 的有效性,并强调了相对于它们不进行搜索的大型模型的性能效率。结果见图 11。
可以看到,在 ImageNet 上,对小型模型进行推理时计算 scaling 会非常有效。在固定的推理计算预算下,对小型模型进行搜索可以超过不进行搜索的大型模型。举个例子,在推理计算量有限时,SiT-L 相对于 SiT-XL 存在优势。但是,将 SiT-B 与其他两个模型进行比较表明,这种有效性取决于小型模型相对强大的基线性能。当小型模型的基线性能明显落后时,scaling 的收益是有限的,会导致结果次优。
在基于文本的设置中也出现这种现象,如表 4 所示。
可以看到,仅使用十分之一的计算量,PixArt-Σ 就超过了不进行搜索的 FLUX-1.dev。而如果使用大约两倍的计算量,PixArt-Σ 还能大幅超过不进行搜索的 FLUX-1.dev。这些结果具有重要的实践意义:在训练中投入的大量计算资源可以通过生成期间的一小部分计算来抵消,从而能更高效地获得更高质量的样本。
#并非所有扩散模型噪声本质相同
这篇文章发现了并非所有扩散噪声生成质量相同这一问题,并设计了反演稳定度这一衡量噪声质量的指标。
自从 ChatGPT o1 在 2024 年 9 月发布后,人们逐渐把研究重点放在了推理时扩展 (Inference-time scaling) 上。对于扩散模型而言,除了在推理时增加步数外,谷歌今年 1 月的研究 Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps 表明,搜索更好的扩散模型噪声也能提升图像质量。但实际上,选取更好的噪声这一想法早在 2024 年 6 月的论文 Not All Noises Are Created Equally:Diffusion Noise Selection and Optimization 中就已提出。可惜的是,这篇论文最后在 ICLR 上被主动撤稿了,详情可参考 OpenReview:https://openreview.net/forum?id=R5xozf2ZoP 。
在这篇博文中,我会简单介绍「并非所有噪声本质相同」这篇论文。之后,我们来尝试批判性分析这篇论文及其审稿意见,并反思该如何做出高质量的工作。
知识准备
想读懂这篇论文,只需了解 DDPM 和 DDIM 这两篇扩散模型经典论文。
DDPM 定义了加噪过程和去噪过程:加噪过程中,训练图像 会加噪 次,最后变成纯噪声 ;去噪过程中,纯噪声 会被去噪 次,最后变成生成图像 。
DDIM 修改了 DDPM 中的部分假设,让扩散模型在去噪时不再有随机性,而唯一的随机性来自于随机噪声 的采样。保证了去噪过程不再有随机性后,我们就可以通过公式反演 (inversion),找出某个真实图像 对应的 。当然,这个反演过程存在误差。
核心思想:更好的噪声带来更高的图像质量
作者认为,过去的工作仅通过优化网络结构或参数来提升扩散模型生成质量,却忽视了噪声的重要性。作者从两个方向探索了噪声空间(即所有可能的噪声构成的集合):能否用某种定量指标选取更好的噪声?能否优化某一现有噪声?这两个方向的应用如下所示。噪声选取能够提升模型在同一个文本提示下的随机生成质量,而噪声优化能够提升现有某噪声对应图像的质量。
用反演稳定度反映噪声质量
根据以上设计思想,这个工作的重点就落在了噪声质量的评估指标上。只要有了一个这样的指标,我们就能用简单的算法实现噪声选取和噪声优化:
- 噪声选取:随机采样大量噪声,然后根据指标选出最好的那一个。
- 噪声优化:把指标变成损失函数,用损失函数优化给定噪声。
作者认为,函数不动点具有某些良好性质。因此,如果一个噪声在做完一轮生成-反演(去噪-加噪)后变化越小,那么这个噪声的质量越好。
所谓不动点,就是令函数 满足 的输入 。这里作者把函数 定义成先生成再反演这一整个过程,把 定义成某个高斯噪声。
具体来说,作者用 相似度来反映噪声的「变化幅度小」,即对于噪声 和其生成-反演结果 ,反演过程的稳定度,或者说噪声的质量,等于 。
有了这样一个指标后,我们可以用下面伪代码所示算法实现噪声选取和噪声采样:
input number of samples K, denoising model M
eps_list = sample_noise(K)
stability_list = []
for i in range(K):
eps = eps_list[i]
img = generate(M, eps)
eps_inversion = inversion(M, img)
stability = cos(eps, eps_inversion)
stability_list.append(stability)
best_idx = argmax(stability_list)
output eps_list[best_idx]
input denoising model M, optimization steps n, noise eps
for i in range(n):
img = generate(M, eps)
eps_inversion = inversion(M, img)
loss = 1 - cos(eps, eps_inversion)
eps = gradient_descent(loss, eps)
output eps
实验结果实验配置
本工作用 SDXL 和 SDXL-turbo 模型做实验。去噪步数分别为 10, 4。
本工作在测试时只需要设置提示词。提示词来自于 Pick-a-pic, DrawBench, GenEval 这三个数据集。
- Pick-a-Pic 收集了真实文生图应用中用户的输入提示词及相同提示词下不同输出的偏好程度。每项数据包含提示词、一对图像、一个表示用户偏好哪个图像的标签。
- Drawbench 的每项数据包含提示词、一对图像、一个反映图像质量的偏好标签、一个反映提示词对齐度的偏好标签。这些提示词能从颜色、物体数量等多个角度测试模型生成属性。
- GenEval 提供了能从单个物体、两个物体、物体数量、颜色、位置等多个角度测试模型的提示词。
为了对比生成图像质量,本工作用了 HPS v2, AES,PickScore, ImageReward 指标。除 AES 只考虑美学分数外,其他几个指标都是通过训练一个匹配人类偏好的模型来评估图像质量。除了给出这些指标的值外,论文还展示了方法的胜率——在同样的提示词下,比较使用/不使用新方法的定量结果,统计使用新方法胜利的频率。
定性定量结果
以下是噪声选取的实验结果。
以下是噪声优化的实验结果。
从定量结果中可以看出,本文的方法一定程度上确实是有效的。但是,这些数值上的提升真的总能反映在视觉效果上吗?这些数值上的提升足够大吗?由于缺乏同类方法的比较,这些实验还是欠缺说服力。
博文开头已经展示了本文的定性结果,这里就不再重复展示了。值得一提的是,本文的方法可以拓展到其他模型上,比如生成 3D 模型的 SV3D。以下是 3D 模型上的可视化结果。
审稿意见分析
在看 ICLR 2025 官方审稿人意见之前,我先给出我的一份简单审稿意见。我把本文的主要内容都在这篇博文里写出来了,读者也可以在阅读下面内容之前尝试给该论文打分。如果读者想着重提升审稿能力,可以先去把这篇论文细读一遍,再去 OpenReview 上审视其他审稿人的意见。
这篇文章发现了并非所有扩散噪声生成质量相同这一问题,并设计了反演稳定度这一衡量噪声质量的指标。定性和定量实验表明该方法确实能提升图像质量。
本文的优点在于新问题的发现:确实之前多数研究者只考虑如何优化模型,而忽视了噪声也是影响图像的重要因素。
本文在叙述和论证上均存在缺陷:
- 在为何要使用反演稳定性这件事上,作者说是受到了不动点理论的启发。但作者完全没有深入分析反演稳定性的原理,本质上是通过实验来直接验证假设的合理性。哪怕不好从数学原理上解释,作者也应该通过列出相关工作,并补充部分探究原理的实验,来巩固该指标的合理性。
- 本文的技术贡献并不多:优化噪声是一个非常常见的策略。
- 在实验中,作者没有和其他方法对比,完全看不出这个方法是否具有优越性。找出一个基准方法其实并不难:将 CFG (classifier-free guidance) 分别设置成 1 和 7.5 就是一个非常容易实现的基准方法。
如果是现在已经知道了优化噪声是一个有潜力的方向,我可能会给这篇论文一个位于接收临界线上的分数。但如果是去年的话,很可能给的是低于临界线的分数。
我们来简单看看官方审稿人的评价。
审稿人 A 给了明确的差评。文章的主要缺陷有:
- 虽然作者宣称这是首次研究噪声空间的论文,但实际上之前已经至少有三篇工作做过类似研究。因此,文章的创新性是有限的。
- DDIM 反演稳定性差一部分是因为反演算法不够好导致的,而和扩散模型无关。最好能够展示更好反演算法下的结果。
- 评估不充分:这篇文章没有展示任何和其他方法的比较。另外,最好能展示诸如 FID,IS 等其他指标。文中的指标提升看起来微不足道,仅有约 1%,感觉不到明显的进步。
- 每次评估都需要执行一次加噪-去噪过程,方法性能较低。
审稿人 A 是这个方向的专家,他用已有的噪声空间优化工作来质疑本文的创新性。我简单看了一下该审稿人引用的工作,确实已经有不少往期工作做了类似的事,但实现方法还是比本文的方法更复杂一点。所以,不能说这篇论文一点创新也没有,作者必须严谨地比较这篇论文和之前的论文,并强调出创新点在哪。
审稿人 B 给了一份评价略低于临界线的短评。他认为这篇论文提出的「反演稳定性」及其实验结果还是非常合理的。他认为论文的主要缺陷是方法质量受到反演算法的影响,并希望作者给出不同反演策略下的结果。
审稿人 C 也给出了略低于临界线的评价。他认为这种使用反演稳定性来评估噪声质量的想法有新颖性,且实验的设置比较合理。在 3D 上的应用或许很有潜力。但他也列出了几条主要缺陷:
- 作者认为反演质量与生成质量强相关。但这个假设只是通过一个具体示例来验证,而缺少理论分析或者系统性的研究。一个可能的分析方法是:测试10万个样本,计算图像质量和反演稳定性的相关性。
- 和上面两个审稿人一样,审稿人 C 也提到了 DDIM 反演本身就不够好的问题。
- 噪声选取和噪声优化都需要大量计算资源。
审稿人 D 也给出了略低于临界线的评价。他同样认为反演稳定性是一种新颖的指标,作者对于方法设计动机阐述得十分清晰。他提到的主要缺陷有:
- 研究反演中的不同点并不是一个新颖的想法,之前至少有三篇工作从迭代优化来寻找反演不动点的文章。这篇文章没有与这些工作做对比。
- 从 K=100 个样本里选取噪声的方法太简单直接了。之前有类似工作用一种更深入的方式选取噪声。
- 缺乏和其他噪声优化论文的比较。
- 计算效率太低。
总结与反思
本工作明确指出了扩散模型不同初始噪声质量不同这一设计动机。为了找出更好的噪声,本工作用一种听起来比较有趣的「反演稳定性」来评估噪声质量。基于这个指标,本工作设计了噪声选取和噪声优化两个算法,用于两类生成任务。
多数审稿人都认可该工作的方法有一定的新颖性,但其论文有着诸多缺陷:
- 没有深入分析反演稳定性的原理,且没有考虑是否应该用更好的反演算法。
- 噪声优化是一种常见的策略。没有与相关工作做对比。
- 指标提升的说服力不足。
- 每次评估反演稳定性都需要执行去噪和加噪过程,运算效率过低。
当然,在目前的眼光看来,如果生成质量足够高,增加推理时间也是可以接受的。但是,论文的前几个缺陷是实实在在的,被评价为低于接收临界线合情合理。作者最后也没有 rebuttal,直接撤稿了。我感觉这篇论文还是比较可惜的,毕竟它的研究方向比较有潜力。
从这篇论文以及最近做推理时扩展的论文可以看出,哪怕对于同一项技术,改变看待问题的角度也算创新。比如同样是更换扩散模型噪声,如果是从推理时扩展的角度来看,那么运算时间长就不是缺陷了;但与之相对,我们需要证明随着运算时间的增加,生成质量会按某种规律逐渐变好,表明花更多的时间是有意义的。
从这篇论文投稿失败的教训中,我们也可以总结出顶会论文应该满足哪些最基本的要求。当论文的方法和之前方法非常相似时,一定要与之前的方法对比,并明确写出本文的创新点在哪。另外,虽然很多时候我们是先有了实验结果再去补充说明方法中的假设,但我们依然可以在论文里尝试解释我们这样做的原因。这种解释可以是通过原理上的推导,也可以是通过解释性实验。
#扩散模型(Diffusion Model)详解
直观理解、数学原理、PyTorch 实现
本文详细介绍了扩散模型的原理、数学基础以及基于PyTorch的实现方法。文章从图像生成任务的背景出发,对比了GAN和VAE等生成模型,深入讲解了扩散模型的正向和反向过程,推导了相关的数学公式,并提供了完整的代码实现和实验结果,帮助读者从直观到深入地理解扩散模型。
在过去的大半年里,以Stable Diffusion为代表的AI绘画是世界上最为火热的AI方向之一。或许大家会有疑问,Stable Diffusion里的这个"Diffusion"是什么意思?其实,扩散模型(Diffusion Model)正是Stable Diffusion中负责生成图像的模型。想要理解Stable Diffusion的原理,就一定绕不过扩散模型的学习。
Stable Diffusion以「毕加索笔下的《最后的晚餐》」为题的绘画结果
在这篇文章里,我会由浅入深地对最基础的去噪扩散概率模型(Denoising Diffusion Probabilistic Models, DDPM)进行讲解。我会先介绍扩散模型生成图像的基本原理,再用简单的数学语言对扩散模型建模,最后给出扩散模型的一份PyTorch实现。本文不会堆砌过于复杂的数学公式,哪怕你没有相关的数学背景,也能够轻松理解扩散模型的原理。
扩散模型与图像生成
在认识扩散模型之前,我们先退一步,看看一般的神经网络模型是怎么生成图像的。显然,为了生成丰富的图像,一个图像生成程序要根据随机数来生成图像。通常,这种随机数是一个满足标准正态分布的随机向量。这样,每次要生成新图像时,只需要从标准正态分布里随机生成一个向量并输入给程序就行了。
而在AI绘画程序中,负责生成图像的是一个神经网络模型。神经网络需要从数据中学习。对于图像生成任务,神经网络的训练数据一般是一些同类型的图片。比如一个绘制人脸的神经网络会用人脸照片来训练。也就是说,神经网络会学习如何把一个向量映射成一张图片,并确保这个图片和训练集的图片是一类图片。
可是,相比其他AI任务,图像生成任务对神经网络来说更加困难一点——图像生成任务缺乏有效的指导。在其他AI任务中,训练集本身会给出一个「标准答案」,指导AI的输出向标准答案靠拢。比如对于图像分类任务,训练集会给出每一幅图像的类别;对于人脸验证任务,训练集会给出两张人脸照片是不是同一个人;对于目标检测任务,训练集会给出目标的具体位置。然而,图像生成任务是没有标准答案的。图像生成数据集里只有一些同类型图片,却没有指导AI如何画得更好的信息。
为了解决这一问题,人们专门设计了一些用于生成图像的神经网络架构。这些架构中比较出名的有生成对抗模型(GAN)和变分自编码器(VAE)。
GAN的想法是,既然不知道一幅图片好不好,就干脆再训练一个神经网络,用于辨别某图片是不是和训练集里的图片长得一样。生成图像的神经网络叫做生成器,鉴定图像的神经网络叫做判别器。两个网络互相对抗,共同进步。
VAE则使用了逆向思维:学习向量生成图像很困难,那就再同时学习怎么用图像生成向量。这样,把某图像变成向量,再用该向量生成图像,就应该得到一幅和原图像一模一样的图像。每一个向量的绘画结果有了一个标准答案,可以用一般的优化方法来指导网络的训练了。VAE中,把图像变成向量的网络叫做编码器,把向量转换回图像的网络叫做解码器。其中,解码器就是负责生成图像的模型。
一直以来,GAN的生成效果较好,但训练起来比VAE麻烦很多。有没有和GAN一样强大,训练起来又方便的生成网络架构呢?扩散模型正是满足这些要求的生成网络架构。
扩散模型是一种特殊的VAE,其灵感来自于热力学:一个分布可以通过不断地添加噪声变成另一个分布。放到图像生成任务里,就是来自训练集的图像可以通过不断添加噪声变成符合标准正态分布的图像。从这个角度出发,我们可以对VAE做以下修改:1)不再训练一个可学习的编码器,而是把编码过程固定成不断添加噪声的过程;2)不再把图像压缩成更短的向量,而是自始至终都对一个等大的图像做操作。解码器依然是一个可学习的神经网络,它的目的也同样是实现编码的逆操作。不过,既然现在编码过程变成了加噪,那么解码器就应该负责去噪。而对于神经网络来说,去噪任务学习起来会更加有效。因此,扩散模型既不会涉及GAN中复杂的对抗训练,又比VAE更强大一点。
具体来说,扩散模型由正向过程和反向过程这两部分组成,对应VAE中的编码和解码。在正向过程中,输入 会不断混入高斯噪声。经过 次加噪声操作后,图像 会变成一幅符合标准正态分布的纯噪声图像。而在反向过程中,我们希望训练出一个神经网络,该网络能够学会 个去噪声操作,把 还原回 。网络的学习目标是让 个去噪声操作正好能抵消掉对应的加噪声操作。训练完毕后,只需要从标准正态分布里随机采样出一个噪声,再利用反向过程里的神经网络把该噪声恢复成一幅图像,就能够生成一幅图片了。
高斯噪声,就是一幅各处颜色值都满足高斯分布(正态分布)的噪声图像。
总结一下,图像生成网络会学习如何把一个向量映射成一幅图像。设计网络架构时,最重要的是设计学习目标,让网络生成的图像和给定数据集里的图像相似。VAE的做法是使用两个网络,一个学习把图像编码成向量,另一个学习把向量解码回图像,它们的目标是让复原图像和原图像尽可能相似。学习完毕后,解码器就是图像生成网络。扩散模型是一种更具体的VAE。它把编码过程固定为加噪声,并让解码器学习怎么样消除之前添加的每一步噪声。
扩散模型的具体算法
上一节中,我们只是大概了解扩散模型的整体思想。这一节,我们来引入一些数学表示,来看一看扩散模型的训练算法和采样算法具体是什么。为了便于理解,这一节会出现一些不是那么严谨的数学描述。更加详细的一些数学推导会放到下一节里介绍。
前向过程
在前向过程中,来自训练集的图像 会被添加 次噪声,使得 为符合标准正态分布。准确来说,「加噪声」并不是给上一时刻的图像加上噪声值,而是从一个均值与上一时刻图像相关的正态分布里采样出一幅新图像。如下面的公式所示, 是上一时刻的图像, 是这一时刻生成的图像,该图像是从一个均值与 有关的正态分布里采样出来的。
多数文章会说前向过程是一个马尔可夫过程。其实,马尔可夫过程的意思就是当前时刻的状态只由上一时刻的状态决定,而不由更早的状态决定。上面的公式表明,计算 ,只需要用到 ,而不需要用到 ,这符合马尔可夫过程的定义。
绝大多数扩散模型会把这个正态分布设置成这个形式:
这个正态分布公式乍看起来很奇怪: 是哪里冒出来的?为什么会有这种奇怪的系数?别急,我们先来看另一个问题:假如给定 ,也就是从训练集里采样出一幅图片,该怎么计算任意一个时刻 的噪声图像 呢?
我们不妨按照公式,从 开始倒推。 其实可以通过一个标准正态分布的样本 算出来:
再往前推几步:
由正态分布的性质可知,均值相同的正态分布「加」在一起后,方差也会加到一起。也就是 与 合起来会得到 。根据这一性质,上面的公式可以化简为:
再往前推一步的话,结果是:
我们已经能够猜出规律来了,可以一直把公式推到 。令 ,则:
有了这个公式,我们就可以讨论加噪声公式为什么是 了。这个公式里的 是一个小于 1 的常数。在DDPM论文中, 从 到 线性增长。这样, 变大, 也越小, 趋于 0 的速度越来越快。最后, 几乎为 0 ,代入 就满足标准正态分布了,符合我们对扩散模型的要求。上述推断可以简单描述为:加噪声公式能够从慢到快地改变原图像,让图像最终均值为 0 ,方差为 。
反向过程
在正向过程中,我们人为设置了步加噪声过程。而在反向过程中,我们希望能够倒过来取消每一步加噪声操作,让一幅纯噪声图像变回数据集里的图像。这样,利用这个去噪声过程,我们就可以把任意一个从标准正态分布里采样出来的噪声图像变成一幅和训练数据长得差不多的图像,从而起到图像生成的目的。
现在问题来了:去噪声操作的数学形式是怎么样的?怎么让神经网络来学习它呢?数学原理表明,当足够小时,每一步加噪声的逆操作也满足正态分布。
其中,当前时刻加噪声逆操作的均值 和方差 由当前的时刻 ,当前的图像 决定。因此,为了描述所有去噪声操作,神经网络应该输入 ,拟合当前的均值 和方差 。
不要被上文的「去噪声」、「加噪声逆操作」绕晕了哦。由于加噪声是固定的,加噪声的逆操作也是固定的。理想情况下,我们希望去噪操作就等于加噪声逆操作。然而,加噪声的逆操作不太可能从理论上求得,我们只能用一个神经网络去拟合它。去噪声操作和加噪声逆操作的关系,就是神经网络的预测值和真值的关系。
现在问题来了:加噪声逆操作的均值和方差是什么?
直接计算所有数据的加噪声逆操作的分布是不太现实的。但是,如果给定了某个训练集输入 ,多了一个限定条件后,该分布是可以用贝叶斯公式计算的(其中 表示概率分布):
等式左边的 表示加噪声操作的逆操作,它的均值和方差都是待求的。右边的 是加噪声的分布。而由于 已知, 和 两项可以根据前面的公式 得来:
这样,等式右边的式子全部已知。我们可以把公式套入,算出给定时的去噪声分布。经计算化简,分布的均值为:
其中, 是用公式算 时从标准正态分布采样出的样本,它来自公式
分布的方差为:
注意, 是加噪声的方差,是一个常量。那么,加噪声逆操作的方差 也是一个常量,不与输入 相关。这下就省事了,训练去噪网络时,神经网络只用拟合 个均值就行,不用再拟合方差了。
知道了均值和方差的真值,训练神经网络只差最后的问题了:该怎么设置训练的损失函数?加噪声逆操作和去噪声操作都是正态分布,网络的训练目标应该是让每对正态分布更加接近。那怎么用损失函数描述两个分布尽可能接近呢?最直观的想法,肯定是让两个正态分布的均值尽可能接近,方差尽可能接近。根据上文的分析,方差是常量,只用让均值尽可能接近就可以了。
那怎么用数学公式表达让均值更接近呢?再观察一下目标均值的公式:
神经网络拟合均值时, 是已知的(别忘了,图像是一步一步倒着去噪的)。式子里唯一不确定的只有 。既然如此,神经网络干脆也别预测均值了,直接预测一个噪声 (其中 为可学习参数),让它和生成 的噪声 的均方误差最小就行了。对于一轮训练,最终的误差函数可以写成
这样,我们就认识了反向过程的所有内容。总结一下,反向过程中,神经网络应该让 个去噪声操作拟合对应的 个加噪声逆操作。每步加噪声逆操作符合正态分布,且在给定某个输入时,该正态分布的均值和方差是可以用解析式表达出来的。因此,神经网络的学习目标就是让其输出的去噪声分布和理论计算的加噪声逆操作分布一致。经过数学计算上的一些化简,问题被转换成了拟合生成 时用到的随机噪声 。
训练算法与采样算法
理解了前向过程和反向过程后,训练神经网络的算法和采样图片(生成图片)的算法就呼之欲出了。
以下是DDPM论文中的训练算法:
让我们来逐行理解一下这个算法。第二行是指从训练集里取一个数据 。第三行是指随机从 里取一个时刻用来训练。我们虽然要求神经网络拟合 个正态分布,但实际训练时,不用一轮预测 个结果,只需要随机预测 个时刻中某一个时刻的结果就行。第四行指随机生成一个噪声 ,该噪声是用于执行前向过程生成 的。之后,我们把 和 传给神经网络 ,让神经网络预测随机噪声。训练的损失函数是预测噪声和实际噪声之间的均方误差,对此损失函数采用梯度下降即可优化网络。
DDPM并没有规定神经网络的结构。根据任务的难易程度,我们可以自己定义简单或复杂的网络结构。这里只需要把 当成一个普通的映射即可。
训练好了网络后,我们可以执行反向过程,对任意一幅噪声图像去噪,以实现图像生成。这个算法如下:
第一行的 就是从标准正态分布里随机采样的输入噪声。要生成不同的图像,只需要更换这个噪声。后面的过程就是扩散模型的反向过程。令时刻从 到 1 ,计算这一时刻去噪声操作的均值和方差,并采样出 。均值是用之前提到的公式计算的:
而方差 的公式有两种选择,两个公式都能产生差不多的结果。实验表明,当 是特定的某个数据时,用上一节推导出来的方差最好。
而当 时,只需要令方差和加噪声时的方差一样即可。
循环执行去噪声操作。最后生成的 就是生成出来的图像。特别地,最后一步去噪声是不用加方差项的。为什么呢,观察公式 。当 时,分子会出现 这一项。 是一个连乘,理论上 是从 1 开始的,在 时没有定义。但我们可以特别地令连乘的第 0 项 。这样, 时方差项的分子 为 0 ,不用算这一项了。
当然,这一解释从数学上来说是不严谨的。据论文说,这部分的解释可以参见朗之万动力学。
数学推导的补充 (选读)
理解了训练算法和采样算法,我们就算是搞懂了扩散模型,可以去编写代码了。不过,上文的描述省略了一些数学推导的细节。如果对扩散模型更深的原理感兴趣,可以阅读一下本节。
加噪声逆操作均值和方差的推导
上一节,我们根据下面几个式子
一步就给出了的均值和方差。
现在我们来看一下推导均值和方差的思路。
首先,把其他几个式子带入贝叶斯公式的等式右边。
由于多个正态分布的乘积还是一个正态分布,我们知道 也可以用一个正态分布公式 表达,它最后一定能写成这种形式:
问题就变成了怎么把开始那个很长的式子化简,算出 和 。
方差 可以从指数函数的系数得来,比较好求。系数为
所以,方差为:
接下来只要关注指数函数的指数部分。指数部分一定是一个关于的 的二次函数,只要化简成 的形式,再除以一下 -2 倍方差,就可以得到均值了。
指数部分为:
只在前两项里有。把和 有关的项计算化简,可以计算出均值:
回想一下,在去噪声中,神经网络的输入是 和 。也就是说,上式中 已知,只有 一个未知量。要算均值,还需要算出 。 和 之间是有一定联系的。 是 在正向过程中第 步加噪声的结果。而根据正向过程的公式倒推:
把这个 带入均值公式,均值最后会化简成我们熟悉的形式。
优化目标
上一节,我们只是简单地说神经网络的优化目标是让加噪声和去噪声的均值接近。而让均值接近,就是让生成 的噪声 更接近。实际上,这个优化目标是经过简化得来的。扩散模型最早的优化目标是有一定的数学意义的。
扩散模型,全称为扩散概率模型(Diffusion Probabilistic Model)。最简单的一类扩散模型,是去噪扩散概率模型(Denoising Diffusion Probabilistic Model),也就是常说的DDPM。DDPM的框架主要是由两篇论文建立起来的。第一篇论文是首次提出扩散模型思想的_Deep Unsupervised Learning using Nonequilibrium Thermodynamics_。在此基础上,_Denoising Diffusion Probabilistic Models_对最早的扩散模型做出了一定的简化,让图像生成效果大幅提升,促成了扩散模型的广泛使用。我们上一节看到的公式,全部是简化后的结果。
扩散概率模型的名字之所以有「概率」二字,是因为这个模型是在描述一个系统的概率。准确来说,扩散模型是在描述经反向过程生成出某一项数据的概率。也就是说,扩散模型 是一个有着可训练参数 的模型,它描述了反向过程生成出数据 的概率。 满足 ,其中 就是我们熟悉的反向过程,只不过它是以概率计算的形式表达:
我们上一节里见到的优化目标,是让去噪声操作 和加噪声操作的逆操作 尽可能相似。然而,这个描述并不确切。扩散模型原本的目标,是最大化 这个概率,其中 是来自训练集的数据。换个角度说,给定一个训练集的数据 ,经过前向过程和反向过程,扩散模型要让复原出 的概率尽可能大。这也是我们在本文开头认识VAE时见到的优化目标。
最大化 ,一般会写成最小化其负对数值,即最小化 。使用和VAE类似的变分推理,可以把优化目标转换成优化一个叫做变分下界(variational lower bound,VLB)的量。它最终可以写成:
这里的表示分布P和Q之间的KL散度。KL散度是衡量两个分布相似度的指标。如果都是正态分布,则它们的KL散度可以由一个简单的公式给出。
其中,第一项 和可学习参数 无关(因为可学习参数只描述了每一步去噪声操作,也就是只描述了 ),可以不去管它。那么这个优化目标就由两部分组成:
1.最小化 表示的是最大化每一个去噪声操作和加噪声逆操作的相似度。
2.最小化 就是已知 时,让最后复原原图 概率更高。
我们分别看这两部分是怎么计算的。
对于第一部分,我们先回顾一下正态分布之间的KL散度公式。设一维正态分布 的公式如下:
则
而对于 ,根据前文的分析,我们知道,待求方差 可以直接由计算得到。
两个正态分布方差的比值是常量。所以,在计算KL散度时,不用管方差那一项了,只需要管均值那一项。
由根据之前的均值公式
这一部分的优化目标可以化简成
DDPM论文指出,如果把前面的系数全部丢掉的话,模型的效果更好。最终,我们就能得到一个非常简单的优化目标:
这就是我们上一节见到的优化目标。
当然,还没完,别忘了优化目标里还有这一项。它的形式为:
只管后面有的那一项(注意,):
这和那些 散度项 时的形式相同,我们可以用相同的方式简化优化目标,只保留 。这样,损失函数的形式全都是 了。
DDPM论文里写 这一项可以直接满足简化后的公式 时的情况,而没有去掉系数的过程。我在网上没找到文章解释这一点,只好按自己的理解来推导这个误差项了。不论如何,推导的过程不是那么重要,重要的是最后的简化形式。
总结
图像生成任务就是把随机生成的向量(噪声)映射成和训练图像类似的图像。为此,扩散模型把这个过程看成是对纯噪声图像的去噪过程。通过学习把图像逐步变成纯噪声的逆操作,扩散模型可以把任何一个纯噪声图像变成有意义的图像,也就是完成图像生成。
对于不同程度的读者,应该对本文有不同的认识。
对于只想了解扩散模型大概原理的读者,只需要阅读第一节,并大概了解:
- 图像生成任务的通常做法
- 图像生成任务需要监督
- VAE通过把图像编码再解码来训练一个解码器
- 扩散模型是一类特殊的VAE,它的编码固定为加噪声,解码固定为去噪声
对于想认真学习扩散模型的读者,只需读懂第二节的主要内容:
- 扩散模型的优化目标:让反向过程尽可能成为正向过程的逆操作
- 正向过程的公式
- 反向过程的做法(采样算法)
- 加噪声逆操作的均值和方差在给定时可以求出来的,加噪声逆操作的均值就是去噪声的学习目标
- 简化后的损失函数与训练算法
对有学有余力对数学感兴趣的读者,可以看一看第三节的内容:
- 加噪声逆操作均值和方差的推导
- 扩散模型最早的优化目标与DDPM论文是如何简化优化目标的
我个人认为,由于扩散模型的优化目标已经被大幅度简化,除非你的研究目标是改进扩散模型本身,否则没必要花过多的时间钻研数学原理。在学习时,建议快点看懂扩散模型的整体思想,搞懂最核心的训练算法和采样算法,跑通代码。之后就可以去看较新的论文了。
在附录中,我给出了一份DDPM的简单实现。欢迎大家参考,并自己动手复现一遍DDPM。
参考资料与学习建议
网上绝大多数的中英文教程都是照搬 https://lilianweng.github.io/posts/2021-07-11-diffusion-models/ 这篇文章的。这篇文章像教科书一样严谨,适合有一定数学基础的人阅读,但不适合给初学者学习。建议在弄懂扩散模型的大概原理后再来阅读这篇文章补充细节。
多数介绍扩散模型的文章对没学过相关数学知识的人来说很不友好,我在阅读此类文章时碰到了大量的问题:为什么前向公式里有个 ?为什么突然冒出一个快速算 的公式?为什么反向过程里来了个贝叶斯公式?优化目标是什么? 是什么?为什么优化目标里一大堆项,每一项的意义又是什么?为什么最后莫名其妙算一个 ?为什么采样时 就不用加方差项了?好不容易,我才把这些问题慢慢搞懂,并在本文做出了解释。希望我的解答能够帮助到同样有这些困惑的读者。想逐步学习扩散模型,可以先看懂我这篇文章的大概讲解,再去其他文章里学懂一些细节。无论是教,还是学,最重要的都是搞懂整体思路,知道动机,最后再去强调细节。
这里还有篇文章给出了扩散模型中数学公式的详细推导,并补充了变分推理的背景介绍,适合从头学起:https://arxiv.org/abs/2208.11970
想深入学习DDPM,可以看一看最重要的两篇论文:Deep Unsupervised Learning using Nonequilibrium Thermodynamics、Denoising Diffusion Probabilistic Models。当然,后者更重要一些,里面的一些实验结果仍有阅读价值。
我在代码复现时参考了这篇文章(https://medium.com/mlearning-ai/enerating-images-with-ddpms-a-pytorch-implementation-cef5a2ba8cb1)。相对于网上的其他开源DDPM实现,这份代码比较简短易懂,更适合学习。不过,这份代码有一点问题。它的神经网络不够强大,采样结果会有一点问题。
附录:代码复现
在这个项目中,我们要用PyTorch实现一个基于U-Net的DDPM,并在MNIST数据集(经典的手写数字数据集)上训练它。模型几分钟就能训练完,我们可以方便地做各种各样的实验。
后续讲解只会给出代码片段,完整的代码请参见 https://github.com/SingleZombie/DL-Demos/tree/master/dldemos/ddpm 。git clone 仓库并安装后,可以直接运行目录里的main.py
训练模型并采样。
获取数据集
PyTorch的torchvision
提供了获取了MNIST的接口,我们只需要用下面的函数就可以生成MNIST的Dataset
实例。参数中,root
为数据集的下载路径,download
为是否自动下载数据集。令download=True
的话,第一次调用该函数时会自动下载数据集,而第二次之后就不用下载了,函数会读取存储在root
里的数据。
mnist = torchvision.datasets.MNIST(root='./data/mnist', download=True)
我们可以用下面的代码来下载MNIST并输出该数据集的一些信息:
import torchvision
from torchvision.transforms import ToTensor
def download_dataset():
mnist = torchvision.datasets.MNIST(root='./data/mnist', download=True)
print('length of MNIST', len(mnist))
id = 4
img, label = mnist[id]
print(img)
print(label)
# On computer with monitor
# img.show()
img.save('work_dirs/tmp.jpg')
tensor = ToTensor()(img)
print(tensor.shape)
print(tensor.max())
print(tensor.min())
if __name__ == '__main__':
download_dataset()
执行这段代码,输出大致为:
length of MNIST 60000
<PIL.Image.Image image mode=L size=28x28 at 0x7FB3F09CCE50>
9
torch.Size([1, 28, 28])
tensor(1.)
tensor(0.)
第一行输出表明,MNIST数据集里有60000张图片。而从第二行和第三行输出中,我们发现每一项数据由图片和标签组成,图片是大小为28x28
的PIL格式的图片,标签表明该图片是哪个数字。我们可以用torchvision
里的ToTensor()
把PIL图片转成PyTorch张量,进一步查看图片的信息。最后三行输出表明,每一张图片都是单通道图片(灰度图),颜色值的取值范围是0~1。
我们可以查看一下每张图片的样子。如果你是在用带显示器的电脑,可以去掉img.show
那一行的注释,直接查看图片;如果你是在用服务器,可以去img.save
的路径里查看图片。该图片的应该长这个样子:
我们可以用下面的代码预处理数据并创建DataLoader
。由于DDPM会把图像和正态分布关联起来,我们更希望图像颜色值的取值范围是[-1, 1]
。为此,我们可以对图像做一个线性变换,减0.5再乘2。
def get_dataloader(batch_size: int):
transform = Compose([ToTensor(), Lambda(lambda x: (x - 0.5) * 2)])
dataset = torchvision.datasets.MNIST(root='./data/mnist',
transform=transform)
return DataLoader(dataset, batch_size=batch_size, shuffle=True)
DDPM 类
在代码中,我们要实现一个DDPM
类。它维护了扩散过程中的一些常量(比如),并且可以计算正向过程和反向过程的结果。
先来实现一下DDPM
类的初始化函数。一开始,我们遵从论文的配置,用torch.linspace(min_beta, max_beta, n_steps)
从min_beta
到max_beta
线性地生成n_steps
个时刻的。接着,我们根据公式,计算每个时刻的alpha
和alpha_bar
。注意,为了方便实现,我们让t
的取值从0开始,要比论文里的少1。
import torch
class DDPM():
# n_steps 就是论文里的 T
def __init__(self,
device,
n_steps: int,
min_beta: float = 0.0001,
max_beta: float = 0.02):
betas = torch.linspace(min_beta, max_beta, n_steps).to(device)
alphas = 1 - betas
alpha_bars = torch.empty_like(alphas)
product = 1
for i, alpha in enumerate(alphas):
product *= alpha
alpha_bars[i] = product
self.betas = betas
self.n_steps = n_steps
self.alphas = alphas
self.alpha_bars = alpha_bars
部分实现会让 DDPM 继承
torch.nn.Module
,但我认为这样不好。DDPM本身不是一个神经网络,它只是描述了前向过程和后向过程的一些计算。只有涉及可学习参数的神经网络类才应该继承 torch.nn.Module
。
准备好了变量后,我们可以来实现DDPM
类的其他方法。先实现正向过程方法,该方法会根据公式 计算正向过程中的 。
def sample_forward(self, x, t, eps=None):
alpha_bar = self.alpha_bars[t].reshape(-1, 1, 1, 1)
if eps is None:
eps = torch.randn_like(x)
res = eps * torch.sqrt(1 - alpha_bar) + torch.sqrt(alpha_bar) * x
return res
这里要解释一些PyTorch编程上的细节。这份代码中,self.alpha_bars
是一个一维Tensor
。而在并行训练中,我们一般会令t
为一个形状为(batch_size, )
的Tensor
。PyTorch允许我们直接用self.alpha_bars[t]
从self.alpha_bars
里取出batch_size
个数,就像用一个普通的整型索引来从数组中取出一个数一样。有些实现会用torch.gather
从self.alpha_bars
里取数,其作用是一样的。
我们可以随机从训练集取图片做测试,看看它们在前向过程中是怎么逐步变成噪声的。
接下来实现反向过程。在反向过程中,DDPM会用神经网络预测每一轮去噪的均值,把 复原回 ,以完成图像生成。反向过程即对应论文中的采样算法。
其实现如下:
def sample_backward(self, img_shape, net, device, simple_var=True):
x = torch.randn(img_shape).to(device)
net = net.to(device)
for t in range(self.n_steps - 1, -1, -1):
x = self.sample_backward_step(x, t, net, simple_var)
return x
def sample_backward_step(self, x_t, t, net, simple_var=True):
n = x_t.shape[0]
t_tensor = torch.tensor([t] * n,
dtype=torch.long).to(x_t.device).unsqueeze(1)
eps = net(x_t, t_tensor)
if t == 0:
noise = 0
else:
if simple_var:
var = self.betas[t]
else:
var = (1 - self.alpha_bars[t - 1]) / (
1 - self.alpha_bars[t]) * self.betas[t]
noise = torch.randn_like(x_t)
noise *= torch.sqrt(var)
mean = (x_t -
(1 - self.alphas[t]) / torch.sqrt(1 - self.alpha_bars[t]) *
eps) / torch.sqrt(self.alphas[t])
x_t = mean + noise
return x_t
其中,sample_backward
是用来给外部调用的方法,而sample_backward_step
是执行一步反向过程的方法。
sample_backward
会随机生成纯噪声x
(对应),再令t
从n_steps - 1
到0
,调用sample_backward_step
。
def sample_backward(self, img_shape, net, device, simple_var=True):
x = torch.randn(img_shape).to(device)
net = net.to(device)
for t in range(self.n_steps - 1, -1, -1):
x = self.sample_backward_step(x, t, net, simple_var)
return x
在sample_backward_step
中,我们先准备好这一步的神经网络输出eps
。为此,我们要把整型的t
转换成一个格式正确的Tensor
。考虑到输入里可能有多个batch,我们先获取batch size n
,再根据它来生成t_tensor
。
def sample_backward_step(self, x_t, t, net, simple_var=True):
n = x_t.shape[0]
t_tensor = torch.tensor([t] * n,
dtype=torch.long).to(x_t.device).unsqueeze(1)
eps = net(x_t, t_tensor)
之后,我们来处理反向过程公式中的方差项。根据伪代码,我们仅在t
非零的时候算方差项。方差项用到的方差有两种取值,效果差不多,我们用simple_var
来控制选哪种取值方式。获取方差后,我们再随机采样一个噪声,根据公式,得到方差项。
if t == 0:
noise = 0
else:
if simple_var:
var = self.betas[t]
else:
var = (1 - self.alpha_bars[t - 1]) / (
1 - self.alpha_bars[t]) * self.betas[t]
noise = torch.randn_like(x_t)
noise *= torch.sqrt(var)
最后,我们把eps
和方差项套入公式,得到这一步更新过后的图像x_t
。
mean = (x_t -
(1 - self.alphas[t]) / torch.sqrt(1 - self.alpha_bars[t]) *
eps) / torch.sqrt(self.alphas[t])
x_t = mean + noise
return x_t
稍后完成了训练后,我们再来看反向过程的输出结果。
训练算法
接下来,我们先跳过神经网络的实现,直接完成论文里的训练算法。
再回顾一遍伪代码。首先,我们要随机选取训练图片 ,随机生成当前要训练的时刻 ,以及随机生成一个生成 的高斯噪声。之后,我们把 和 输入进神经网络,尝试预测噪声。最后,我们以预测噪声和实际噪声的均方误差为损失函数做梯度下降。
为此,我们可以用下面的代码实现训练。
import torch
import torch.nn as nn
from dldemos.ddpm.dataset import get_dataloader, get_img_shape
from dldemos.ddpm.ddpm import DDPM
import cv2
import numpy as np
import einops
batch_size = 512
n_epochs = 100
def train(ddpm: DDPM, net, device, ckpt_path):
# n_steps 就是公式里的 T
# net 是某个继承自 torch.nn.Module 的神经网络
n_steps = ddpm.n_steps
dataloader = get_dataloader(batch_size)
net = net.to(device)
loss_fn = nn.MSELoss()
optimizer = torch.optim.Adam(net.parameters(), 1e-3)
for e in range(n_epochs):
for x, _ in dataloader:
current_batch_size = x.shape[0]
x = x.to(device)
t = torch.randint(0, n_steps, (current_batch_size, )).to(device)
eps = torch.randn_like(x).to(device)
x_t = ddpm.sample_forward(x, t, eps)
eps_theta = net(x_t, t.reshape(current_batch_size, 1))
loss = loss_fn(eps_theta, eps)
optimizer.zero_grad()
loss.backward()
optimizer.step()
torch.save(net.state_dict(), ckpt_path)
代码的主要逻辑都在循环里。首先是完成训练数据 , ,噪声的采样。采样 的工作可以交给PyTorch的DataLoader完成,每轮遍历得到的x
就是训练数据。的采样可以用torch.randint
函数随机从[0, n_steps - 1]
取数。采样高斯噪声可以直接用torch.randn_like(x)
生成一个和训练图片x
形状一样的符合标准正态分布的图像。
for x, _ in dataloader:
current_batch_size = x.shape[0]
x = x.to(device)
t = torch.randint(0, n_steps, (current_batch_size, )).to(device)
eps = torch.randn_like(x).to(device)
之后计算 并将其和 输入进神经网络 net 。计算 的任务会由DDPM
类的sample_forward
方法完成,我们在上文已经实现了它。
x_t = ddpm.sample_forward(x, t, eps)
eps_theta = net(x_t, t.reshape(current_batch_size, 1))
得到了预测的噪声eps_theta
,我们调用PyTorch的API,算均方误差并调用优化器即可。
loss_fn = nn.MSELoss()
optimizer = torch.optim.Adam(net.parameters(), 1e-3)
...
loss = loss_fn(eps_theta, eps)
optimizer.zero_grad()
loss.backward()
optimizer.step()
去噪神经网络
在DDPM中,理论上我们可以用任意一种神经网络架构。但由于DDPM任务十分接近图像去噪任务,而U-Net又是去噪任务中最常见的网络架构,因此绝大多数DDPM都会使用基于U-Net的神经网络。
我一直想训练一个尽可能简单的模型。经过多次实验,我发现DDPM的神经网络很难训练。哪怕是对于比较简单的MNIST数据集,结构差一点的网络(比如纯ResNet)都不太行,只有带了残差块和时序编码的U-Net才能较好地完成去噪。注意力模块倒是可以不用加上。
由于神经网络结构并不是DDPM学习的重点,我这里就不对U-Net的写法做解说,而是直接贴上代码了。代码中大部分内容都和普通的U-Net无异。唯一要注意的地方就是时序编码。去噪网络的输入除了图像外,还有一个时间戳t。我们要考虑怎么把t的信息和输入图像信息融合起来。大部分人的做法是对t进行Transformer中的位置编码,把该编码加到图像的每一处上。
import torch
import torch.nn as nn
import torch.nn.functional as F
from dldemos.ddpm.dataset import get_img_shape
class PositionalEncoding(nn.Module):
def __init__(self, max_seq_len: int, d_model: int):
super().__init__()
# Assume d_model is an even number for convenience
assert d_model % 2 == 0
pe = torch.zeros(max_seq_len, d_model)
i_seq = torch.linspace(0, max_seq_len - 1, max_seq_len)
j_seq = torch.linspace(0, d_model - 2, d_model // 2)
pos, two_i = torch.meshgrid(i_seq, j_seq)
pe_2i = torch.sin(pos / 10000**(two_i / d_model))
pe_2i_1 = torch.cos(pos / 10000**(two_i / d_model))
pe = torch.stack((pe_2i, pe_2i_1), 2).reshape(max_seq_len, d_model)
self.embedding = nn.Embedding(max_seq_len, d_model)
self.embedding.weight.data = pe
self.embedding.requires_grad_(False)
def forward(self, t):
return self.embedding(t)
class ResidualBlock(nn.Module):
def __init__(self, in_c: int, out_c: int):
super().__init__()
self.conv1 = nn.Conv2d(in_c, out_c, 3, 1, 1)
self.bn1 = nn.BatchNorm2d(out_c)
self.actvation1 = nn.ReLU()
self.conv2 = nn.Conv2d(out_c, out_c, 3, 1, 1)
self.bn2 = nn.BatchNorm2d(out_c)
self.actvation2 = nn.ReLU()
if in_c != out_c:
self.shortcut = nn.Sequential(nn.Conv2d(in_c, out_c, 1),
nn.BatchNorm2d(out_c))
else:
self.shortcut = nn.Identity()
def forward(self, input):
x = self.conv1(input)
x = self.bn1(x)
x = self.actvation1(x)
x = self.conv2(x)
x = self.bn2(x)
x += self.shortcut(input)
x = self.actvation2(x)
return x
class ConvNet(nn.Module):
def __init__(self,
n_steps,
intermediate_channels=[10, 20, 40],
pe_dim=10,
insert_t_to_all_layers=False):
super().__init__()
C, H, W = get_img_shape() # 1, 28, 28
self.pe = PositionalEncoding(n_steps, pe_dim)
self.pe_linears = nn.ModuleList()
self.all_t = insert_t_to_all_layers
if not insert_t_to_all_layers:
self.pe_linears.append(nn.Linear(pe_dim, C))
self.residual_blocks = nn.ModuleList()
prev_channel = C
for channel in intermediate_channels:
self.residual_blocks.append(ResidualBlock(prev_channel, channel))
if insert_t_to_all_layers:
self.pe_linears.append(nn.Linear(pe_dim, prev_channel))
else:
self.pe_linears.append(None)
prev_channel = channel
self.output_layer = nn.Conv2d(prev_channel, C, 3, 1, 1)
def forward(self, x, t):
n = t.shape[0]
t = self.pe(t)
for m_x, m_t in zip(self.residual_blocks, self.pe_linears):
if m_t is not None:
pe = m_t(t).reshape(n, -1, 1, 1)
x = x + pe
x = m_x(x)
x = self.output_layer(x)
return x
class UnetBlock(nn.Module):
def __init__(self, shape, in_c, out_c, residual=False):
super().__init__()
self.ln = nn.LayerNorm(shape)
self.conv1 = nn.Conv2d(in_c, out_c, 3, 1, 1)
self.conv2 = nn.Conv2d(out_c, out_c, 3, 1, 1)
self.activation = nn.ReLU()
self.residual = residual
if residual:
if in_c == out_c:
self.residual_conv = nn.Identity()
else:
self.residual_conv = nn.Conv2d(in_c, out_c, 1)
def forward(self, x):
out = self.ln(x)
out = self.conv1(out)
out = self.activation(out)
out = self.conv2(out)
if self.residual:
out += self.residual_conv(x)
out = self.activation(out)
return out
class UNet(nn.Module):
def __init__(self,
n_steps,
channels=[10, 20, 40, 80],
pe_dim=10,
residual=False) -> None:
super().__init__()
C, H, W = get_img_shape()
layers = len(channels)
Hs = [H]
Ws = [W]
cH = H
cW = W
for _ in range(layers - 1):
cH //= 2
cW //= 2
Hs.append(cH)
Ws.append(cW)
self.pe = PositionalEncoding(n_steps, pe_dim)
self.encoders = nn.ModuleList()
self.decoders = nn.ModuleList()
self.pe_linears_en = nn.ModuleList()
self.pe_linears_de = nn.ModuleList()
self.downs = nn.ModuleList()
self.ups = nn.ModuleList()
prev_channel = C
for channel, cH, cW in zip(channels[0:-1], Hs[0:-1], Ws[0:-1]):
self.pe_linears_en.append(
nn.Sequential(nn.Linear(pe_dim, prev_channel), nn.ReLU(),
nn.Linear(prev_channel, prev_channel)))
self.encoders.append(
nn.Sequential(
UnetBlock((prev_channel, cH, cW),
prev_channel,
channel,
residual=residual),
UnetBlock((channel, cH, cW),
channel,
channel,
residual=residual)))
self.downs.append(nn.Conv2d(channel, channel, 2, 2))
prev_channel = channel
self.pe_mid = nn.Linear(pe_dim, prev_channel)
channel = channels[-1]
self.mid = nn.Sequential(
UnetBlock((prev_channel, Hs[-1], Ws[-1]),
prev_channel,
channel,
residual=residual),
UnetBlock((channel, Hs[-1], Ws[-1]),
channel,
channel,
residual=residual),
)
prev_channel = channel
for channel, cH, cW in zip(channels[-2::-1], Hs[-2::-1], Ws[-2::-1]):
self.pe_linears_de.append(nn.Linear(pe_dim, prev_channel))
self.ups.append(nn.ConvTranspose2d(prev_channel, channel, 2, 2))
self.decoders.append(
nn.Sequential(
UnetBlock((channel * 2, cH, cW),
channel * 2,
channel,
residual=residual),
UnetBlock((channel, cH, cW),
channel,
channel,
residual=residual)))
prev_channel = channel
self.conv_out = nn.Conv2d(prev_channel, C, 3, 1, 1)
def forward(self, x, t):
n = t.shape[0]
t = self.pe(t)
encoder_outs = []
for pe_linear, encoder, down in zip(self.pe_linears_en, self.encoders,
self.downs):
pe = pe_linear(t).reshape(n, -1, 1, 1)
x = encoder(x + pe)
encoder_outs.append(x)
x = down(x)
pe = self.pe_mid(t).reshape(n, -1, 1, 1)
x = self.mid(x + pe)
for pe_linear, decoder, up, encoder_out in zip(self.pe_linears_de,
self.decoders, self.ups,
encoder_outs[::-1]):
pe = pe_linear(t).reshape(n, -1, 1, 1)
x = up(x)
pad_x = encoder_out.shape[2] - x.shape[2]
pad_y = encoder_out.shape[3] - x.shape[3]
x = F.pad(x, (pad_x // 2, pad_x - pad_x // 2, pad_y // 2,
pad_y - pad_y // 2))
x = torch.cat((encoder_out, x), dim=1)
x = decoder(x + pe)
x = self.conv_out(x)
return x
convnet_small_cfg = {
'type': 'ConvNet',
'intermediate_channels': [10, 20],
'pe_dim': 128
}
convnet_medium_cfg = {
'type': 'ConvNet',
'intermediate_channels': [10, 10, 20, 20, 40, 40, 80, 80],
'pe_dim': 256,
'insert_t_to_all_layers': True
}
convnet_big_cfg = {
'type': 'ConvNet',
'intermediate_channels': [20, 20, 40, 40, 80, 80, 160, 160],
'pe_dim': 256,
'insert_t_to_all_layers': True
}
unet_1_cfg = {'type': 'UNet', 'channels': [10, 20, 40, 80], 'pe_dim': 128}
unet_res_cfg = {
'type': 'UNet',
'channels': [10, 20, 40, 80],
'pe_dim': 128,
'residual': True
}
def build_network(config: dict, n_steps):
network_type = config.pop('type')
if network_type == 'ConvNet':
network_cls = ConvNet
elif network_type == 'UNet':
network_cls = UNet
network = network_cls(n_steps, **config)
return network
实验结果与采样
把之前的所有代码综合一下,我们以带残差块的U-Net为去噪网络,执行训练。
if __name__ == '__main__':
n_steps = 1000
config_id = 4
device = 'cuda'
model_path = 'dldemos/ddpm/model_unet_res.pth'
config = unet_res_cfg
net = build_network(config, n_steps)
ddpm = DDPM(device, n_steps)
train(ddpm, net, device=device, ckpt_path=model_path)
按照默认训练配置,在3090上花5分钟不到,训练30~40个epoch即可让网络基本收敛。最终收敛时loss在0.023~0.024左右。
batch size: 512
epoch 0 loss: 0.23103461712201437 elapsed 7.01s
epoch 1 loss: 0.0627968365987142 elapsed 13.66s
epoch 2 loss: 0.04828845852613449 elapsed 20.25s
epoch 3 loss: 0.04148937337398529 elapsed 26.80s
epoch 4 loss: 0.03801360730528831 elapsed 33.37s
epoch 5 loss: 0.03604260584712028 elapsed 39.96s
epoch 6 loss: 0.03357676289876302 elapsed 46.57s
epoch 7 loss: 0.0335664684087038 elapsed 53.15s
...
epoch 30 loss: 0.026149748386939366 elapsed 204.64s
epoch 31 loss: 0.025854381563266117 elapsed 211.24s
epoch 32 loss: 0.02589433005253474 elapsed 217.84s
epoch 33 loss: 0.026276464049021404 elapsed 224.41s
...
epoch 96 loss: 0.023299352884292603 elapsed 640.25s
epoch 97 loss: 0.023460942271351815 elapsed 646.90s
epoch 98 loss: 0.023584651704629263 elapsed 653.54s
epoch 99 loss: 0.02364126600921154 elapsed 660.22s
训练这个网络时,并没有特别好的测试指标,我们只能通过观察采样图像来评价网络的表现。我们可以用下面的代码调用DDPM的反向传播方法,生成多幅图像并保存下来。
def sample_imgs(ddpm,
net,
output_path,
n_sample=81,
device='cuda',
simple_var=True):
net = net.to(device)
net = net.eval()
with torch.no_grad():
shape = (n_sample, *get_img_shape()) # 1, 3, 28, 28
imgs = ddpm.sample_backward(shape,
net,
device=device,
simple_var=simple_var).detach().cpu()
imgs = (imgs + 1) / 2 * 255
imgs = imgs.clamp(0, 255)
imgs = einops.rearrange(imgs,
'(b1 b2) c h w -> (b1 h) (b2 w) c',
b1=int(n_sample**0.5))
imgs = imgs.numpy().astype(np.uint8)
cv2.imwrite(output_path, imgs)
一切顺利的话,我们可以得到一些不错的生成结果。下图是我得到的一些生成图片:
大部分生成的图片都对应一个阿拉伯数字,它们和训练集MNIST里的图片非常接近。这算是一个不错的生成结果。
如果神经网络的拟合能力较弱,生成结果就会差很多。下图是我训练一个简单的ResNet后得到的采样结果:
可以看出,每幅图片都很乱,基本对应不上一个数字。这就是一个较差的训练结果。
如果网络再差一点,可能会生成纯黑或者纯白的图片。这是因为网络的预测结果不准,在反向过程中,图像的均值不断偏移,偏移到远大于1或者远小于-1的值了。
总结一下,在复现DDPM时,最主要是要学习DDPM论文的两个算法,即训练算法和采样算法。两个算法很简单,可以轻松地把它们翻译成代码。而为了成功完成复现,还需要花一点心思在编写U-Net上,尤其是注意处理时间戳的部分。
#基于扩散模型的实时规划框架DiffuserLite
在人工智能领域,扩散模型(Diffusion Models,简称DMs)凭借其强大的生成能力,在图像、文本以及轨迹规划等任务中展现了巨大的潜力。然而,当扩散模型应用于决策规划领域时,其低效的迭代采样成为了一个主要瓶颈。许多现有方法,如 Diffuser 和 Decision Diffuser (DD),在生成高质量轨迹的同时,由于复杂的建模过程,其决策频率往往低于 1Hz。这种性能远不能满足机器人控制、游戏 AI 等实际场景对实时响应的要求。
为解决上述问题,DiffuserLite 应运而生。它通过引入一种创新性的 渐进式精细规划(Progressive Refinement Planning, PRP) 方法,在减少冗余信息建模的同时,实现了更高效的轨迹生成。令人印象深刻的是,DiffuserLite 的决策频率达到了 122Hz,比主流框架快 112 倍以上,并在多个基准任务中实现了 State-of-the-Art(SOTA) 性能。本文将深入讲解 DiffuserLite 的核心原理与技术亮点,剖析其在实际应用中的潜力,并探讨这一框架可能带来的广泛影响。
问题
在决策规划领域,DiffuserLite 所面临的任务可以描述为一个基于离线强化学习(Offline Reinforcement Learning)的优化问题,其核心目标是生成一条满足目标属性的轨迹。
系统的状态由离散时间动力学方程控制,公式如下:
其中:
- 表示系统在时间 t 的状态。
- 为在时间 t 的动作。
- f(.)表示系统的动态模型。
Diffusion Planning 的目标是找到一条最符合目标的轨迹:
其中:
- d 是一个度量轨迹属性与目标属性之间距离的函数。
- C是一个评价器(Critic)。
在离线强化学习的背景下,轨迹的属性通常被定义为其累积奖励:
核心在于从离线数据中生成符合目标属性的轨迹,并从中提取执行动作
扩散模型
扩散模型被引入来解决轨迹生成问题。其通过对轨迹分布的建模,能够生成满足目标条件的长时序轨迹。与传统的逐步生成方法不同,扩散规划直接生成完整轨迹,避免了步进式方法中常见的累积误差。现有扩散规划方法的一个主要瓶颈在于高复杂度的采样过程,这需要多次前向传播和复杂的去噪处理,导致决策频率极低。
DiffuserLite
DiffuserLite 是一种高效轻量级的扩散规划框架,旨在解决现有扩散规划方法中因冗余信息建模而导致的低效率问题。它通过引入 渐进式精细规划(Progressive Refinement Planning, PRP) 方法,减少了不必要的计算。此外,DiffuserLite 使用 DiT(Transformer 变体) 替代传统的 UNet 模型。DiffuserLite 的规划过程分为以下几个阶段:
- 初始粗略规划:在较大的时间间隔内多次生成关键状态点,忽略中间冗余细节。
- 评价:使用评价器(Critic)选择最优轨迹。
- 逐步精细化:在每一层细化关键状态点之间的轨迹(执行多次),利用评价器选择最有轨迹。
- 快速执行:通过简化的模型架构和条件采样技术,快速生成动作决策。
渐进式精细规划
渐进式精细规划(PRP) 是 DiffuserLite 的核心创新之一,用于在生成轨迹时逐步减少冗余建模,提升效率和准确性。PRP 的核心理念是将轨迹生成过程分解为多个逐层优化的阶段,从最初的粗略规划逐步细化,直到生成精确的完整轨迹。在扩散规划中,完整轨迹的生成通常涉及大量冗余信息,尤其是远端状态的细节对于当前决策的影响微乎其微。例如,在长时间跨度的轨迹规划中,远端的状态可能由于环境噪声或动态变化而无法准确到达,因此建模这些细节不仅无益,反而增加了计算负担。冗余信息导致模型需要处理更高维度、更复杂的概率分布,这显著降低了推理速度。PRP
- 粗略规划:
初始阶段仅生成关键点(Key Points),间隔较大的状态点被视为轨迹的主要参考,而中间状态被忽略。 - 逐层精细化:
从关键点轨迹出发,逐层填补两点之间的细节,逐步将轨迹分辨率提高到所需的精度。
示例: - 优化终结:最后一层完成全轨迹的生成,生成的轨迹既包括整体的长远性,又有足够的局部细节。
例如:对于规划范围为 128 的轨迹,第一层只生成 [0,32,64,96,128] 这几个点。第二层生成[0, 8, 16, 24, 32],第三层生成[0,1,2,3,4,5,6,7,8]。相比于传统的one-shot方法,能够快速生成出理想的轨迹。
训练
将完整轨迹按层次划分为子轨迹,每层生成的轨迹采用扩散模型拟合其概率分布,优化目标为最小化噪声预测误差。为指导模型生成目标轨迹,DiffuserLite 使用了无分类器指导(CFG)。
Critic设计
Critic 是 DiffuserLite 中的关键部分。Critic的两个核心作用:
- 提供生成条件:在扩散模型训练过程中,Critic 评估轨迹的属性 C(x),并作为条件输入指导模型生成目标轨迹。
- 选择最优轨迹:在推理阶段,Critic 用于从候选轨迹中选择最优轨迹。
在稀疏奖励任务中,直接使用累积奖励可能不够有效,因此可以加入最终状态的值函数 :
在特定任务(如机器人控制或行为定制)中,可以将其设计为目标任务的特定指标,例如轨迹平滑性、目标对齐度、或者状态变化的能量消耗。
实验
1. 实验目的
实验的主要目的是验证 DiffuserLite 在多种任务中的性能,包括:
- 效率:决策频率是否显著提升。
- 准确性:生成轨迹是否符合任务目标。
- 任务适应性:是否能够处理长时序、稀疏奖励等复杂任务。
2. 实验设置
数据集与环境
DiffuserLite 在以下标准基准任务中进行评估:
- Gym-MuJoCo:机器人运动控制任务。
- 任务包括 HalfCheetah、Walker2D 等,主要关注连续控制性能。
- FrankaKitchen:复杂任务完成环境。
- 包含多个子任务组合,如打开微波炉、关门等,考验轨迹规划能力。
- Antmaze:长时序导航任务。
- 需要在稀疏奖励场景中生成可行的轨迹,特别测试远端轨迹规划能力。
评估指标
- 决策频率(Hz):每秒生成的决策次数。
- 任务成功率:轨迹是否成功达到目标状态。
- 奖励累计值:生成轨迹的累计奖励。
3. 实验结果
性能比较
DiffuserLite 与主流方法(如 Diffuser、Decision Diffuser)进行了全面对比。结果如下:
- 决策频率:
- DiffuserLite 的频率达到了 122 Hz,显著高于 Diffuser(约 1 Hz)和 Decision Diffuser(约 0.8 Hz)。
- 任务成功率:
- 在 FrankaKitchen 任务中,DiffuserLite 的成功率达到了 85%,比基准方法高出约 10%。
- 累积奖励:
- 在 Gym-MuJoCo 任务中,DiffuserLite 的累计奖励接近理论最优值,展现了优异的控制性能。
消融实验
为了验证各组件的贡献,进行了以下消融实验:
- 无 PRP(渐进式精细规划):
- 不使用 PRP,直接生成完整轨迹。结果显示效率和准确性显著下降。
替换模型架构:
- 用传统的 UNet 替换 DiT(Diffusion Transformer),推理速度降低约 30%。
移除 Critic 指导:
- 不使用 Critic 提供的轨迹属性条件,任务成功率降低了 15%。
总结
DiffuserLite 在扩散规划领域实现了一项重要突破,通过引入 渐进式精细规划(PRP) 和灵活的 Critic 和属性设计,显著提升了轨迹生成的效率和适应性。实验结果表明:
- 高效性:DiffuserLite 的决策频率达到了 122Hz,是现有方法的百倍以上,能够满足实时任务需求。
- 准确性:在 Gym-MuJoCo、FrankaKitchen、Antmaze 等复杂任务中,DiffuserLite 均展现了优异的任务成功率和轨迹规划能力。
- 适应性:通过 PRP 和 Critic 的结合,DiffuserLite 能够有效处理长时序和稀疏奖励任务,展现了极强的任务扩展性。
整体来看,DiffuserLite 解决了扩散规划的实时性问题。
#AdaVD
零训练成本!中科大创新扩散模型概念擦除方法,先验保护较SOTA提升 10 倍
当前,AI 图像生成技术迅猛发展,各类图文生成模型让用户能凭借简单文字描述创作出精美的图像。然而,这也引发了诸多问题,比如有人借此剽窃艺术风格、丑化 IP 角色和名人,甚至生成不安全图像。
如何以低成本且高效的方式,精准快速地从生成图像中去除这些不想要的概念,成为亟待解决的难题。
这一概念擦除任务有两个关键要求:一是擦除效果,需在生成过程中彻底清除与目标概念相关的语义;二是先验保护,即擦除目标概念时尽量不影响无关概念的生成。但现有方法难以平衡二者。
为此,中国科学技术大学联合曼彻斯特大学等机构的研究团队,提出了一种名为自适应值分解器(AdaVD)的 training-free 方法。在无需额外训练的前提下,它实现了对目标概念的精准擦除,同时最大程度地保护了先验知识,相较于现有 SOTA 方法,先验保护能力提升 2 到 10 倍。
该成果已被CVPR2025接收。
- 论文标题:Precise, Fast, and Low-cost Concept Erasure in Value Space: Orthogonal Complement Matters
- 论文链接:https://arxiv.org/pdf/2412.06143
- 代码链接:https://github.com/WYuan1001/AdaVD
现有方法的困境
现有概念擦除方法主要分为基于训练(training-based)和无训练(training-free)两类。基于训练的方法需要对模型参数进行微调,尽管擦除效果较好,但成本高昂,且处理速度较慢,难以满足在线 T2I 平台的实时需求。
此外,这类方法往往难以在擦除目标概念的同时保护非目标概念。而无训练方法,如 NP、SLD 和 SuppressEOT,虽然能够实现快速擦除,但各自存在不足。
例如,NP 的擦除精度有限,SuppressEOT 需要用户手动指定目标概念的位置,而 SLD 在保护非目标概念的先验信息方面表现不佳,影响图像的整体质量。
AdaVD 如何破局
核心原理:正交补运算
作为一种无需训练的概念擦除方法,AdaVD 通过经典的正交补运算,在扩散模型 UNet 的值空间内进行精准擦除。
具体而言,它将原始文本提示所对应的值投影到目标概念的正交补空间,从而剥离目标概念的语义,同时最大程度地保留其他内容。此外,AdaVD 采用逐 token 计算方式,以确保对每个 token 的擦除精度,实现更灵活细致的概念擦除。
创新机制:自适应擦除移位
为了在保证擦除效果的同时进一步增强先验知识保护,AdaVD 还引入了自适应擦除移位机制(Adaptive Token Shift)。该机制基于文本 token 与目标概念 token 之间的语义相关性计算移位因子,从而动态调整擦除强度。
若某个 token 与目标概念的语义关联较低,AdaVD 会减少对其的擦除,从而在精准擦除目标概念的同时,最大程度地保留先验知识。
效果惊艳,实力验证
实验结果显示,AdaVD 在多个概念擦除任务中均表现出色,涵盖 IP 角色、艺术风格、NSFW 内容及名人相关概念等场景。与其他方法相比,AdaVD 在 CLIP Score(CS)和 FID 评分上均表现优异,前者反映擦除效果,后者衡量先验保护能力。
此外,在多概念擦除实验中,AdaVD 也展现出卓越的性能。例如,在同时擦除“Snoopy”“Mickey” 和 “Spongebob” 时,其 CS 和 FID 均优于其他方法,即便同时擦除 40 个概念,也能有效保护非目标概念,展现出强大稳定的擦除和先验保护能力。
效率与可解释性兼具
在时间消耗层面,AdaVD 由于无需模型微调,在运行速度上远超基于训练的方法。虽然比 SLD 略慢,但换来的却是更优的擦除性能,整体性价比更高。
此外,研究团队通过可视化擦除部分,深入探究了 AdaVD 的可解释性。结果显示,AdaVD 能够精准定位语义空间中与目标概念相关的语义信息。
对于无关概念,其对应的擦除部分呈现出无意义的信息特征。这一现象进一步证实,AdaVD 不仅可以精准擦除目标概念,还能实现擦除效果与先验保护的良好平衡。
多场景应用拓展
AdaVD 的能力不仅限于概念擦除,还可拓展至多个图像生成任务。例如,在隐式概念擦除中,它能去除“rainy”“foggy”等隐含概念;在图像编辑任务中,可精准移除“glasses”“mustache”等外观特征;在属性抑制任务中,可去除“red”等耦合颜色概念,如让苹果或玫瑰褪色。
此外,AdaVD 还能与多种扩散模型兼容,如 Chilloutmix、DreamShaper、RealisticVision 和 SD v2.1,适用性极为广泛。
#AttentionDistillation
扩散模型炼出新绝技!注意力蒸馏技术:图像生成效果全面碾压传统方法
本文介绍了一种名为“注意力蒸馏”的新技术,通过新颖的注意力蒸馏损失函数和改进的分类器引导采样方法,显著提升了扩散模型在视觉特征转移任务中的性能,包括风格迁移、纹理合成等,生成效果全面优于传统方法。
论文链接:https://arxiv.org/pdf/2502.20235
项目链接:https://xugao97.github.io/AttentionDistillation/
亮点直击
- 分析了之前即插即用注意力特征方法的局限性,并提出了一种新颖的注意力蒸馏损失,用于重现参考图像的视觉特征,取得了显著优越的结果。
- 开发了attention distillation引导采样,这是一种改进的分类器引导方法,将注意力蒸馏损失整合到去噪过程中,大大加快了合成速度,并支持广泛的视觉特征转移和合成应用。
生成效果一览
给定参考图,文生图
艺术风格迁移
纹理合成
总结速览
解决的问题
- 现有生成扩散模型在图像风格和语义理解方面虽然有进展,但在将参考图像的视觉特征转移到生成图像中时,使用即插即用注意力特征的方法存在局限性。
提出的方案
- 提出了一种新颖的注意力蒸馏损失,用于在理想和当前风格化结果之间计算损失,并在隐空间中通过反向传播优化合成图像。
- 开发了一种改进的分类器引导方法,即注意力蒸馏引导采样,将注意力蒸馏损失整合到去噪采样过程中。
应用的技术
- 利用预训练扩散网络的自注意力特征。
- 通过注意力蒸馏损失优化生成图像。
- 将注意力蒸馏损失整合到去噪采样过程中,以改进分类器引导方法。
达到的效果
- 大大加快了合成速度。
- 支持广泛的视觉特征转移和合成应用。
- 在将示例的风格、外观和纹理转移到新图像的合成中表现出色。
方法
预备知识
隐空间扩散模型(LDM),如Stable Diffusion,由于其对复杂数据分布的强大建模能力,在图像生成方面达到了最先进的性能。在LDM中,首先使用预训练的VAE 将图像 压缩到一个学习到的隐空间中。随后,基于UNet的去噪网络 被训练用于在扩散过程中预测噪声,通过最小化预测噪声与实际添加噪声 之间的均方误差来实现。
其中 表示条件, 表示时间步长。去噪 UNet 通常由一系列卷积块和自注意力/交叉注意力模块组成,所有这些都集成在残差架构的预测分支中。
KV注入在图像编辑 [2, 7, 58]、风格迁移 [10, 25, 29] 和纹理合成 [70] 中被广泛使用。它建立在自注意力机制之上,并将扩散模型中的自注意力特征用作即插即用的属性。自注意力机制的公式为:
在注意力机制的核心,是基于查询 和键( )之间的相似性计算权重矩阵,该矩阵用于对值( )进行加权聚合。KV注入通过在不同的合成分支之间复制或共享KV特征来扩展这一机制。其关键假设是 特征代表图像的视觉外观。在采样过程中,将合成分支中的 特征替换为示例的相应时间步长的 特征,可以实现从源图像到合成目标的外观转移。
注意力蒸馏损失
尽管KV注入取得了显著的效果,但由于残差机制的影响,它在保留参考的风格或纹理细节方面表现不足;例如,参见下图3(a)。KV注入仅作用于残差,这意味着信息流(红色箭头)随后受到恒等连接的影响,导致信息传递不完整。因此,采样输出无法完全再现所需的视觉细节。
本项工作提出了一种新的损失函数,通过在自注意力机制中重新聚合特征来提取视觉元素,因此称之为注意力蒸馏(AD)损失。利用预训练的T2I扩散模型Stable Diffusion的UNet,从自注意力模块中提取图像特征。如下图2(a)所示,首先根据目标分支的 ,从参考分支重新聚合KV特征( 和 )的视觉信息,这与KV注入相同。
将此注意力输出视为理想的风格化。然后,我们计算目标分支的注意力输出,并计算相对于理想注意力输出的L1损失,这定义了AD损失:
可以使用提出的AD损失通过梯度下降来优化随机隐空间噪声,从而在输出中实现生动的纹理或风格再现;例如,参见上图3(b)。这归因于优化中的反向传播,它不仅允许信息在(残差)自注意力模块中流动,还通过恒等连接流动。通过持续优化, 和 之间的差距逐渐缩小,使得注意力越来越准确,最终特征被正确聚合以产生所需的视觉细节。
根据最近的实验分析 [7, 29, 58],经验性地选择UNet的最后6个自注意力层来计算AD损失。此外,在优化过程中,通过线性减少输入到UNet的时间步长 从 到 0来模拟扩散模型的采样过程。从不同的随机隐空间噪声开始,并在100步内优化它们。请注意,在整个优化过程中,UNet预测的噪声完全被丢弃,不断更新相同的隐空间变量。
为了更好地理解我们的AD损失,展示了多次运行的优化结果,如下图4所示。这些结果表明:i)AD损失有效地蒸馏出高质量的风格和纹理视觉特征;ii)AD损失自适应于不同的空间结构,展示了多次运行中的多样性。
内容保留优化
通过AD损失蒸馏的纹理和风格,可以进一步使用内容损失将合成内容与另一个参考图像对齐。这种优化允许图像合成在保留目标内容的同时转换一个图像的视觉元素,实现风格迁移、外观迁移等任务。
如上图2(b)所示,定义内容损失与AD损失类似,也基于自注意力机制,充分利用扩散模型中对图像的深刻理解。特别是,计算目标查询 和参考查询 之间的L1损失,构成了内容损失:
在实现中,同样选择最后6个自注意力层来计算内容损失,这与AD损失保持一致。内容保留优化的目标是:
优化后,经过预训练的VAE将优化后的隐空间code解码为图像空间。
注意力蒸馏引导采样
上述合成是使用反向传播优化的。在本节中,介绍如何将注意力蒸馏损失以改进的分类器引导方式纳入扩散模型的采样过程中。
根据[13],分类器引导在去噪过程中改变去噪方向,从而生成来自的样本,其公式可以表示为:
其中, 是时间步长, 表示提示, 和 分别指去噪网络和LDM中的隐空间变量。 控制引导强度。受 [17]启发,使用基于注意力蒸馏损失的能量函数来引导扩散采样过程。
具体来说,在DDIM采样[55]过程中,时间步长 的隐空间变量 根据从去噪网络估计的更新方向 被转换为 :
其中, 。将方程(6)中的 替换为 :
其中, 是与 相关的常数。通过用 替换为 来定义能量函数。然而,我们发现配置引导强度是一个棘手的问题;详见下文的验分析。为了解决这个问题,我们引入了Adam优化器[35]来自动管理强度并计算梯度。为简化起见,将 近似为 ,从而能够对 进行初始DDIM采样,然后进行简单的优化更新。AD损失以隐空间的 和噪声扰动的参考隐空间 作为输入,如上图2(c)所示。通过AD损失优化更新 为:
通过AD损失的引导,可以在时间步条件下计算隐空间的损失,而不是将隐空间转换为图像空间并计算图像级损失,如最近的工作所做的那样 [4, 41, 49]。使用Adam优化器能够建立一个通用的学习率,从而减轻设定引导强度的挑战。请注意,方程 (4) 中提出的内容损失也可以与AD损失一起添加到采样过程中,以进一步保留内容参考图像的结构。
改进的VAE解码
实验证据 表明,隐空间扩散模型中使用的VAE在感知上是有损的。对于需要高频局部细节的任务,例如纹理合成,可以选择在示例图像 上微调VAE解码器 的权重 ,以使用 L 1 损失来增强其重建质量,按照[3]的方法:
其中,表示VAE编码器。下图5展示了一些重建和采样的结果,显示了经过微调的VAE在感知质量上的提升。
实验
应用和对比
接下来,将注意力蒸馏损失应用于各种视觉特征转移任务,并将结果与每个应用中的最新方法进行比较。
风格和外观转移。 遵循Gatys等人 [22] 的著名工作精神,通过前文中描述的优化方法实现风格和外观转移。方法与CSGO 、StyleShot、StyleID、StyTR2和NST进行风格转移的比较,以及与Cross-Image Attention和SpliceViT进行外观转移的比较。下图6展示了定性比较结果。在风格转移中,本文方法有效地捕捉到高质量、一致的风格特征,同时保留了内容图像的语义结构。这在第3和第4行的素描风格中尤为明显。相比之下,尽管基线方法保留了原始结构,但在风格上存在显著差异。在外观转移中,本文的方法也表现出优越性,避免了Cross-image Attention中颜色过饱和的问题。
特定风格的文本到图像生成。 如前文所述,可以在扩散采样中应用我们的AD损失,从而实现特定风格的文本到图像生成。将参考图像设置为所需的风格图像。下图7展示了一些生成结果,并与其他方法进行比较,包括Visual Style Prompting、InstantStyle和RB-Modulation。下图7中的结果表明,本文的方法在文本语义上与现有方法相当,同时在风格一致性上与参考图像相比具有显著优势。除了上述方法之外,进一步结合ControlNet,以在各种模态(例如深度和Canny边缘)上进行额外条件的特定风格文本到图像生成。下图8展示了一些生成的例子。
受控纹理合成。 本文方法可以应用于纹理优化,如前文所示。受[7]的启发,在计算注意力蒸馏损失时进一步引入了mask引导,从而限制了由 查询的值,实现了受控的纹理合成。具体来说,给定一个源纹理 ,其对应的源分割图 ,以及目标分割图 ,首先将 和 展平,并将它们下采样以匹配注意力特征的分辨率,得到 和 。然后,计算引导 mask ;通过这个 mask,在注意力计算中限制了由 聚合的视觉信息,使其仅关注相应的区域:
为了进一步确保与目标mask的对齐,用从源图像对应标记区域随机抽取的像素填充目标mask作为初始化。然后,将此初始化视为内容参考,并基于查询特征添加内容损失,如前文所述。上图1和下图9展示了我们的受控纹理合成结果。
与基于补丁的神经纹理优化方法GCD相比,本文结果在保持可比纹理细节的同时展现了更平滑的物体边缘。相比之下,GCD存在颜色混叠的伪影问题;参见图9的第二行。最近,Self-Rectification引入了一种“懒编辑”控制,用于生成非平稳纹理。为了实现相同的目标,利用SDEdit来保留用户编辑的布局图像的结构。然后,将提出的AD损失和内容损失整合到采样过程中。如图9所示,Self-Rectification输出的纹理过渡更为平滑,而我们的结果更好地遵循了原始纹理示例。
纹理扩展。 由于补丁来源有限,使用传统方法合成超高分辨率纹理非常困难。在此,将注意力蒸馏引导的采样应用于MultiDiffusion模型,使纹理扩展到任意分辨率。尽管SD-1.5是在尺寸为512×512的图像上训练的,但令人惊讶的是,当结合注意力蒸馏时,它在大尺寸纹理合成中表现出了强大的能力。下图10展示了将纹理扩展到512×1536的尺寸与GCD和GPDM的比较。本文方法在如此具有挑战性的任务中显示出显著的优势。
消融研究
在本节中,展示了关于我们方法的两个方面的消融研究结果:i) 内容保持优化中内容损失权重的影响,以及ii) 在注意力蒸馏引导采样中管理引导强度的优化器。
内容损失权重。 如下图11所示,变化内容损失权重对转移结果有着有趣的影响。例如,在风格转移中,左侧展示的抽象风格示例说明了的调整如何导致不同程度的抽象化,为艺术创作提供了灵活性。在外观转移中,由于扩散网络的精确语义理解,右侧展示的人脸图像转移随着的变化表现出平滑的身份过渡。
优化器。 下图12展示了优化器在管理引导强度中的重要性。实验测试了一种手动设置引导强度的简单策略,以控制对隐空间变量的梯度更新规模。然而,手动变化这个强度通常无法产生合理的结果:示例中的纹理或外观特征通常会丢失,如下图12的最后三列所示。相比之下,引入Adam优化器来管理隐空间优化产生的结果与输入示例的视觉特征非常匹配(下图12的第2到4列)。此外,在每个采样时间步中增加优化迭代次数通常可以提高生成结果的质量,尽管这也带来了额外的计算时间。在实践中,将迭代次数设置为2,以实现高效平衡,提供高质量的结果。
用户偏好研究
为了验证定性分析,在三个转移任务上进行了包含30个问题的用户研究(每个任务针对6个选定的竞争对手各设置5个问题)。在每个问题中,展示两个结果:一个来自我们的方法,一个来自竞争对手。用户被要求根据提供的说明和标准选择更好的一个。从50名参与者那里收集了1500个回应,整体偏好得分总结在下图13中。我们的方法在很大程度上持续优于其他选项。
结论
本文提出了一种统一的方法来处理各种视觉特征转移任务,包括风格/外观转移、特定风格的图像生成和纹理合成。该方法的关键是一种新颖的注意力蒸馏损失,它计算理想风格化与当前风格化之间的差异,并逐步修改合成。本文方法克服了以往工作的局限性,实验验证了其优越性。
参考文献
[1] Attention Distillation: A Unified Approach to Visual Characteristics Transfer
#解决扩散模型中的优化难题
重建 vs 生成
本文提出了一种名为VA-VAE的方法,通过将视觉词元分析器的潜在空间与预训练的视觉基础模型对齐,解决了潜在扩散模型中重建与生成之间的优化难题。基于该方法构建的LightningDiT模型在ImageNet 256x256生成任务上取得了最佳性能,FID得分1.35,并在64个epoch内达到2.11的FID得分,显著提升了训练效率。
题目:Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models作者:Jingfeng Yao, Xinggang Wang
作者单位:华中科技大学Paper:https://arxiv.org/abs/2412.04852
Code:https://github.com/hustvl/LightningDiT
01 介绍
两阶段的潜在扩散模型中存在优化难题:在visual tokenizer中增加每个标记的特征维度,虽能提升重建质量,但要达到相近的生成性能,却需要大得多的扩散模型和更多训练迭代。因此,现有系统常常只能采用次优解决方案,要么因tokenizer中的信息丢失而产生视觉伪影,要么因计算成本高昂而无法完全收敛。作者认为这种困境源于学习不受约束的高维潜在空间的固有困难。
为了解决这一问题,作者建议在训练视觉词元分析器时,将潜在空间与预先训练的视觉基础模型对齐。提出的VA-VAE(视觉基础模型结合变分自动编码器)显著扩展了潜在扩散模型的重建生成边界,使高维潜在空间中的Diffusion Transformers(DiT) 能够更快地收敛。为了充分发挥VA-VAE的潜力,构建了一个增强型DiT基线,改进了训练策略和架构设计,称为LightningDiT。在ImageNet 256x256 生成上实现了最佳 (SOTA) 性能,FID得分为1.35,同时在短短64个epoch内就达到了2.11的FID得分,展现了卓越的训练效率——与原始DiT相比,收敛速度提高了21倍以上。
相关工作
- 可视化生成的tokenizer
visual tokenizer包括以变分自编码器(VAE)为代表的连续型和 VQVAE、VQGAN 等离散型。离散型词元分析器虽然能提高重建保真度,但编码对照本利用率低下,对生成性能产生不利影响。连续型tokenizer通过增加词元分析器的特征维度会提高重建质量,但会降低生成性能,还需要大幅增加训练成本,当前缺乏对连续型 VAE 优化的有效解决方案。
- 扩散Transformer的快速收敛
扩散Transformer(DiT)目前是潜在扩散模型最常用的实现方式,存在收敛速度慢的问题,往研究提出多种加速方法,本文则从优化视觉词元分析器学习的潜在空间入手,在不修改扩散模型的情况下实现更快收敛,并对 DiT 进行了训练策略和架构设计优化。
02 方法
- 网络架构
VA-VAE基于VQGAN模型架构,通过视觉基础模型对齐损失(VF损失)优化潜在空间。VF损失由边缘余弦相似度损失(Marginal Cosine Similarity Loss) 和边缘距离矩阵相似度损失(Marginal Distance Matrix Similarity Loss) 组成,是一个即插即用模块,在不改变模型架构和训练流程的情况下解决优化困境。
- 边际余弦相似度损失
将视觉标记器编码器输出的图像潜在特征投影后,与冻结的视觉基础模型输出的特征计算余弦相似度,通过 ReLU 函数和设置边际值,使相似度低于边际值的特征对损失有贡献,从而聚焦于对齐差异较大的特征对。
- 边际余弦相似度损失
将视觉标记器编码器输出的图像潜在特征投影后,与冻结的视觉基础模型输出的特征计算余弦相似度,通过ReLU函数和设置边际值,使相似度低于边际值的特征对损失有贡献,从而聚焦于对齐差异较大的特征对。
03 结果
visual tokenizer采用LDM的架构和训练策略,使用VQGAN网络结构和KL损失,训练三种不同的f16标记器(无VF损失,VF损失(MAE),VF损失(DINOv2))。生成模型采用LightningDiT,在ImageNet上以256分辨率训练,设置相关参数和训练策略。
使用 VF 损失在f16d32和f16d64收敛速度提升明显
8种不同tokenizer的重建和生成的评估
与现有扩散模型对比
可视化效果
#Nabla-GFlowNet
扩散模型奖励微调新突破: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 这一常用的文生图扩散模型上显示出相较于其他方法的优势。