改进的分布扩散模型

TL;DR iDDM 通过延迟多粒子训练与时间依赖评分规则，让分布扩散模型在 ImageNet-256 上单模型 4 步 FID 达 4.48，且采样步数增加不退化，无需蒸馏。

问题

背景

当前 diffusion 生成领域对少步采样 与单模型质量保持 的关注度很高, 主流方案依赖教师蒸馏或多阶段训练, 推理部署成本高。

现有方法局限

Distributional Diffusion Models (DDM) 用分布式去噪器替代均值预测, 理论上更适配 few-step 生成, 但原版 DDM 存在两个具体局限:

• 多粒子训练开销：为估计分布需同时采样 K 个 particle, 显存与计算随 K 线性增长, 难以扩展到 DiT 级 latent backbone。
• 固定 scoring rule 超参：全局固定的 scoring rule 在 4 步与 50 步采样预算下不能同时最优, 迫使单一 trade-off。

难点与重要性

Distributional 目标无闭式梯度, 依赖 score matching 或 implicit differentiation 近似; 高分辨率 latent 空间维度高、多模态强, 少量粒子易 collapse。业界对单模型、多 NFE 预算、无蒸馏的 few-step 生成器有实际部署需求, 因为避免教师模型和多阶段训练可显著降低工程与算力成本。

行业类比

类似 Speculative Decoding 在单模型内动态切换采样步数而不牺牲多样性, DDM 改进方向是让随机去噪器在不同采样预算下保持分布质量。

核心洞察

• 将多粒子训练推迟到晚期 Transformer 层，大幅降低 DDM 的计算开销，同时保留分布建模能力。传统 DDM 在所有层使用多粒子，导致训练成本随粒子数线性增长，难以扩展到高分辨率生成。该方法仅在深层进行粒子展开，利用浅层共享特征，从而在保持分布预测质量的同时，使训练成本可与单粒子模型相当。这一工程折衷是 DDM 首次在 **ImageNet-256** 规模上实用的关键。
• 引入时间依赖的 **scoring rule** 调度，替代全局固定超参数，使单一模型在不同采样预算下性能稳定。之前 DDM 使用固定 scoring rule 系数，导致需要在少步采样和多步采样之间取舍；通常少步模型在多步时 FID 会变差。该方法根据扩散过程的不同动力学阶段动态调整 scoring rule 权重（借鉴 Biroli2024 的 regime 分析），使得 4 步和 50 步采样共享同一组参数，且 FID 均保持较低，避免了常见的采样步数增加导致性能下降的问题。
• 整体训练方案无需教师模型、自蒸馏或 JVP 计算，单阶段即可获得有竞争力的少步生成器。许多少步扩散模型（如蒸馏方法）依赖教师模型或多阶段训练，增加了工程复杂性和计算成本。该方法从零开始单阶段训练，直接优化分布匹配目标，在 class-conditional ImageNet-256^2 上达到 4.48 FID@4步，50步 2.38 FID，且模型是随机生成器，采样多样性与质量兼顾。这种简洁的训练流程对实际部署和快速迭代有显著价值。

方法

输入与 DDMs 基础

DDMs 将标准去噪器从条件均值预测替换为 分布去噪器，学习 p(x_1 | x_t) 的随机近似，而非仅预测条件均值。训练目标采用 scoring rule（如 energy score），输入为带噪潜变量 x_t、时间步 t 与条件 c。

关键模块

1. 延迟粒子扩展 ：传统 DDM 为了估计分布，每个 Transformer 层都需采样多个粒子，开销随粒子数线性增长。本文仅在后期 Transformer 层 引入多粒子采样，前期保持单粒子计算，显著降低训练与推理开销。
2. 时间依赖 scoring rule 调度 ：原 DDM 全局固定 scoring rule 超参数，只能在少步与多步采样之间折衷。本文依据 Biroli2024 揭示的扩散动态区域（生成阶段 / 锐化阶段），对不同时间步 t 动态调整 scoring rule 参数，使单个模型适配不同采样预算。
3. 结合 DiT-XL/2 latent 架构，在 VAE 潜空间单阶段从零训练，无需教师、自蒸馏或 JVP 计算。

输出

模型作为随机少步生成器，在 ImageNet-256² 上 4 NFE 取得 FID 4.48，50 NFE 2.38，且 FID 不随步数增加而退化；同一配方可迁移至 text-to-image 生成。

与同类方法差异：通过在架构层延迟粒子计算、在训练目标层引入时间动态，首次让 DDM 在单模型、无蒸馏条件下实现少步至多步一致的生成质量。

实验

实验设计

在 class-conditional ImageNet-256 上，以 DiT-XL/2 为骨干，训练改进版 DDM (iDDM)。核心策略：将多粒子训练延迟到 late transformer layers，并采用 时间依赖的 scoring rule 调度 替代全局固定超参数。单阶段从零训练，不使用 teacher、自蒸馏或 JVP。

关键发现

• 4 步采样 FID 4.48，50 步 FID 2.38，FID 随 NFE 增加不退化。
• 延迟粒子扩展降低了多粒子训练的开销，使其在 DiT 规模上可行。
• 时间依赖调度平衡了不同采样预算下的质量，方法可迁移到文本到图像。

基线对比与工程启示

相比依赖蒸馏或一致性损失的少步模型，iDDM 无需 teacher 或蒸馏，单模型即达竞争性 FID；在 4→50 NFE 内单调改善，避免了高 NFE 下的性能饱和或反弹。工程上，该设计简化训练流程，为需要灵活采样预算的部署场景提供了可行替代。

行业影响

落地场景

Improved Distributional Diffusion Models (iDDM) 将扩散模型采样步数压缩至 4-50 步且 FID 稳定（4 步 4.48，50 步 2.38），适合低延迟生成场景：实时图像/视频生成、交互式设计、电商商品图合成、广告创意生成等。对比一致性模型或渐进蒸馏，iDDM 无需教师模型或蒸馏阶段，训练流程更简单。

商业价值

• 降低推理成本：4 步生成达到 4.48 FID，较常规 50 步以上扩散模型，单图算力需求降低一个数量级。
• 减少训练开销：单阶段从零训练，无需自蒸馏或 JVP，节省训练资源与工程复杂度。
• 体验提升：低延迟支持实时交互，如在线创意生成与预览。

与现有工作流集成

iDDM 基于 DiT-based latent diffusion，可在现有 pipeline 中替换去噪器，加入延迟粒子扩展 与 时间依赖 scoring rule 调度，不改变数据流或推理框架。模型输出分布而非点估计，适合多样性需求任务。代码与预训练模型已开源（GitHub），便于直接微调或蒸馏。

具体落地案例：

1. 电商商品图生成：商家上传单张商品图，实时生成多角度、多背景展示图，加速 listing 制作，降低拍摄成本。
2. 内容平台创作者工具：输入文字描述即时生成配图/封面，支持风格调整，提升创作效率。

局限

• 尽管推迟粒子扩展到 late transformer layers 降低了多粒子训练的计算开销，但 DDM 类方法本质上仍需要多个粒子来近似条件分布，相比标准扩散模型的单点预测，训练和采样时的内存与计算成本依然更高。论文未系统讨论粒子数对最终生成质量与训练稳定性的影响，例如粒子数过少可能导致分布近似偏差，过多则又回到开销问题，实际部署时需要在资源与效果之间反复权衡。
• 时间依赖的 scoring rule schedule 虽然基于 Biroli2024 的动力学区域设计，但引入了更多超参数：不同时间段的 scoring rule 类型与权重需要手工设定。论文只在 ImageNet-256 和文生图任务上验证了该 schedule 的有效性，缺乏在其他数据集、分辨率或不同扩散模型架构上的泛化证据。从业者迁移到新任务时可能需要重新调参，增加了工程适配成本。
• 实验聚焦于 class-conditional ImageNet-256 和初步的 text-to-image 验证，未涉及更高分辨率（如 512² 及以上）或更大规模数据集。与 GAN、Consistency Models 等少步生成方法的直接对比不足，特别是在样本多样性、模式覆盖和训练稳定性等维度。此外，DDM 的随机采样机制在某些需要确定性输出的场景（如图像编辑、反演）可能不如一致性模型直接，本文也未探讨此类下游任务的兼容性。