[CVPR 2024] DiverseDiT: Breaking Feature Homogeneity for SOTA Diffusion Transformers

TL;DR

DiverseDiT is a breakthrough in Diffusion Transformer (DiT) training that identifies representation diversity as the secret sauce for high-quality image synthesis. By introducing long-range residual connections and a specialized diversity loss, it achieves superior results to REPA (which uses external encoders) while being faster to train and more scalable. It effectively hits 1.52 FID on ImageNet 256x256 in just 200 epochs.

Problem & Motivation: The "Copy-Paste" Problem in DiTs

Despite the success of DiT-based models (like Sora or Stable Diffusion 3), our understanding of their internal representation learning remains opaque. A common issue in deep transformers is representational homogenization: layers eventually start learning the same features, leading to redundancy and "wasted" parameters.

Previous "Alignment" methods like REPA tried to fix this by forcing DiT layers to match features from a pre-trained DINOv2 or MAE encoder. However, the authors of DiverseDiT ask a fundamental question: Is the external guidance necessary, or can we fix the internal dynamics directly?

Through Centered Kernel Alignment (CKA) analysis, they discovered that:

1. Diversity across blocks naturally increases during successful training.
2. Aligning too many blocks to external models actually reduces overall diversity, harming performance.
3. The real benefit of external alignment was simply helping blocks specialize.

Methodology: Engineering Diversity

Based on these insights, the authors proposed two "plug-and-play" components that don't require external models.

1. Long Residual Connections

Standard DiTs use local residuals (Layer $ioi+1$). DiverseDiT adds "Long" residuals that connect early layers to deep layers (e.g., Block $i$ to Block $L-i$).

• The Intuition: By mixing "fresh" low-level features with deep high-level features, it prevents the input to deep blocks from becoming a homogenous soup.

2. Triple-Threat Diversity Loss

The model is optimized with a specific ${\mathcal{L}}_{div}$ composed of three parts:

• Orthogonality Loss: Forces the mean features of different blocks to be perpendicular in latent space.
• Mutual Information (MI) Minimization: Uses a cosine-similarity proxy to ensure blocks are statistically independent.
• Feature Dispersion: Encourages the model to use all available channels, preventing channel "shutdown."

Fig 1: The DiverseDiT framework showing the Long Residual Connections and the Diversity Loss points.

Experiments & Results: More for Less

DiverseDiT was tested on ImageNet at 256px and 512px resolutions. The results are striking because they show convergence acceleration.

• Efficiency: DiverseDiT (XL model) at 200 epochs (FID 1.52) beats the original SiT-XL at 1400 epochs (FID 2.06).
• Scalability: Whether using a Base (B), Large (L), or Extra Large (XL) backbone, the diversity-promoting components consistently dropped FID scores.
• One-Step King: When applied to "MeanFlow" (a one-step generation baseline), DiverseDiT set a new SOTA with an FID of 2.99, significantly outperforming Apple's Shortcut models and other consistency-based methods.

Fig 2: Visual comparison showing that DiverseDiT produces clearer structures and fewer artifacts at much earlier training iterations.

Ablation Insight: Why does it work?

The authors found that using a constant weight for diversity loss causes the model to diverge. Instead, they used an Adaptive Weight $(w)$. If the features become too different, the loss scales down, allowing the model to maintain the necessary "shared knowledge" required to reconstruct a coherent image.

Critical Analysis & Conclusion

DiverseDiT is a refreshing take on representation learning in generative models. It moves away from the "bigger is better" trend of external alignment and focuses on architectural Inductive Bias.

Takeaway: If you are training a large-scale Diffusion Transformer, don't just stack layers. Use long-range skip connections and explicitly penalize feature redundancy.

Limitations:

• While it works great for ImageNet, we haven't seen its performance on massive-scale uncurated web-crawled data (LAION-sized).
• The "Adaptive Weight" range $[0.1,0.5]$ seems somewhat empirical; a more robust theoretical derivation for this window would be beneficial.

DiverseDiT proves that Diversity is Efficiency. By forcing the network to work "harder" to learn distinct features, we get better images in a fraction of the time.