Source-linked AI summary
MPDiT: Multi-Patch Global-to-Local Transformer Architecture For Efficient Flow Matching and Diffusion Model
Quan Dao, Dimitris Metaxas
TL;DR
Diffusion and flow-matching transformers remain expensive because isotropic token processing imposes substantial computation. MPDiT introduces a coarse-to-fine multi-patch architecture with revised time and class embeddings, reducing GFLOPs while preserving strong ImageNet generation quality. The paper reports improved efficiency and convergence, but extension to large-scale text-to-image and text-to-video settings remains open.
Problem
Diffusion models remain computationally expensive to train and sample, motivating more efficient transformer designs.
Method
MPDiT processes large patches in early blocks for global context, upsamples to smaller patches, and refines local details in later blocks, while revising time and class embeddings.
Results
Up to 50% GFLOPs reduction is reported while maintaining original performance, and the revised embeddings improve FID by 10 points.
Takeaways & Limitations
MPDiT improves training efficiency, memory usage, and sampling speed while maintaining strong generative performance.
Takeaways & Limitations
Extending MPDiT to large-scale text-to-image and text-to-video models remains open and requires substantial computational resources.
Abstract
from arXiv · showhide
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
1. Introduction
Diffusion models offer strong visual-generation performance but remain costly to train and sample. MPDiT addresses this by combining a global-to-local transformer architecture with revised time and class embeddings, validated on ImageNet.
- Diffusion models achieve strong visual-generation performance but remain computationally expensive to train and sample.
- MPDiT processes large-patch tokens early for global context, then progressively upsamples to small-patch tokens for local refinement.
- An FNO time embedding and multi-token class embedding revise diffusion-transformer conditioning to enrich temporal and class representations.
- The architecture applies global-local modeling at the network level rather than inside individual attention layers.
- Extensive ImageNet experiments evaluate the proposed architectural and conditioning designs.
2. Related Works
Prior efficiency efforts target training strategies, latent compression, token reduction, and alternative backbones. MPDiT instead focuses on the diffusion-transformer backbone and extends global-local modeling to the architectural level.
- Diffusion training remains computationally demanding because of slow convergence, motivating timestep sampling, loss reweighting, and representation-based acceleration methods.
- VAE designs reduce latent token counts through aggressive spatial compression, whereas MPDiT focuses on backbone design.
- Transformer-based diffusion models offer scalability, while linear attention and state-space alternatives seek to reduce computational or memory costs.
- MaskDiT demonstrates that aggressive token masking can degrade generation quality and may require full-token fine-tuning for recovery.
- MPDiT extends global-local modeling from individual attention layers to the architectural level to seek efficiency without the usual performance reduction.
3. Method
MPDiT uses a coarse-to-fine token hierarchy, expanding global representations before refining local details, and supplements this design with FNO time and multi-token class embeddings.
- 3.1. Preliminaries: Latent flow matching trains a velocity model on noisy latent representations, timesteps, and optional conditioning signals.
- 3.1. Preliminaries: The generative process integrates backward from Gaussian noise to a clean latent, then decodes it into an image.
- 3.2. Multi-patch Transformer: MPDiT uses p = 4 patches in early blocks, reducing ImageNet-256 tokens from 256 to 64 for efficient global modeling.
- 3.2. Multi-patch Transformer: An upsample block expands 64 tokens to 256, combines them with fine-grained spatial features through a skip connection, and sends them to final refinement blocks.
- 3.2. Multi-patch Transformer: A three-level hierarchy with p ∈ {8, 4, 2} is proposed for 512^2 resolution, while {4, 2} suffices for ImageNet-256.
- 3.3. Revisiting Time and Class Embedding: The method shares the time-embedding module across transformer blocks and separates time conditioning from learnable prefix class tokens.
- 3.3. Revisiting Time and Class Embedding: FNO time embedding models continuous flow dynamics using a gridded time signal and MixedFNO blocks.
- 3.3. Revisiting Time and Class Embedding: Multi-token class embedding prepends several learnable class tokens to image tokens, improving performance and convergence.
4. Experiment
Experiments on ImageNet evaluate MPDiT's generative quality, efficiency, and architectural components. The results show that multi-patch processing and redesigned embeddings improve efficiency and FID with limited degradation or unchanged computational cost in several ablations.
- Experimental Setup: ImageNet experiments evaluate FID, Inception Score, Precision, Recall, GFLOPs, and parameter count on 50K samples at 256 × 256 resolution.Evaluation uses Euler sampling with 250 steps.
- Main Results: MPDiT-XL achieves a non-cfg FID of 7.36 and cfg FID of 2.05 after 240 training epochs, while SiT reaches 9.35 after 1400 epochs.The comparison reports both generation quality and training duration.
- Component Ablation: Shared AdaIN reduces parameters from 130M to 90M, approximately 30%, while keeping GFLOPs unchanged and increasing FID by only 0.4.The ablation replaces per-block AdaIN with a shared AdaIN strategy.
- Component Ablation: Multi-token class embedding reduces FID by roughly 7 points, while FNO-based time embedding provides a further improvement of about 4 FID points.The ablation attributes these changes to richer class representations and smoother, more informative temporal structure.
- Patch-Scale Ablation: Using k = 6 fine-resolution blocks causes less than a 1 point FID drop relative to isotropic DiT while providing substantial efficiency gains.This result holds for both B and XL configurations.
- Class-Token Ablation: Using 16 class tokens substantially improves FID without changing GFLOPs, whereas 32 tokens increase computational cost for only marginal additional improvement.The result indicates a trade-off between class-conditioning capacity and computation.
5. Conclusion
MPDiT combines redesigned time and class embeddings with a global-to-local architecture to improve diffusion efficiency and generative performance. The paper reports substantial gains in FID and computational efficiency, while identifying large-scale multimodal generation as an open direction.
- Simple time and class embedding modifications improve FID by 10 points.
- The global-local MPDiT architecture reduces GFLOPs by up to 50% while maintaining original performance.
- MPDiT improves training efficiency, memory usage, and sampling speed.
- Extending MPDiT to large-scale text-to-image and text-to-video models remains an open direction requiring substantial computational resources.
Supplementary Material
The supplementary material covers additional MPDiT results, convergence and training-time comparisons, implementation details for Fourier neural operator time embeddings, and qualitative results.
- The supplementary presents additional MPDiT results and training-time and convergence comparisons against DiT and SiT baselines.
- It also provides Fourier neural operator time-embedding implementation details and additional qualitative results.
6. ImageNet 512 results
On ImageNet 512, MPDiT uses hierarchical patch configurations to achieve strong FID with substantially reduced computation. The reported experiments use standardized training and sampling settings, while variants expose a performance-efficiency trade-off.
- The ImageNet 512 experiments train for 120 epochs on a single A100 node with a total batch size of 256.
- Sampling uses Euler sampling with 250 NFEs, classifier-free guidance scale 1.375, and evaluation on 50,000 generated images.
- MPDiT-XLk=6 achieves an FID of 2.47, outperforming all baselines with 228.4 GFLOPs.It uses patch size 4 in the first 22 transformer blocks and patch size 2 in the last 6.
- MPDiT{22,6} surpasses DiT/SiT baselines after 120 training epochs while requiring 228.4 GFLOPs.
- MPDiT{18,6,4} uses approximately 26% of DiT/SiT GFLOPs, achieving FID 11.77 without guidance and FID 3.13 with classifier-free guidance.The three-level hierarchy uses patch sizes 8, 4, and 2 across successive block groups.
7. Convergence & Efficiency Analysis
MPDiT substantially reduces training and inference costs while preserving strong ImageNet performance. Across 256 and 512 resolutions, it reaches comparable or better reported FID with much lower training compute, faster convergence, faster sampling, and lower memory use.
- MPDiT-XL reaches FID 2.05 on ImageNet 256 using 8.8% of DiT/SiT total training compute, corresponding to approximately 11.36× faster convergence.
- MPDiT-XL reaches FID 2.47 on ImageNet 512 using 8.7% of DiT/SiT total training compute, corresponding to approximately 11.5× faster convergence.
- MPDiT achieves more than 2× faster sampling than DiT and SiT under the same GPU and number of function evaluations.It also consumes less memory during inference.
- At 256 × 256 resolution, MPDiT fits batch size 1024 on one node with eight A100 GPUs, unlike DiT/SiT under identical settings.
- At 512 × 512 resolution, MPDiT fits batch size 256 on the same hardware, enabling training without multi-node clusters or higher-memory GPUs.
8. FNO Time Embedding Details
The FNO time embedding uses a 1D grid as time features; replacing it with cosine-sine sinusoidal features prevented convergence in the authors’ experiments.
- The implementation uses a 1D grid as the time features for FNO time embedding.The authors report that replacing this grid with traditional cosine-sine sinusoidal features made the model unable to converge.
9. More Qualitative Results
The paper provides additional qualitative ImageNet samples and implementation details for the FNO time embedding, including spectral convolutions, local convolutions, and Fourier-domain processing. The qualitative images are non-cherry-picked samples generated with Euler 250 steps and CFG scale 4.
- Qualitative results: The qualitative images are non-cherry-picked samples generated with Euler 250 steps and CFG scale 4.
- FNO implementation: The FNO timestep embedder lifts a scalar timestep into a width-dimensional signal, applies three spectral and local-convolution blocks, then averages and projects to model dimension.The implementation uses linear lifting and projection, three spectral convolution layers, and three local 1x1 convolutions.
- FNO implementation: The time signal is formed by centering a 1D grid around the timestep before Fourier-domain processing.The grid spans -1 to 1 and is shifted by t for each batch element.
- FNO implementation: The spectral convolution computes an FFT, applies learned complex weights to selected modes, and uses an inverse FFT to return to the spatial domain.The code stores real and imaginary Fourier weights and converts the output back to the original precision.
- Qualitative results: Additional qualitative results cover ImageNet classes including snail, loggerhead turtle, peacock, macaw, balloon, mushroom, volcano, and bubble.The figures also include box turtle and golden retriever classes.