备注好友：
方向-学校/公司-姓名/昵称
【AIGC 先锋科技】交流群
扩散 Transformer （DiT）在研究中受到了广泛关注。然而，它们存在收敛速度慢的问题。在本文中，作者旨在在不进行任何架构修改的情况下加速DiT训练。
作者发现在训练过程中存在以下问题：首先，某些训练策略在不同数据集上表现不稳定。其次，在特定时间步长的监督效果有限。
针对这些问题，作者提出了以下贡献：
首先，作者提出了一个新的视角来解释策略失败的原因。具体而言，作者略微扩展了信噪比（SNR）的定义，并建议观察SNR的概率密度函数（PDF）以理解策略的数据鲁棒性本质。
其次，作者进行了大量实验，并报告了超过100个实验结果，以从PDF的视角实证地总结出一个统一的加速策略。最后，作者开发了一种新的监督方法，进一步加速了DiT的训练过程。
基于这些，作者提出了 FasterDiT ，这是一个极其简单且实际可行的设计策略。
通过少量代码修改，它在1000k迭代时，在ImageNet 256分辨率上实现了2.30 FID，与DiT（2.27 FID）相当，但训练速度快7倍。
1 Introduction
随着Sora 的出现，其基础模型 Diffusion Transformers（DiT） [36]，引发了广泛的研究兴趣。DiT以其惊人的灵活性和可扩展性而著称，在图像和视频生成方面展现了卓越的能力。然而，与视觉 Transformer （Vision Transformers） [15]面临类似挑战一样，DiT与高训练成本相关联。其收敛速度仍然较慢，在ImageNet生成任务上达到256分辨率需要超过4700个GPU小时。这一显著的计算需求凸显了在大规模训练中提高训练效率的必要性。
提高训练效果的有效方法之一是调节训练过程中不同时间步的信号噪声比（SNR）分布。去噪生成模型 [22, 29] 创建了从噪声到数据的转换，随着时间步（t）的推进，有条不紊地将噪声转换为数据。在这个过程中，SNR 逐渐从零增加到无穷大。给定一个生成过程
。假设输入数据
理想情况下呈正态分布，方差为 1，SNR 通常定义为方差的比值
。在每次训练步骤中，对于一对输入数据和噪声，作者随机选择一个
进行训练。调节噪声调度 [34, 25, 10]，损失权重 [33, 20] 和时间步采样策略 [16] 是常见的在训练过程中修改 SNR 分布的方法。它们被证明是提高训练效率和有效性的有效途径。
然而，这些方法并不总是在不同数据集上有效。例如，Stable Diffusion 3 (SD3) [16] 提出了一种通过对数正态分布（lognorm）[1] 进行时间步长采样策略来获得更好的训练。对数正态采样是一种出色的先导策略，并在后续工作中被证明是有效的[17]。
但是，它也存在一些局限性。根据SD3的探索，对数正态(-0.5, 1)在ImageNet[13]上的Frechet Inception距离（FID）结果明显优于对数正态(0.5, 0.6)。然而，作者发现这个结论是特定于所使用的训练数据的特定解决方案。如图1所示，随着信号强度的持续降低，对数正态(-0.5, 1)的FID结果逐步恶化，最终成为最差的结果。这强调了在不同的数据条件下评估这些方法鲁棒性的必要性。
在本文中，作者旨在为解释问题提供更全面的视角。作者的第一个贡献是建议从训练过程中的SNR（信噪比）的概率密度函数（PDF）的角度来解释性能的稳健性。具体而言，这些方法通常根据时间步长
来调整SNR的分布。尽管相同的时间步长具有相同的相对SNR（
），但实际SNR在数据信号增强时会增加[25, 10]。作者认为，SNR的原定义仅反映了同一数据分布中的相对信号噪声比。然而，训练过程中与数据相关的绝对SNR的分布是决定训练有效性的关键因素。因此，作者稍微扩展了SNR的定义并可视化了其概率密度函数。
作者的第二个贡献在于进行了广泛的实验，并在论文中报告了一百个实验结果，以实证分析训练性能和鲁棒性之间的关系与PDF。作者在DiT [36, 31]上分析了常用基本 Pipeline 之间的数据鲁棒性差异。作者发现方法鲁棒性和性能之间存在权衡，并建议从PDF的角度更直观地设计训练过程。
另一种提高训练的方法是修改单步预测或监督 [33]。作者的第三项贡献是引入了一种新的速度预测基础方法的速度预测监督方法。使用相同的调度，不同的预测目标 [33] 或不同的监督方法 [22, 14] 也可能导致不同的训练结果。最近，速度已被证明是一种优越的预测目标。具体来说，除了传统的均方误差（MSE）损失，作者还引入了速度方向的监督。作者发现这种方法在概念上很简单，并且显著加速了训练过程。
总之，本文提出了一种新的训练SNR PDF的视角，以提供对生成模型训练效率的直接解释（第2节）。作者进行了广泛的实验，并报告了近一百个实验结果，旨在从实证上推导出训练性能、鲁棒性及其与SNR PDF之间的关联（第3节）。
随后，作者将这些观察应用到改进DIT过程，并引入了一种新的监督方法，开发了FasterDiT，以显著提高训练效率（第4节）。FasterDiT在ImageNet上的FID达到2.30，分辨率为256，与原始DiT的FID 2.27相当，但收敛速度快7倍。作者希望作者的探索为生成模型训练的未来研究提供有价值的洞察。
2 Probability Density Function of SNR during Training
作者假设在训练过程中，对不同信噪比（SNRs）的关注分布是训练效率和效果的关键决定因素。在此，有两个问题需要解决。首先，以前将SNR定义为
忽略了数据信号强度的影响，这在训练过程中被证明是重要的 [25, 34]。
其次，需要一种统一且直观的方法来分析训练SNR分布。因此，在本节中，作者对现有SNR定义进行了一定程度的修改。随后，作者利用训练过程中SNR的概率密度函数（PDF）将噪声调度、损失加权以及时间步长采样策略整合到一个连贯的框架中。
Preliminary
为了便于理解，作者对生成模型的公式和训练流水线进行简要介绍[22, 34, 28, 29]。由于作者关注的是SNR分布，作者从高层次看待流匹配和扩散模型，类似于以前的工作[16]。给定数据
和高斯噪声
，作者将噪声和数据的传输定义为方程1。在一个离散扩散过程中[22, 34]，
是一个从0到1000的整数。在一个连续流过程中[28, 29]，
是一个连续值，介于[0, 1]之间。
和
是与
相关的系数，根据不同的时间表定义。
在训练过程中，噪声数据
和时间步长
被输入到生成模型中。模型需要预测特定的目标（噪声、速度等）。例如，在最常用的DDPM[22]流水线中，预测目标是噪声
，损失函数定义为方程2。在每个迭代中，对于一张图像，只训练一个
，而不是整个时间步长。因此，一些方法选择改变
的采样[16]，或者根据
给出不同的损失权重[20, 34]。
Formulate Probability Density Function of SNR
基于数据的SNR
相关工作[24, 14, 33, 38, 20]强调了SNR的重要性。通常，它们假设数据是理想归一化分布，并将其定义为方差的比值（方程3）。这一定义与数据无关，仅与时间步长
中的系数相关。
然而，实际数据的分布是无法描述的。它与图像本身的性质密切相关。在本研究中，作者引入了一个与图像
相关的系数
，它直接缩放信号。
的值受到图像的各种特性（如高频率和低频率范围、分辨率、方差等）的影响。已知，图像的方差越大，其信噪比（SNR）越高，如图4所示。此外，对于特定数据集，作者假设方差的变化对图像的其他性质（即
）影响很小，即
是
的一个独立变体，并近似将
视为常数
。
作者的目标是探索不同数据SNR变化如何导致训练的改变。然而，不断更改训练数据既不能量化也不能方便。上述转换（方程5）的优势在于，作者可以使用相同的数据集，但具有不同的std，来模拟不同的C(I)变化。这简化了检测数据鲁棒性的方案，并使作者无需反复更改训练数据。
作者考虑一个带有时间步长采样函数
和损失权重函数
的训练过程。在大多数情况下，
和
是非负的。为了将它们统一到一个单一的
分布中，作者定义一个新的概率密度函数
如下：
在SNR训练过程中，作者试图获得SNR的概率密度函数。它与时间步长t无关。假设分布SNR(t)服从Y，且t遵循f_t(t)的分布。求解Y的分布可以转化为概率转换问题。从数学上讲，它可以定义为等式7。
估计SNR的概率密度函数在实际应用中，上述函数并非总是可用的。因此，作者使用大量的离散样本来近似方程7。在这里，作者提供其估计算法。对于所有训练过程，作者按照算法1所示的方式可视化它们的概率密度函数。
现在，作者已经开发了一种简单且可行的方法来生成SNR的PDF。它直接反映了在训练过程中对不同SNR值的关注程度。具体来说，PDF值越高，在训练过程中对应SNR值训练的频率就越高。接下来，作者将进行广泛的实验来分析它与不同训练数据和训练策略之间的关系。
3 What can we learn from SNR PDF?
回到作者在第一部分讨论的问题图1。作者主张，仅基于t的战略倾向于只关注生成过程中的相对较高或较低噪声条件，但仅考虑不同信号强度下的绝对信噪比。它们为提高训练提供了令人惊叹的解决方案，但缺乏全面的评估。在本部分，作者将从SNR（信噪比）的概率密度函数（PDF）的角度探讨这个问题，以识别这些策略之间的共同点，并划分出一个更全面的设计空间。
Experiment Settings
作者首先介绍实验设置。作者主要采用来自两个先前重要工作的扩散 Transformer （DiT）训练 Pipeline ，分别是DiT [36]和SiT [31]。作者从它们中选择了四种不同的噪声调度，包括DDPM具有线性贝塔调度[22]，DDPM具有余弦贝塔调度[34]，流量匹配具有线性调度（正则化流量）[29]，以及流量匹配具有余弦调度[31]。
在随后的实验中，作者最初将每个噪声计划的预测目标设置为噪声（ε）以确保公平的比较。随后，作者通过缩放数据到各种标准差（std）来修改输入强度。然后，作者将流量匹配的预测目标设置为与原始设置相同的速度，以比较速度预测和噪声预测结果的相似性和差异。每个实验均在ImageNet [13]上进行，分辨率为128。作者对每个模型进行100,000次迭代训练，并使用FID-10k指标进行比较分析。每个实验均使用了8块H800 GPU。
作者的主要观察目标有两个：
（1）在不同信号强度下的性能和鲁棒性；
（2）PDF状态与性能和鲁棒性之间的潜在联系。
Insights from PDF
不同的数据信号强度会导致不同的训练效果。在图2中，所有计划的预测目标都统一设置为噪声，以确保公平的比较。曲线展示了不同输入强度如何影响训练结果。值得注意的是，单个噪声计划的性能随数据强度变化而波动。例如，线性贝塔调度器的DDPM在
时，FID达到65.38，而在
时，显著提高至47.04（图2a）。
不同的调度方案在数据鲁棒性方面存在显著差异。在图2中，随着数据强度的变化，不同的调度方案表现出显著不同的鲁棒性曲线。具体来说，随着标准差增加，DDPM-线性[22]方案的FID范围观察到为18.03（图2a）。相比之下，DDPM-余弦[34]方案的FID范围要大得多，为50.02（图2c），这突显了不同调度方案之间性能的巨大差异。
带有广泛SNR关注的时间表可能具有强大的性能。在图3中，作者呈现了第2节中概述的SNR PDFs。对于每个噪声时间表，作者显示了三个与不同信号强度相应的PDF曲线：低强度（标准差为0.3），高强度（标准差为3.3）和最佳FID对应的强度。这些PDF曲线从视觉上代表了训练过程中在不同噪声间隔上的重点。通过分析不同噪声时间表曲线形状的实证分析，作者发现PDF曲线形状的更大波动性，表明对噪声间隔考虑更广泛，与时间表的稳定性提高相关（图3a&b）。
不同调度下的最优SNR范围似乎相似。直接观察不同调度下标准差的变化并不能立即得出一致的结论。例如，在图2中，DDPM-线性调度（图2a）在标准差约为2.5时实现最优生成性能，而DDPM-余弦调度（图2c）在标准差为0.9时达到峰值。然而，从PDF曲线的分析中可能可以获得一些洞见。在作者的实验中，如图3所示，在统一预测目标下，更好的FID性能与PDF的平均值位于指定的灰色区域内相关。这一观察为作者提供了一个新的视角，丰富了作者对之前结果的理解。
性能与鲁棒性之间存在权衡。在训练过程中使用加权的主要目的是加强特定SNR水平上的关注。作者实现了一个对数正态函数（lognorm）[1]来调整两种不同的时间表（图2a&b）的步进采样，结果如图4所示。作者的分析表明，虽然定义良好的采样策略可以提高性能，但也可能引入风险。例如，DDPM-线性时间表实现了最优FID为47.04，表明在18.02的范围内具有鲁棒性能。
然而，当作者使用lognorm加强关注（图4a），尽管最佳FID提高到36.9，但也导致了更不稳定的结果，性能范围扩大到80（图4b）。
Stream 匹配与v-预测得到更强的性能。最近， Stream 匹配[28, 29]被视为生成模型的更简洁、高效流水线。在之前的讨论中，为了进行公平的比较，作者将预测目标设置为噪声预测。如图5所示，作者发现 Stream 匹配与v-预测实现了更稳定的输出性能，显示出在数据稳定性方面的增强鲁棒性。例如，线性 Stream 在噪声预测下FID范围为17.54，而在v-预测下为11.79（图5a）。这一特性在余弦时间表的 Stream 中更加明显，范围分别为34.76和12.39（图4(b)）。此外，作者还发现v-预测的整体性能通常超过噪声预测。
仍存在与v-预测的流匹配中的权衡。此外，作者观察到尽管v-预测通常更鲁棒，但性能和稳定性之间仍存在权衡。如图4(c)和d所示，当使用对数正态来集中线性流的重点时，其性能类似于噪声预测，其中上限改进但稳定性降低。作者假设这样的权衡可能广泛存在。
总结起来，从PDF的角度来看，高效的训练需要满足两个条件：
(1) PDF应具有集中关注的特性；
(2) SNR的关注区域需要落在正确的范围内。基于简单的观察，作者试图改进DIT训练。
4 Improving DiT Training.
在本文中，作者的目标并非设计新的模型架构以实现最先进的性能。相反，作者旨在探索一种更简单、更易解释和更普遍适用的训练方法，用于深度 Transformer （DiTs）[36]。
在第三节中，作者通过信号噪声比（SNR）概率密度函数（PDF）分析展示了DiT训练过程中性能和鲁棒性之间的权衡。为了加速训练，作者需要在训练过程中关注正确的SNR。
最初，作者选择流匹配[29]并使用v-预测作为噪声计划，因为其鲁棒性和优越性能。然后，作者通过调整训练数据的方差（std）来将PDF调整到关注训练过程中的最优SNR。进一步，作者使用对数正态函数[1]来突出调整区域的重点。最后，作者引入了一种新的、直接的控制策略，显著提高了训练效果。
Improving Multiple Step Balance
在这里，作者首先将标准差从0.5调整到1.2，如图6（左）所示。在这个扫描过程中，作者避免使用其他技术，如加权。结果表明，当
约为0.70时，DiT训练具有更好的生成性能。实际上，作者将目标
设置为0.82，这是合适的，并且与以前工作的设置[31; 36]一致。然而，作者认为，当输入数据发生变化，例如不同的分辨率，选择
将有所不同。这种策略类似于以前的工作[10]，但作者在训练SNR PDF的新视角下进行讨论。
然后，作者通过人为的浓度来提高PDF性能。具体来说，作者在时间步长采样中使用logit正态函数[1]，其中
和
分别设置为0和1。值得注意的是，将
设置为0的原因是PDF已经正确偏移（参见图3）。这可能有助于避免第1节中讨论的不稳定性，并提高训练性能。
Improving Single Step Supervision
有效的单步监督对于训练生成模型至关重要。例如，使用不同的预测目标可能导致不同的预测结果[33]。在原始的DIT模型中，同时预测噪声和sigma的策略被用于最大化性能[36]。通常，均方误差（MSE）损失函数被用于监督这些目标。
最近，在 Stream 匹配[28; 29]的背景下，速度预测的目标具有更具体的物理意义，代表从噪声到数据的流量。基于此，作者假设监督速度方向可以作为一种有效的监督策略。
Comparison with Previous Methods
由于资源限制，作者在ImageNet上使用256分辨率训练FasterDiT模型1000k次迭代，并将其与当前最先进的生成模型（见表1）进行比较。作者的实验是在8块H800 GPU上进行的。值得注意的是，作者工作的重点是探索超越结构增强的训练策略。因此，FasterDiT使用了与DiT相同的架构。如图1所示，在400k次迭代时，FasterDiT的FID-50k得分达到11.9，明显优于原始DiT模型（FID-50K 19.5）和其增强版SiT（FID-50k 17.2）[31]。使用与DiT相同的架构，FasterDiT取得了相似的性能水平，但只需1000k次迭代即可收敛，其FID得分为2.30。
5 Related Work
Generative Models and Diffusion Transformers
去噪扩散概率模型[22]逐渐取代了GANs[12; 19]，因为其更强的性能和更稳定的训练，成为主流的生成模型。其中，[14]首先在ImageNet[13]上证明了其有效性。 [38; 37]将扩散过程从扩散空间扩展到VAE潜在空间，实现了高性能和高分辨率的生成。同时，一些方法[31; 44; 17]也使用流匹配[28; 29]技术训练生成模型以替代扩散路径，这些方法具有更简单的数学特性和更快的学习效果。
其中，U-Net [39] 是生成模型中最常用的架构。最近，由于 Transformer 的可扩展性优势 [15; 42]，一些基于 Transformer 的生成模型也出现了。然而，与视觉 Transformer [15] 类似，Diffusion Transformers [36] 收敛速度较慢。
Fast Training of Generative Models
Diffusion Transformers（DiT）的加速工作可以简单地分为两个类别。一个是架构修改。MaskDiT [43] 和 MDT [18] 结合了 Mask 图像建模预训练和扩散训练，以加速。同样，SD-DiT [44] 进一步将DiT训练与类似MoCo的对比学习相结合。CAN [5] 提出了一个动态权重用于条件，以加速扩散模型[36, 2]的训练。
另一种策略涉及改变非模型设计方法。例如，使用不同的噪声时间表可以在各种分辨率的图像上实现更好的结果。使用不同的预测目标，如噪声、数据或速度，也可以直接影响训练的有效性[33]。调整损失函数和训练采样权重也可以直接影响训练结果[20, 16]。然而，它们构成了一个设计空间，其中高度直接相互依赖，这使得模型设计非常复杂。
6 Conclusion
在本文中，作者讨论了加速Diffusion Transformer的训练策略。首先，作者对SNR的定义进行了一定程度的泛化，并通过训练过程中SNR概率密度函数的分析，对各种训练策略进行了统一分析。
作者发现在扩散模型的训练中，鲁棒性和性能之间存在权衡。不同的噪声时间表表现出不同程度的鲁棒性。在训练过程中使用加权可以提高性能天花板，但可能会降低过程的鲁棒性。
此外，作者发现使用方向损失作为辅助损失函数进行速度预测，可以显著提高训练性能。
基于这些观察，作者在ImageNet256上使用DiT-XL-2进行了实验，并观察到训练速度的巨大加速。作者将这种简单的训练方法称为FasterDiT。
作者希望作者的工作能激发对生成模型训练策略的进一步探索。
[0]. FasterDiT: Towards Faster Diffusion Transformers Training without Architecture Modification.