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.