Scaling Diffusion Transformers to 16 Billion Parameters

本文介绍了 DiT-MoE，它是扩散 Transformer 的一种稀疏版本，具有可扩展性和与密集网络相竞争的性能，同时具有高度优化的推理能力。DiT-MoE 包括两个简单的设计：共享专家路由和专家级平衡损失，从而捕捉常见知识并减少不同路由专家之间的冗余。在应用于条件图像生成时，对专家特化的深入分析得出了一些有趣的观察结果：（i）专家选择显示出对空间位置和去噪时间步的偏好，而对不同类条件信息不敏感；（ii）随着 MoE 层的加深，专家选择逐渐从特定的空间位置转向分散和平衡。（iii）专家特化倾向于在早期时间步骤更集中，然后在一半后逐渐均匀。我们将其归因于扩散过程，首先对低频空间信息进行建模，然后对高频复杂信息进行建模。基于以上指导，一系列 DiT-MoE 实验在推理期间需要较少的计算负载，但实现了与密集网络相当的性能。更令人鼓舞的是，我们用合成图像数据展示了 DiT-MoE 的潜力，在 512×512 分辨率设置下，扩散模型的参数规模达到了 16.5B，获得了新的 SoTA FID-50K 分数 1.80。项目页面：https://github.com/feizc/DiT-MoE。

本论文旨在提出一种可扩展且具有高度优化推理的稀疏扩散Transformer模型，以及应用于条件图像生成的实验研究。主要问题是如何在保持性能的同时减少计算负担。

该模型包括两个简单的设计：共享专家路由和专家级平衡损失，从而捕捉常见知识并减少不同路由专家之间的冗余。实验结果表明，该模型在推理过程中需要更少的计算负载，同时能够达到与密集网络相当的性能。

实验结果显示，专家选择对空间位置和去噪时间步显示出偏好，而对于不同的类条件信息则不敏感。随着MoE层数的增加，专家的选择逐渐从特定的空间位置转向分散和平衡。专家的专业化倾向于在早期时间步骤上更集中，然后在一半之后逐渐变得均匀。该模型在512x512分辨率设置下，通过合成图像数据，成功地扩展了一个165亿参数的扩散模型，并获得了新的SoTA FID-50K得分1.80。作者还提供了开源代码。

近期相关研究包括：《Efficient Transformers: A Survey》、《Routing Transformer: Towards Efficient and Accurate Neural Machine Translation》、《MoE Transformer》等。