Search papers, labs, and topics across Lattice.
This paper introduces MPDiT, a multi-patch transformer architecture for diffusion and flow-matching models that uses larger patches in early blocks for global context and smaller patches in later blocks for local refinement. This hierarchical design reduces computational cost by up to 50% in GFLOPs while maintaining strong generative performance. The authors also propose improved time and class embedding designs to accelerate training convergence, validated on ImageNet.
Halve the training cost of your diffusion transformer without sacrificing generative performance by using multi-patch hierarchies.
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 \url{https://github.com/quandao10/MPDiT}
