Optimizing DDPM Sampling with Shortcut Fine-Tuning
Ying Fan, Kangwook Lee
ICML 2023
cs.LG
2023-01-31
把扩散生成当成强化学习微调采样器,不再模仿反向过程;CIFAR-10 上 10 步采样 FID 2.28,优于 1000 步原模型的 3.03。
扩散模型(DDPM)生成一张图要跑上千步去噪,这是它落地时最大的拖累。压缩步数的办法几乎都遵循同一个思路:把完整的反向去噪过程尽量忠实地用更少步数复刻出来,手段包括更准的噪声估计、更聪明的子采样时间表,或换用非高斯噪声。无论怎么调,都是在复刻一个固定的反向过程。
作者换了角度:把从噪声生成图像看成控制问题,纯噪声是起点、自然图像是终点。反向过程更像一组示范轨迹(imitation learning),示范有用但很少是最优的。既然如此,何必死守这条轨迹?能不能让采样器自己找一条更短的捷径。问题于是变成:直接把生成分布拉到接近真实数据分布,而不是去拟合反向过程的每一步。
衡量的尺子换成积分概率度量(IPM,Integral Probability Metric)。IPM 是一族衡量两个分布差异的指标,Wasserstein 距离、总变差、MMD 都属于这一族。目标直接写成最小化生成分布 pθ0 与数据分布 q0 之间的 IPM,不再关心每一步是否忠于反向过程。
梯度怎么算,文章给了两条路。第一条 SFT 走 GAN 式的可微路径,把整条 T 步采样链一起反向传播,问题类似 RNN:梯度消失、梯度爆炸、显存爆。第二条 SFT-PG 更关键,作者证明(定理 4.1)对 IPM 做梯度下降等价于做策略梯度(policy gradient,强化学习里直接更新策略参数、拿轨迹回报当信号的那一类方法)。做法是把扩散采样构造成一个有限时域的 MDP:每一步的条件分布 pθt(xt|xt+1) 当作策略,状态转移是恒等映射,只在最后一步给奖励,奖励值就是判别器(critic)对最终图像的打分。这样一来,梯度只用到判别器的值而非它的梯度,可选的判别器范围更宽,也不会撞上长链的梯度消失和爆炸,代价是策略梯度方差更高,需要用 baseline 函数 V 来降方差,和 REINFORCE、GAE 的套路一致。
两个理论支撑让它能稳着训。定理 4.2 给出 IPM 的一个代理函数,保证在信任域里多走几步生成器更新仍能单调改善,精神类似 TRPO;落地是小学习率配梯度范数裁剪。判别器的正则上,作者发现可以复用 baseline 的损失去约束判别器的值(baseline regularization,B),比 WGAN 那套梯度惩罚(GP)更松、对生成器变化更敏感,实测效果更好。
toy 数据集 swiss roll 上,先把 DDPM 训到 10 步再微调,微调后的 10 步采样器连 1000 步原模型都赢了:
| 模型 | W2(p0,q0) ×10-2(越低越好) |
| DDPM,10 步 | 8.29 |
| DDPM,100 步 | 2.36 |
| DDPM,1000 步 | 1.78 |
| SFT-PG(B),10 步 | 0.64 |
真正有分量的是图像基准。在 CIFAR-10(32×32)和 CelebA(64×64)上,从预训练的 1000 步 DDPM 出发,固定 FastDPM 的 10 步采样表、只微调均值网络:
| 方法(均 10 步) | CIFAR-10 FID | CelebA FID |
| 原始 DDPM 子采样 | 34.76 | 36.69 |
| FastDPM | 29.43 | 28.98 |
| Analytic-DPM | 22.94 | 28.99 |
| SN-DDPM | 16.33 | 20.60 |
| SFT-PG(B) | 2.28 | 2.01 |
| 1000 步原模型(参照) | 3.03 | 3.26 |
只用了 10 步,微调后的 FID 不仅把当时最强的快速采样器 SN-DDPM 远远甩开,还低于 1000 步全步模型。和确定性的 DDIM 系比,SFT-PG 在 NFE=10 时 FID 2.28 是表里最好的;NFE=8 时 2.64,与 progressive distillation 的 2.57 相当。算力上,progressive distillation 在 CIFAR-10 上要 8 块 TPUv4 跑约一天,SFT-PG 只要 4 块 RTX 2080Ti 跑约 6 小时。
这是已知最早把强化学习用来训练扩散模型的工作。放到 2023 年的语境里,真正站得住的贡献是那个视角转换:别再把反向过程当成必须复刻的标尺,它只是一条未必最优的示范轨迹,直接对结果分布做优化,采样器就能自己长出捷径。这个想法后来在「扩散+RL」的研究线上被反复重新发现,也是它最近被翻出来讨论的原因。
对从业者,务实的一面是 CIFAR-10 这种小尺度上,10 步就能换回超过 1000 步的质量,算力还比蒸馏便宜。但要清醒,这是 32×32 的小图、2023 年的模型体量,没有在大型文生图上验证过,直接搬去 Stable Diffusion 级别的管线还为时过早。
作者自己承认,微调只动了均值网络 μθ,方差固定,学方差留作未来工作。哪怕采样时不再跟踪梯度,收集训练序列仍要跑完整 T' 步推理,所以推理速度还是比 GAN 这种一步出图的慢。
原理上也有几处没压实。策略梯度本身方差大,baseline 能压,但对 baseline 网络的拟合质量敏感;初始化时生成分布若离数据太远,容易陷到次优;判别器是作者手搭的小 CNN,信号好坏全看它,今天主流用预训练特征算 MMD、有现成的好判别器,这条路当时还没走通。通用性也没被检验:所有结果只到 CIFAR-10/CelebA 这一档,没上过更大的模型。