扩散模型采样误差有了专门修正器:ImageNet上FID从9.58降到2.07,无引导生成也能接近CFG
核心概要
扩散模型在采样时会累积自身预测误差,一个以基础模型预测为输入、用同一干净目标训练的修正器,在ImageNet类条件生成上把DiT-XL/2的FID从9.58降到2.07,无需引导即可接近CFG水平。
深度剖析
扩散模型采样时,去噪器每一步的输入来自它自己上一步的不完美预测,而不是训练时见过的加噪真实图像,这种训练与推理之间的错位会让误差沿去噪轨迹累积,产生伪影和结构不连贯。 此前改善样本质量主要靠分类器无关引导(CFG),它提升质量但会降低多样性并造成过饱和;另一条路是增加采样步数或使用更精细的数值积分,但积分误差变小后收益迅速趋于平台。 在ImageNet类条件生成上,修正采样在匹配函数评估次数下把FID从DiT-XL/2的9.58降到2.07、SiT-XL/2的8.55降到1.92、SiTREPA-XL/2的5.53降到1.31,而CFG参考值分别为2.51、2.11和1.35。
修正器把去噪器的干净样本预测当作输入,去预测同一个干净目标,因此修正成为一个显式的学习任务,训练输入由基础模型在线提供。 标准去噪训练只学习从加噪观测到干净图像的映射,修正器则在这个初始预测之上继续重建,专注于剩余的不准确部分。 在DiT、SiT、JiT等架构上,从公开检查点初始化去噪器和修正器并联合微调,SiTREPA仅需5万次更新,约为其400万次预训练的1.25%。
训练只监督一次修正,但推理时可以把修正器反复作用在自己的输出上,在固定噪声水平下迭代细化预测,从而在采样更新之前就抑制早期伪影。 此前增加推理算力通常用于增加去噪步数,而修正把同样的算力花在改进每一步的预测上。 在匹配总推理成本时,迭代修正比单纯增加去噪步数更有效;在文本到图像生成上,无CFG时GenEval从28.30升到50.09、DPG-Bench从70.71升到76.82,函数评估次数约为245。
修正带来的额外推理开销可以通过蒸馏消除:让学生模型匹配教师经过多次修正后的预测,学生只需每个采样步一次前向传播。 权重共享可以让去噪器和修正器共用主干、只保留独立的轻量输入嵌入,把参数量从1.35B降到0.68B,但每次修正迭代仍需额外前向传播。 在SiT-XL/2的ImageNet-256类条件生成上,学生模型FID为2.04,接近教师修正后的1.92,同时把函数评估次数减少约67%。
启示与展望
报告覆盖的是ImageNet类条件与无条件生成、文本到图像生成,以及DiT、SiT、SiTREPA、JiT等架构上的实验;修正器从公开检查点初始化并联合微调,训练更新量相对预训练很小。读者若关心无引导生成或CFG的多样性代价,可先看类条件与无条件ImageNet的FID、精度、召回对照,再看文本到图像的GenEval与DPG-Bench。
修正器只在基础去噪器的预测上训练,却能反复作用于自身输出,这一外推行为目前只有经验解释,理论理解仍是开放问题。修正会提高去噪误差,同时缩小训练与推理分布之间的差距,两者如何共同影响样本质量尚未完全厘清。修正有时会产生过于规则的结构,例如圆盘状披萨配料或规则网格排列的草莓籽。冻结去噪器只训练修正器时,更多修正迭代后FID会回升,联合训练则继续下降。
