用大块先看全局、小块再抠细节，让扩散模型更快更省算力
MPDiT: Multi-Patch Global-to-Local Transformer Architecture For Efficient Flow Matching and Diffusion Model
- 早期用大块处理全局信息，后期用小块细化局部细节
- 在ImageNet上计算量减少50%GFLOPs，生成质量不降反升
- 适合追求高效训练的扩散模型研究者和开发者
Transformer架构，尤其是扩散Transformer（DiT），因其相比卷积UNet的优异性能，已广泛应用于扩散模型与流匹配模型。然而，传统DiT采用各层相同数量的分块令牌处理，导致训练时计算开销较大。本文提出一种多分块全局到局部的Transformer结构：早期模块使用较大分块以捕捉粗粒度全局上下文，后期模块采用较小分块以精细重构局部特征。该层级化设计使计算量最多降低50% GFLOPs，同时保持良好生成性能。此外，我们还改进了时间与类别嵌入的设计，加速训练收敛。在ImageNet数据集上的大量实验验证了所提架构的有效性。代码已开源：https://github.com/quandao10/MPDiT
原文摘要 · Abstract (English)
Transformer architectures, particularly Diffusion Transformers (DiTs), have become widely used in diffusion and flow-matching models due to their strong performance compared to convolutional UNets. However, the isotropic design of DiTs processes the same number of patchified tokens in every block, leading to relatively heavy computation during training process. In this work, we introduce a multi-patch transformer design in which early blocks operate on larger patches to capture coarse global context, while later blocks use smaller patches to refine local details. This hierarchical design could reduces computational cost by up to 50% in GFLOPs while achieving good generative performance. In addition, we also propose improved designs for time and class embeddings that accelerate training convergence. Extensive experiments on the ImageNet dataset demonstrate the effectiveness of our architectural choices. Code is released at: https://github.com/quandao10/MPDiT
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。