Improved Distributional Diffusion Models
Tommaso Martorella, Alexandre Galashov, Felix Krause, Stefan Andreas Baumann, Valentin De Bortoli, Arthur Gretton, Björn Ommer
cs.CV, cs.LG
2026-09-29
推迟粒子扩展并调度目标，DiT-XL在ImageNet上4步FID 4.48、50步2.38，单阶段从零训。
少步采样时，条件分布 p(x1|xt) 方差很大。标准 diffusion 和 flow matching 只拟合条件均值 E[x1|xt]，粗网格上这一刀切掉了后验形状，样本会糊、会塌。常见补救是换更好的 ODE 求解器、把多步模型蒸馏成少步，或用 consistency / flow-map 直接训少步生成器。
De Bortoli 等人 2025 年的 Distributional Diffusion Models（DDM）换了目标：去噪器额外吃一份噪声 ξ，学的是对 p(x1|xt) 的随机近似，训练用 generalized energy score。λ=0、β=2 时退回普通 MSE。它上不了现代图像规模，卡在两处。每个训练样本要跑 m 个粒子，前向成本随 m 线性涨；全局固定的 (λ, β) 把保真和多样性锁成一套，4 步和 50 步没法共用。
iDDM 动两处，backbone 是 latent DiT。
延迟粒子扩展。早期 Transformer 层对 xt 只跑一遍，到后几层才复制成 m 份并注入 ξ。粒子共享后段权重，差别只来自 ξ。DiT-XL 共 28 层，从第 24 层起扩展、m=4 时，单步训练成本大约是朴素 DDM 的 36%，相对 flow matching 约 1.5×；XL 在 4×A100 上逐步实测 1.41×。ξ 注入试过通道拼接、残差相加、AdaNorm、register token。固定内宽的通道拼接在 4 步 FID 上最好，AdaNorm 和 token 在少步上直接垮。开 CFG 之后分布层不宜太多：DiT-B 上 2 到 4 层最好，最终 B 规模 ℓstart=10（共 12 层）。
时间相关 scoring rule。后验在噪声端宽、数据端窄，固定超参两边都委屈。Biroli 等人 2024 年把反向过程分成 speciation、collapse，再加一个更靠数据端的几何分离阈值。早期偏分布目标，晚期偏回归。默认线性日程 s(t)=1-t、βmin=0.1，不必真去估那些阈值；按 regime 分段的 dyn-reg 日程效果接近。训练时 t 改成偏向噪声端的 jit 分布，分布损失在那儿更要紧。
推理每步抽一个 ξ，随机速度场走 Euler。训练是单阶段从零开始，不需要教师、自蒸馏或 JVP。
DiT-B、ImageNet-256、400k 步、同一套设定重训。评测 FID 用 50k 样本，CFG 按模型与步数网格搜索，训练不用 CFG。Shortcut 的 50 步栏实际是 64 步，因为它只支持 2 的幂。
| 方法 | 50步 FID | 4步 FID |
| Shortcut | 9.85 | 15.02 |
| IMM | 7.57 | 24.85 |
| MFM | 7.72 | 9.17 |
| 朴素 DDM | 10.04 | 43.68 |
| iDDM | 4.57 | 13.13 |
相对朴素 DDM，4 步从 43.68 降到 13.13，50 步从 10.04 降到 4.57。训练算力对齐时（iDDM 260k 步对 FM 400k 步），4 步 FID 是 14.33 对 26.70，50 步仍略差（5.62 对 4.97）。4 步上 FID、FD-DINOv2、KID、Inception Score、precision、recall 同时好于两边基线。
DiT-XL/2、200 epoch：
| 方法 | 训练量 | 50步 FID | 4步 FID |
| MeanFlow-XL/2 | 240 ep | 3.29 | 2.93 |
| iMF-XL/2（48层） | 800 ep | 1.43 | 1.51 |
| IMM-XL/2 | 3840 ep | 1.82 | 2.51 |
| MFM（多阶段） | 800+80+28 | 5.74 | 1.97 |
| 朴素 DDM-XL/2 | 200 ep | 4.71 | 23.25 |
| iDDM-XL/2 | 200 ep | 2.38 | 4.48 |
MeanFlow 从 4 步到 50 步 FID 从 2.93 升到 3.29，MFM 从 1.97 升到 5.74。iDDM 从 4.48 降到 2.38，同一份权重覆盖 4 到 50 步。逐步训练成本 1.41× FM，MeanFlow 1.96×，iMF 6.20×；总训练量 MeanFlow 大约 1.67×，iMF 大约 17.6×。
同一套做法迁到 1.6B 文生图。MS-COCO 4 步 FID 41.05，匹配 FM 基线是 78.20；8 步 19.40 对 26.33；16 步几乎打平（11.27 对 11.20）。4 步 CLIP、GenEval、HPS v2.1 掉得也更少。16 步时 T2I-CompBench 的 Spatial 仍弱于 FM（0.12 对 0.29）。
如果目标是从零训、单阶段、4 到 50 步共用一份权重的随机少步生成器，Table 4 里只有这套配方同时满足。蒸馏（sCD、ADD）在 1 到 2 步更强，但要预训练教师和第二阶段。iMF 和 IMM 数字更好，训练贵一个数量级。MeanFlow 和 MFM 在 4 到 8 步 FID 更低，加步会变差。
每步大约 1.4× FM，不是免费。代码和预训练权重已公开。
论文自己写了：regime 锚点来自平均场分析，有限规模只是近似；平滑日程比硬切更稳；1 到 2 步上蒸馏和 fastforward 仍更强。最优 (λ(t), β(t)) 仍开放。B 规模消融里，只调度 β、λ=1、在 t=0.1 硬切，4 步 FID 到 13.06，比默认线性日程还好，因 t>0.1 后方差不再被约束而未采用。日程选择仍偏启发式。
文生图只把 256² 预训练模型 finetune 了 20k 步到 512²。Spatial 等组合指标在多步时弱于 FM。和 iMF 的对照不是同深度、同训练量。延迟扩展在无 CFG 时层越多越好，开 CFG 后反过来，这个交互还没有干净解释。