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
Figures & tables
Figure 1: The generated samples from MPDiT-XL with the cfg-scale w=3 at epoch 160.
Figure 2: Architecture of MPDiT , which consists of (a) the Global-Local MultiPatch Diffusion Transformer, (b) DiT Block with shared time embedding, (c) The Upsample Module and (d) The FNO Time Embedding
Model
Epochs
GFLOPs ↓
FID ↓
sFID ↓
IS ↑
Pre ↑
Rec ↑
ADM [ 16 ]
–
–
10.94
6.02
100.98
0.69
0.63
ADM-U
–
–
7.49
5.13
127.49
0.72
0.63
ADM-G, ADM-U
–
–
3.94
6.14
215.84
0.83
0.53
CDM
–
–
4.88
–
158.71
–
–
LDM-8 [ 52 ]
–
–
15.51
–
79.03
0.65
0.63
LDM-8-G
–
–
7.76
–
209.52
0.84
0.35
Table 1: Quantitative performance of MPDiT on ImageNet 256.
Model
N
k
Model Dim D
GFLOPs ↓
GFLOPsDiTGFLOPsMPDiT↓
MPDiT-B
12
6
768
16.6
72.1%
MPDiT-XL
28
6
1152
59.3
49.9%
Table 2: Model configuration and computational cost of MPDiT.
Method
Params (M)
GFLOPs ↓
FID ↓
DiT-B/2
130 .0
23.0
34.84
+ Shared AdaIN (α,γ,β)
90.3
22.9
35.31
+ Multitokens Class (m=16)
101.9
24.3
28.56
+ FNO Time Embedding
101.2
24.3
24.52
+ MPDiT (k=6) (default)
104.8
16.6
24.74
Table 3: Ablation on MPDiT components. All models are trained for 80 epochs under the same settings. Note that in the second row (shared AdaIN), we still apply AdaIN to the combined class and time embeddings as original DiT.
Table 6
Figure 3: Qualitative Result of Imagenet 512 with cfg=4
Model
Params(M)
GFLOPs
FID ↓
Prec ↑
Rec ↑
BigGAN [ 3 ]
160
-
8.43
0.88
0.29
StyleGAN-XL [ 56 ]
166
-
2.41
0.77
0.52
MaskGIT [ 5 ]
227
-
7.32
0.78
0.50
VQ-GAN [ 19 ]
227
-
26.52
0.73
0.31
VAR-d36-s [ 61 ]
2300
-
2.63
-
-
ADM-U [ 16 ]
731
2813
3.85
0.84
0.53
Table 8: Quantitative results of ImageNet 512 with MPDiT k=6
Method
Epoch
Params(M)
GFLOPs ↓
FID ↓
No cfg
DiT-XL/2
600
675
524.7
11.93
MPDiT {22,6}
120
482
228.4
9.24
MPDiT {18,6,4}
120
491
138.2
11.77
With cfg
DiT-XL/2
600
675
524.7
3.04
Table 9: Performace of different MPDiT-XL variants on ImageNet 512
Diffusion Transformers use attention-based networks to generate and transform images, video, audio, and other structured data through iterative denoising.
Diffusion Transformers (DiTs) have achieved state-of-the-art performance in image and video generation, but their success comes at the cost of heavy computation. This inefficiency is largely due to the fixed tokenization process, which uses constant-sized patches throughout the entire denoising phase, regardless of the content's complexity. We propose dynamic tokenization, an efficient test-time strategy that varies patch sizes based on content complexity and the denoising timestep. Our key insight is that early timesteps only require coarser patches to model global structure, while later iterations demand finer (smaller-sized) patches to refine local details. During inference, our method dynamically reallocates patch sizes across denoising steps for image and video generation and substantially reduces cost while preserving perceptual generation quality. Extensive experiments demonstrate the effectiveness of our approach: it achieves up to 3.52× and 3.2× speedup on FLUX-1.Dev and Wan 2.1, respectively, without compromising the generation quality and prompt adherence.
Diffusion Transformers (DiTs) have emerged as a core architecture in generative modeling due to their scalability and adaptability to multimodal tasks. DiTs comprise isotropic transformer blocks, and learn representations progressively across depth, where the denoising objective drives later layers to focus on fine-detail reconstruction. This results in degraded representation quality and an imbalanced encoder-decoder behavior. Prior approaches such as representation alignment (REPA) mitigate this by encouraging stronger early representations via training regularization. Alternatively, U-Net-style DiT architectures introduce explicit multi-scale encoder-decoder structures for improved convergence. But they build on standard U-Net wisdom via learnable operators for spatial downsampling, which are not well-suited to transformer architectures, introducing inefficiencies and compatibility issues with components such as cross-attention and representation regularization. In this work, we propose UDT, a U-Net diffusion transformer that combines the representation power of DiTs with the encoding-decoding benefits of U-Nets, through data-adaptive token merging for downsampling and upsampling, while preserving the DiT token dimension. Our baseline UDT architecture outperforms existing U-Net DiTs and achieves performance comparable to REPA across all model sizes. Furthermore, using architectural optimization and REPA, UDT outperforms SiT's 7.9 FID at 1400 epochs (w/o CFG) within 40 epochs (~ 40x faster convergence) for XL model size on 256x256 ImageNet. Finally, it achieves strong image generation performance with CFG, reaching FID of 1.38 (320 epochs) with SD-VAE and 1.35 (500 epochs) with VA-VAE, providing a new backbone for DiTs with strong empirical benefits.
Diffusion Transformers achieve strong image generation performance, but most operate in compressed latent spaces. Pixel-space diffusion avoids this information loss, yet existing approaches map each raw image patch to a single token, forcing one representation to handle both global communication and fine-grained details. We address this issue by proposing a new architecture, DuSPiT, a Dual-branch SubPatch Pixel Transformer. This model separates global structural reasoning from local appearance modeling. DuSPiT uses a compact base branch for efficient global reasoning and a parallel, high-capacity pixel branch, organized into subpatch groups, to preserve detailed appearance, with the two branches interacting through cross-attention. Our results show that DuSPiT generates images with richer details and stronger fine-grained structures, while also achieving a better quality--efficiency trade-off than prior pixel-space diffusion transformers.