论文题目FasterDiT: Towards Faster Diffusion Transformers Training without Architecture Modification无需修改架构即可实现更快的扩散变形金刚培训会议NIPS2024摘要扩散变压器(Difference TransformersDIT)的研究引起了人们的极大关注。然而它们面临着收敛速度缓慢的问题。在本文中我们的目标是在不进行任何架构修改的情况下加速DIT培训。我们在培训过程中发现了以下问题首先某些培训策略在不同的数据上并不总是表现良好。其次在特定的时间步骤上监督的有效性有限。作为回应我们提出了以下贡献(1)我们引入了一个新的视角来解释这些策略的失败。具体地说我们略微扩展了信噪比(SNR)的定义并建议观察SNR的概率密度函数(PDF)来理解该策略数据稳健性的本质。(2)进行了大量的实验报告了上百个实验结果从概率密度函数的角度经验总结出了一个统一的加速策略。(3)开发了一种新的督导方法进一步加快了DIT的培训进程。在此基础上提出了一种非常简单实用的设计策略--FasterDiT。在只需少量代码修改的情况下在ImageNet上以256×256的分辨率和1000次迭代获得2.30FID与DIT(2.27FID)相当但训练速度快7倍。FasterDiT无需修改架构让扩散 Transformer 训练快 7 倍一、背景与动机自 OpenAI 发布 Sora 以来其基础模型——扩散 TransformerDiffusion TransformerDiT——引发了学界的广泛关注。DiT 凭借出色的灵活性与可扩展性在图像生成和视频生成领域均展现了卓越的能力。然而DiT 存在一个显著的瓶颈收敛速度极慢。在 ImageNet 256 分辨率的训练任务上DiT 需要消耗超过 4700 GPU 小时H800才能收敛。这一巨大的计算开销严重阻碍了大规模训练实验的推进。为了改善训练效率研究者们提出了多种方法包括调整噪声调度noise scheduling修改损失权重loss weighting改变时间步采样策略timestep sampling这些方法的本质是调控训练过程中不同时间步对应的信噪比Signal-to-Noise RatioSNR分布。然而这些方法并不总是可靠的。[此处配图Figure 1 左图]以 SD3 提出的 lognorm 采样策略为例lognorm(−0.5, 1) 在 ImageNet 上优于 lognorm(0.5, 0.6)但随着数据信号强度持续降低lognorm(−0.5, 1) 的 FID 反而逐渐变差最终成为表现最差的方案。同一策略不同数据截然相反的结论。这一现象揭示了当前方法的核心问题缺乏跨数据条件的鲁棒性研究者往往只针对特定数据找到一个局部最优解而没有从本质上理解策略有效性的根源。本文的目标正是在不修改任何模型架构的前提下找到一种简单、可解释、普适的 DiT 训练加速策略。二、重新审视信噪比引入 PDF 视角2.1 传统 SNR 定义的局限在扩散模型和流匹配Flow Matching的训练框架中含噪数据可以统一表示为传统定义的 SNR 为这个定义只与时间步 t 相关完全忽略了数据本身的信号强度。然而现实中不同图像的信号强度如方差差异很大同一时间步在高信号数据和低信号数据上实际对应的 SNR 是完全不同的。2.2 扩展的数据相关 SNR 定义本文引入了一个与图像 I 相关的系数 K(I)将 SNR 扩展为其中是图像方差C(I) 是近似常数。这一扩展说明数据方差越大同一时间步对应的绝对 SNR 越高。这个视角带来了一个关键的实验便利通过对同一数据集缩放标准差std可以在不更换数据集的情况下定量模拟不同数据分布的影响从而系统研究训练策略的鲁棒性。2.3 训练 SNR 的概率密度函数PDF为了统一刻画噪声调度、损失权重和时间步采样策略论文将时间步的综合概率密度函数定义为其中为采样函数为损失权重函数。由此SNR 在训练中的概率密度函数可以通过概率变换推导并用大量离散采样估计见论文 Algorithm 1。PDF 的物理含义非常直观PDF 值越高的 SNR 区域在训练中被关注的频率越高。通过可视化 PDF可以直接看出不同训练策略把注意力放在哪里。[此处配图Figure 3]四种噪声调度DDPM-Linear、DDPM-Cosine、Flow-Linear、Flow-Cosine在不同 std 下的 SNR PDF 可视化。可以看到随着 std 增大PDF 整体右移不同调度的 PDF 宽度焦点宽窄差异显著。三、从 PDF 中能学到什么——百余组实验的系统分析本文在 ImageNet 128 分辨率上进行了超过 100 组实验采用 4 种噪声调度使用 8 张 H800 GPU每组训练 100,000 次迭代用 FID-10k 作为评估指标。3.1 数据信号强度显著影响训练效果即使使用相同的噪声调度不同信号强度std的数据会导致截然不同的训练结果。[此处配图Figure 2]以 DDPM-Linear 为例std0.3 时 FID 为 65.38而 std2.5 时 FID 改善到 47.04二者差距高达 18 以上。这充分说明数据信号强度是训练效果的重要影响因素。3.2 不同调度的鲁棒性差异显著随着 std 的变化不同噪声调度表现出截然不同的稳定性DDPM-LinearFID 随 std 变化的范围为18.03相对稳健DDPM-CosineFID 变化范围高达50.02极不稳定从 PDF 角度来看SNR 焦点更宽广PDF 方差更大的调度往往具有更好的鲁棒性Figure 3 中 DDPM-Linear vs. DDPM-Cosine 的 PDF 宽度对比。3.3 性能与鲁棒性之间存在权衡加权策略如 lognorm可以集中 PDF 的焦点从而提升最优性能的上限但同时会牺牲鲁棒性。[此处配图Figure 4]以 DDPM-Linear 为例不加权时最优 FID 为 47.04FID 变化范围为 18.02加入 lognorm(0,1) 加权后最优 FID 提升到 36.9但 FID 变化范围扩大到了80。性能上限提升了但稳定性大幅下降。这一权衡是普遍存在的研究者需要在追求极致性能和保证跨数据稳定之间做出权衡。3.4 Flow Matching v-prediction更鲁棒的选择将预测目标从 noise噪声换成 velocity速度可以显著提升鲁棒性[此处配图Figure 5]Linear Flownoise-prediction 的 FID 范围为 17.54换用 v-prediction 后降至11.79Cosine Flownoise-prediction 的 FID 范围为 34.76换用 v-prediction 后降至12.39不仅如此v-prediction 的整体性能也普遍优于 noise-prediction。即便如此v-prediction 下性能与鲁棒性的权衡依然存在加 lognorm 后性能上限提升但稳定性下降。3.5 核心结论从 PDF 视角来看高效训练需要同时满足两个条件PDF 应具有集中的焦点concentrated focusSNR 焦点需要落在正确的区间correct SNR range四、FasterDiT三步极简加速方案基于上述分析本文提出FasterDiT其核心思想是选择鲁棒性强的基础配置再通过数据移位和采样集中将 PDF 精准对齐到最优 SNR 区间最后引入方向损失进一步加速收敛。整个方案仅需少量代码改动无需任何架构修改。[此处配图Figure 7 右侧 Algorithm 2 代码块]步骤一数据移位Multi-step Balance / Data Shifting首先在 std0.5 到 1.2 的范围内扫描找到最优的标准差使 PDF 焦点移动到最优 SNR 区间。实验发现当 std 约为 0.70 时DiT 获得最佳生成性能。论文将目标 std 设为0.82与 SiT、DiT 的默认设置保持一致。[此处配图Figure 7(a) 左图Standard Deviation Modulating]在训练代码中这一步只需一行data data * (target_std / data_std)步骤二集中采样Concentrating在数据移位完成后PDF 焦点已对齐到正确区间此时再使用lognorm(0, 1)进行时间步采样进一步集中 PDF 焦点提升性能上限。注意mu 设为 0而非 SD3 中的 −0.5因为 PDF 已通过数据移位对齐到正确位置无需再做偏移从而避免了 Section 1 中描述的不稳定问题。t lognorm(mu0, sigma1, sizet.shape[0]) x_t t * data (1 - t) * noise pred model(x_t, t, y)步骤三改进监督——速度方向损失Velocity Direction Loss在流匹配框架下预测目标 velocity具有明确的物理意义——它代表从噪声到数据的流速方向。本文假设对速度方向的监督可以进一步提升训练效果。因此在标准 MSE 损失之外额外引入基于余弦相似度的方向损失最终损失为在代码中仅需两行v data - noise loss (pred - v)**2.mean() 1 - cosine_similarity(pred, v, dim1).mean()消融实验验证[此处配图Figure 7(b) 表格Ablations for FasterDiT Training]两项改进均带来了显著提升缺一不可。五、实验结果5.1 ImageNet 256×256 主实验[此处配图Table 1]FasterDiT 使用与 DiT 完全相同的 DiT-XL/2 架构在 ImageNet 256×256 上取得以下结果仅用1000k 迭代cfg1.5即达到 FID2.30与 DiT 需要 7000k 迭代才能达到的 FID2.27 相当训练速度提升约 7×训练 400k 迭代时FID-50k 达到11.9远优于 DiT19.5和 SiT17.2分别提升 5.3 和 5.6训练 200k 迭代时FID-50k 达到17.5相比 SiT 同等迭代数的 27.1 提升了9.6训练 2000k 迭代时FID 进一步达到2.03超越 SiT2.06[此处配图Figure 1 右图Performance of FasterDiT 表格]5.2 高分辨率生成ImageNet 512×512[此处配图Table 2]FasterDiT 在 512×512 分辨率上同样取得了一致的加速效果DiT-B/2200k 迭代时 FID-10k 从 77.11 →58.33提升18.78DiT-L/2200k 迭代时 FID-10k 从 67.29 →49.36提升17.935.3 跨架构泛化性[此处配图Table 3]FasterDiT 的训练策略不局限于 DiT在其他扩散架构上同样有效U-ViT-LFID-10k 从 50.22 →37.12提升 13.10LDM-UNetFID-10k 从 66.73 →60.07提升 6.665.4 生成效果可视化[此处配图Figure 8、Figure 9、Figure 10]FasterDiT-XL/2 仅训练 1000k 迭代CFG4.0即可生成高质量的 ImageNet 类别图像涵盖动物、建筑、自然景观等多种类别视觉质量与 DiT 相当。六、与相关工作的对比加速 DiT 训练的方法大致可分为两类架构修改类MaskDiT 和 MDT 结合掩码图像建模预训练加速训练SD-DiT 引入对比学习CAN 提出动态条件权重。这些方法均需修改模型结构增加了实现复杂度。非架构类本文方向调整噪声调度、预测目标、损失权重等。这一设计空间庞大且相互依赖难以系统分析。FasterDiT 正是在这一方向上首次提供了一个基于 SNR PDF 的统一分析框架并据此得出简洁有效的实践方案。七、局限性与未来工作论文作者坦诚指出了 FasterDiT 目前的局限性缺乏大规模实验未在 2K 高分辨率图像、文本到图像生成、视频生成等场景上验证文本-图像架构的潜在不稳定性在 SD3 等需要联合处理文本和视觉 token 的架构中来自 T5、VAE 等不同来源的特征可能引入潜在的不稳定因素需要进一步研究八、总结FasterDiT 的核心贡献可以概括为三点贡献内容意义新视角扩展 SNR 定义引入训练 SNR PDF 分析框架统一解释各类训练策略揭示鲁棒性本质新认知百余组实验揭示性能-鲁棒性权衡规律为训练策略设计提供直觉指导新方法数据移位 lognorm 采样 速度方向损失仅改几行代码训练速度提升 7×FasterDiT 的美妙之处在于它没有提出任何新的网络结构没有引入额外的预训练阶段只是在理解训练过程本质的基础上做了极其简洁的三步调整就让 DiT 的训练效率提升了整整 7 倍。这为后续扩散模型的训练策略研究提供了极具价值的思路与参考。