MPDiT：用于高效流量匹配和扩散模型的多补丁全局到局部 Transformer 架构

Transformer 架构，特别是扩散 Transformer (DiT)，由于其比卷积 UNet 更强的性能，已广泛应用于扩散和流匹配模型。然而，DiT 的各向同性设计在每个块中处理相同数量的补丁标记，导致训练过程中计算量相对较大。在这项工作中，我们引入了多补丁 Transformer 设计，其中早期块在较大补丁上运行以捕获粗略的全局上下文，而后期块使用较小补丁来细化局部细节。这种分层设计可以将 GFLOP 的计算成本降低高达 50%，同时实现良好的生成性能。此外，我们还提出了时间和类别嵌入的改进设计，以加速训练收敛。对 ImageNet 数据集的大量实验证明了我们架构选择的有效性。代码发布于：https://github.com/quandao10/MPDiT