DiverseDiT: Towards Diverse Representation Learning in Diffusion Transformers
Organizations: Fudan University · Shanghai Academy of AI for Science · Shanghai Innovation Institute
Abstract
Recent breakthroughs in Diffusion Transformers (DiTs) have revolutionized the field of visual synthesis due to their superior scalability. To facilitate DiTs' capability of capturing meaningful internal representations, recent works such as REPA incorporate external pretrained encoders for representation alignment. However, the underlying mechanisms governing representation learning within DiTs are not well understood. To this end, we first systematically investigate the representation dynamics of DiTs. Through analyzing the evolution and influence of internal representations under various settings, we reveal that representation diversity across blocks is a crucial factor for effective learning. Based on this key insight, we propose DiverseDiT, a novel framework that explicitly promotes representation diversity. DiverseDiT incorporates long residual connections to diversify input representations across blocks and a representation diversity loss to encourage blocks to learn distinct features. Extensive experiments on ImageNet 256x256 and 512x512 demonstrate that our DiverseDiT yields consistent performance gains and convergence acceleration when applied to different backbones with various sizes, even when tested on the challenging one-step generation setting. Furthermore, we show that DiverseDiT is complementary to existing representation learning techniques, leading to further performance gains. Our work provides valuable insights into the representation learning dynamics of DiTs and offers a practical approach for enhancing their performance.
Figures & tables
| Model | Iter . | FID ↓ | sFID ↓ | IS ↑ | Prec. ↑ | Rec. ↑ |
| SiT-B | 400k | 36.80 | 6.77 | 40.09 | 0.51 | 0.63 |
| + (Ours) | 400k | 28.05 | 6.04 | 50.66 | 0.57 | 0.63 |
| REPA-B | 400k | 22.99 | 6.70 | 64.73 | 0.59 | 0.65 |
| + (Ours) | 400k | 17.29 | 6.56 | 79.92 | 0.62 | 0.65 |
| SiT-L | 400k | 18.77 | 5.27 | 71.44 | 0.64 | 0.63 |
| + (Ours) | 400k | 16.10 | 5.08 | 79.47 | 0.66 | 0.64 |
| Method | Epochs | FID ↓ | sFID ↓ | IS ↑ | Pre. ↑ | Rec. ↑ |
| ADM-U [ 14 ] | 400 | 3.94 | 6.14 | 186.70 | 0.82 | 0.52 |
| VDM++ [ 29 ] | 560 | 2.40 | - | 225.30 | - | - |
| CDM [ 25 ] | 2160 | 4.88 | - | 211.80 | - | - |
| LDM-4 [ 51 ] | 200 | 3.60 | - | 247.70 | 0.87 | 0.48 |
| SD-DiT [ 75 ] | 480 | 3.23 | - | - | - | - |
| MakDiT [ 73 ] | 1600 | 2.28 | 5.67 | 276.70 | 0.80 | 0.62 |
| Method | Epochs | FID ↓ | sFID ↓ | IS ↑ | Pre. ↑ | Rec. ↑ |
| ADM-G [ 14 ] | 400 | 2.85 | 5.86 | 221.70 | 0.84 | 0.53 |
| VDM++ [ 29 ] | - | 2.65 | - | 278.10 | - | - |
| MakDiT [ 73 ] | 800 | 2.50 | 5.10 | 256.30 | 0.83 | 0.56 |
| DiT-XL/2 [ 49 ] | 600 | 3.04 | 5.02 | 240.80 | 0.84 | 0.54 |
| SiT-XL/2 [ 44 ] | 600 | 2.62 | 4.18 | 252.20 | 0.84 | 0.57 |
| REPA [ 70 ] | 200 | 2.08 | 4.19 | 274.60 | 0.83 | 0.58 |
| Model | Iter . | Step . | FID ↓ | sFID ↓ | IS ↑ | Prec. ↑ | Rec. ↑ |
| MF-B/2 | 400K | 1 | 9.44 | 6.21 | 152.55 | 0.77 | 0.42 |
| + (Ours) | 400K | 1 | 8.51 | 5.95 | 158.84 | 0.78 | 0.44 |
| MF-L/2 | 400k | 1 | 8.73 | 6.11 | 161.69 | 0.79 | 0.40 |
| + (Ours) | 400k | 1 | 7.15 | 5.91 | 199.66 | 0.81 | 0.41 |
| MF-XL/2 | 400k | 1 | 5.94 | 6.10 | 213.13 | 0.83 | 0.41 |
| + (Ours) | 400k | 1 | 5.69 | 5.88 | 214.92 | 0.83 | 0.41 |
| Method | Params | Step | NFE | FID ↓ |
| iCT-XL/2 [ 58 ] | 675M | 1 | 1 | 34.24 |
| Shortcut-XL/2 [ 15 ] | 675M | 1 | 1 | 10.60 |
| IMM-XL/2 [ 74 ] | 675M | 1 | 2 | 7.77 |
| MF-XL/2 [ 17 ] | 676M | 1 | 1 | 3.43 |
| MF-XL/2 + DispLoss [ 62 ] | 676M | 1 | 1 | 3.21 |
| MF-XL/2 + ours | 713M | 1 | 1 | 2.99 |
| Component | FID ↓ | sFID ↓ | IS ↑ | Prec. ↑ | Rec. ↑ |
| SiT-B + full | 28.05 | 6.04 | 50.66 | 0.57 | 0.63 |
| w/o diversity | 32.77 | 6.42 | 44.85 | 0.54 | 0.63 |
| w/o residual | 33.72 | 6.53 | 43.97 | 0.53 | 0.63 |
| REPA-B + full | 17.29 | 6.56 | 79.92 | 0.62 | 0.65 |
| w/o diversity | 20.66 | 6.66 | 72.72 | 0.61 | 0.64 |
| w/o residual | 18.18 | 6.62 | 75.49 | 0.61 | 0.65 |
| Component | FID ↓ | sFID ↓ | IS ↑ | Prec. ↑ | Rec. ↑ |
| REPA-B + full | 17.29 | 6.56 | 79.92 | 0.62 | 0.65 |
| only | 18.97 | 6.64 | 75.44 | 0.61 | 0.65 |
| only | 17.70 | 6.62 | 78.34 | 0.62 | 0.65 |
| only | 20.85 | 6.78 | 68.74 | 0.61 | 0.65 |
| Range | FID ↓ | sFID ↓ | IS ↑ | Prec. ↑ | Rec. ↑ |
| Adap | 28.05 | 6.04 | 50.66 | 0.57 | 0.63 |
| Adap | 30.59 | 6.34 | 48.07 | 0.55 | 0.63 |
| Adap | 31.85 | 6.43 | 45.98 | 0.54 | 0.63 |
| Constant | diverge | - | - | - | - |
| Component | FID ↓ | sFID ↓ | IS ↑ | Prec. ↑ | Rec. ↑ |
| SiT-B | 36.80 | 6.77 | 40.09 | 0.51 | 0.63 |
| + Ours | 28.05 | 6.04 | 50.66 | 0.57 | 0.63 |
| ++ DispLoss [ 62 ] | 24.98 | 6.01 | 57.04 | 0.59 | 0.63 |
| +++ SRA [ 27 ] | 21.95 | 5.92 | 64.64 | 0.60 | 0.64 |
Appendix figures & tables29 assets
Supplementary material from the paper’s appendix.
Appendix
| Table 1 (SiT-B) | Table 1 (SiT-L) | Table 1 (SiT-XL) | Table 2 (SiT-XL) | Table 3 (SiT-XL) | |
| SiT + Ours | |||||
| Input dim. | 32 32 4 | 32 32 4 | 32 32 4 | 32 32 4 | 64 64 4 |
| Num. layers | 12 | 24 | 28 | 28 | 28 |
| Hidden dim. | 768 | 1,024 | 1,152 | 1,152 | 1,152 |
| Num. heads | 12 | 16 | 16 | 16 | 16 |
| REPA + Ours |
| Model | Alignment? | Encoder . | Align Depth . | Iter . | FID ↓ | sFID ↓ | IS ↑ | Prec. ↑ | Rec. ↑ |
| SiT-B | |||||||||
| ✗ | ✗ | ✗ | 50k | 89.65 | 12.58 | 18.28 | 0.34 | 0.43 | |
| ✗ | ✗ | ✗ | 100k | 40.46 | 6.15 | 36.26 | 0.52 | 0.49 | |
| ✗ | ✗ | ✗ | 200k | 31.42 | 5.87 | 58.08 | 0.61 | 0.53 | |
| ✗ | ✗ | ✗ | 400k | 12.69 | 5.29 | 106.21 | 0.71 | 0.54 | |
| ✗ | ✗ | ✗ | 450k | 12.04 | 5.24 | 110.19 | 0.71 | 0.54 |
| Model | Iter . | FID ↓ | sFID ↓ | IS ↑ | Prec. ↑ | Rec. ↑ |
| SiT-B | 400k | 43.46 | 7.53 | 36.80 | 0.60 | 0.64 |
| + (Ours) | 400k | 33.18 | 6.92 | 45.09 | 0.67 | 0.65 |
| REPA-B | 400k | 30.13 | 7.79 | 53.92 | 0.68 | 0.64 |
| + (Ours) | 400k | 23.82 | 7.76 | 64.62 | 0.70 | 0.63 |
| SiT-L | 400k | 22.75 | 5.78 | 64.05 | 0.73 | 0.63 |
| + (Ours) | 400k | 19.19 | 5.74 | 71.78 | 0.76 | 0.61 |
| Model | Iter . | Encoder . | Align Depth . | CFG . | FID ↓ | sFID ↓ | IS ↑ | Prec. ↑ | Rec. ↑ |
| REPA-B | |||||||||
| 400K | DINOv2-B | 5 | 1.0 | 22.99 | 6.70 | 64.73 | 0.59 | 0.65 | |
| 400K | DINOv2-B | 5 | 1.35 | 12.47 | 5.95 | 107.38 | 0.67 | 0.61 | |
| + Ours | |||||||||
| 400K | DINOv2-B | 5 | 1.0 | 17.29 | 6.56 | 79.92 | 0.62 | 0.65 | |
| 400K | DINOv2-B | 5 | 1.35 | 8.33 | 5.84 | 134.16 | 0.70 | 0.63 |
| Model | Iter . | FID ↓ | sFID ↓ | IS ↑ | Prec. ↑ | Rec. ↑ |
| SiT-B | 400k | 32.77 | 6.95 | 50.85 | 0.67 | 0.62 |
| + (Ours) | 400k | 23.96 | 6.45 | 63.07 | 0.73 | 0.61 |
| REPA-B | 400k | 21.27 | 7.34 | 78.25 | 0.73 | 0.62 |
| + (Ours) | 400k | 16.44 | 7.30 | 92.11 | 0.75 | 0.63 |
| SiT-L | 400k | 14.85 | 5.41 | 91.49 | 0.78 | 0.60 |
| + (Ours) | 400k | 12.44 | 5.41 | 101.08 | 0.79 | 0.58 |
| Component | FID ↓ | sFID ↓ | IS ↑ | Prec. ↑ | Rec. ↑ |
| SiT-B Baseline | 23.28 | 6.00 | 65.23 | 0.61 | 0.60 |
| SiT-B + full | 16.21 | 5.45 | 84.00 | 0.66 | 0.60 |
| w/o diversity | 20.07 | 5.72 | 73.65 | 0.63 | 0.60 |
| w/o residual | 20.76 | 5.77 | 69.23 | 0.61 | 0.61 |
| REPA | 12.47 | 5.85 | 107.38 | 0.61 | 0.62 |
| REPA-B + full | 8.34 | 5.64 | 134.16 | 0.70 | 0.63 |
| SiT-XL | 5 | 10 | 15 | 20 | all | |
| FID ↓ | 18.77 | 16.85 | 16.10 | 16.01 | 15.84 | 15.77 |
| IS ↑ | 71.44 | 77.62 | 79.47 | 82.05 | 83.95 | 85.64 |
| Time (h) | 18.66 | 19.90 | 21.02 | 23.45 | 25.96 | 28.50 |
| Component | FID ↓ | sFID ↓ | IS ↑ | Prec. ↑ | Rec. ↑ |
| SiT-B + full | 28.05 | 6.04 | 50.66 | 0.57 | 0.63 |
| only | 31.32 | 6.45 | 47.09 | 0.56 | 0.63 |
| only | 29.97 | 6.21 | 48.23 | 0.57 | 0.63 |
| only | 36.12 | 6.64 | 45.04 | 0.55 | 0.62 |
| Component | FID ↓ | sFID ↓ | IS ↑ | Prec. ↑ | Rec. ↑ |
| REPA Baseline | 12.47 | 5.85 | 107.38 | 0.61 | 0.62 |
| REPA-B + full | 8.34 | 5.64 | 134.16 | 0.70 | 0.63 |
| only | 10.98 | 5.78 | 115.03 | 0.68 | 0.62 |
| only | 10.78 | 5.76 | 115.95 | 0.69 | 0.63 |
| only | 8.59 | 5.77 | 131.69 | 0.70 | 0.63 |
| Component | FID ↓ | sFID ↓ | IS ↑ | Prec. ↑ | Rec. ↑ |
| REPA Baseline | 12.47 | 5.85 | 107.38 | 0.61 | 0.62 |
| SiT-B Baseline | 23.28 | 6.00 | 65.23 | 0.61 | 0.60 |
| + Ours | 16.21 | 5.45 | 84.00 | 0.66 | 0.60 |
| ++ DispLoss [ 62 ] | 13.73 | 5.76 | 95.31 | 0.68 | 0.60 |
| +++ SRA [ 27 ] | 11.25 | 5.37 | 108.15 | 0.69 | 0.61 |
Diffusion Transformer
Diffusion Transformers use attention-based networks to generate and transform images, video, audio, and other structured data through iterative denoising.
See all 468 papers