DiverseDiT: Towards Diverse Representation Learning in Diffusion Transformers

AI summary

Authors: Mengping Yang, Zhiyu Tan, Binglei Li, Xiaomeng Yang, Hesen Chen, Hao Li Affiliations: Fudan University; Shanghai Academy of AI for Science; Shanghai Innovation Institute arXiv: 2603.04239v2 [cs.CV], 15 Jun 2026 Code: https://github.com/kobeshegu/DiverseDiT

Overview

Research area: Computer vision, generative modeling, and representation learning in Diffusion Transformers (DiTs).

Technical level: Advanced. The paper assumes familiarity with diffusion/flow-matching models, transformer block architectures, representation alignment methods (REPA), and similarity metrics such as Centered Kernel Alignment (CKA).

Scope: The paper systematically analyzes how internal representations evolve inside DiT blocks, argues that cross-block representation diversity drives effective learning, and proposes the DiverseDiT framework that promotes such diversity without external pretrained encoders, validated on ImageNet 256×256 and 512×512 across multiple model scales and both multi-step and one-step generation.

What This Paper Is About

Diffusion Transformers are strong image generators, and recent methods like REPA improve them by aligning internal features with a pre-trained external vision encoder — but nobody had explained why such alignment helps. The authors first measure how representations change across DiT blocks during training, and find that the key factor is how different each block's representation is from the others. Based on that insight, they build DiverseDiT, which deliberately encourages blocks to learn distinct, complementary features using only internal mechanisms rather than external guidance models.

Key Contributions

1. A systematic analysis of representation dynamics in DiTs. The authors use CKA to track representation similarity across blocks during training, and study how single-block alignment, multi-block alignment, and multi-encoder alignment affect those similarities. They state this is the first work to elucidate this representation relationship and explain why external alignment techniques work — effective representation learning corresponds to increased representation diversity across blocks.

2. The DiverseDiT framework. Two components are introduced: long residual connections that inject earlier blocks' outputs into later blocks to diversify each block's input, and a representation diversity loss that penalizes similarity between block outputs.

3. A composite representation diversity loss with an adaptive weighting scheme. The loss combines an orthogonality loss, a mutual-information minimization proxy, and a feature dispersion loss, with weights that adapt based on the loss value to avoid divergence.

4. Broad empirical validation. DiverseDiT is applied to SiT, REPA, and MeanFlow backbones at B/2, L/2, and XL/2 scales on ImageNet 256×256 and 512×512, in multi-step and one-step settings, and is shown to be complementary to DispLoss and SRA.

Main Findings

• Representation diversity across blocks naturally increases during training. CKA heatmaps for SiT at training steps 5K, 50K, 200K, and 450K become progressively more diagonal, meaning different blocks become less similar to one another.

• External alignment increases the target block's dissimilarity from other blocks. In the REPA heatmaps, regions around the aligned block show less similarity than the corresponding regions in SiT. Aligning earlier blocks (Block 5, Block 8) performs better than aligning later blocks (Block 10), consistent with the original REPA findings.

• Aligning more blocks or using more encoders does not necessarily help. Using Block [2,5,8] or [3,6,9] for guidance did not produce comparable improvements, and in the Block [2,5,8] case FID was slightly worse. Aligning different encoders on different blocks ([5/DINOv2 + 10/MAE]) also gave limited benefit. The authors hypothesize conflicting constraints prevent individual blocks from specializing.

• Consistent gains across model scales without CFG (ImageNet 256×256, 400K iterations). SiT-B FID improved from 36.80 to 28.05; REPA-B from 22.99 to 17.29; SiT-L from 18.77 to 16.10; REPA-L from 9.57 to 8.47; SiT-XL from 17.43 to 12.42; REPA-XL from 8.73 to 8.09. Inception Score improved in every case, for example REPA-B from 64.73 to 79.92.

• Smaller backbone plus DiverseDiT can beat a larger baseline. With 400K iterations, REPA-B with the method reaches FID 17.29, better than SiT-L at 18.77 for the same iterations. REPA-L with the method (FID 8.47) outperforms REPA-XL (FID 8.73), with IS 123.03 versus 118.68.

• Competitive SoTA results with far fewer epochs (ImageNet 256×256, with CFG). DiverseDiT reaches FID 1.89 at 80 epochs and 1.52 at 200 epochs, versus REPA at 1.96 (200 epochs) and 1.80 (800 epochs), SiT-XL/2 at 2.06 (1400 epochs), and DiT-XL/2 at 2.27 (1400 epochs). REG achieves a slightly better FID of 1.36 but requires 800 epochs, which the authors describe as four times the training cost of theirs.

• Strong ImageNet 512×512 results with CFG. DiverseDiT obtains FID 2.21 at 80 epochs and 1.99 at 200 epochs, compared with REPA at 2.08 (200 epochs), SiT-XL/2 at 2.62 (600 epochs), and DiT-XL/2 at 3.04 (600 epochs).

• Improvements carry over to one-step generation without CFG. Applied to MeanFlow, FID improves from 9.44 to 8.51 for MF-B/2, from 8.73 to 7.15 for MF-L/2, and from 5.94 to 5.69 for MF-XL/2; IS for MF-B/2 rises from 152.55 to 158.84 and for MF-L/2 from 161.69 to 199.66.

• New state of the art for one-step generation with CFG. MeanFlow-XL/2 with the method reaches FID 2.99 with 713M parameters and 1 step, versus MeanFlow-XL/2 at 3.43 (676M), MeanFlow-XL/2 + DispLoss at 3.21 (676M), and earlier one-step methods such as IMM-XL/2 at 7.77, Shortcut-XL/2 at 10.60, and iCT-XL/2 at 34.24.

• Both components matter (ablation). Removing the diversity loss raises FID from 28.05 to 32.77 on SiT-B and from 17.29 to 20.66 on REPA-B. Removing the long residual connections raises FID to 33.72 on SiT-B and to 18.18 on REPA-B.

• All three loss terms contribute. On REPA-B, the full combination yields FID 17.29 and IS 79.92; using only the orthogonality loss gives 18.97, only the mutual-information loss gives 17.70, and only the dispersion loss gives 20.85. Adding any single component still outperforms the REPA-B baseline.

• The adaptive loss range matters, and a constant weight diverges. The range [0.1, 0.5] gives the best result (FID 28.05), while [0.2, 0.7] gives 30.59 and [0.3, 0.9] gives 31.85. A constant weight leads to divergence, which the authors attribute to excessive representation discrepancy.

• Complementary to prior internal methods (400K iterations, no CFG). Starting from SiT-B at FID 36.80, adding DiverseDiT gives 28.05, adding DispLoss on top gives 24.98, and further adding SRA gives 21.95 — better than REPA's 22.99 at the same iteration count, without any external guidance.

Methodology in Plain English

The authors first run a diagnostic study rather than immediately proposing a method. They take SiT-B/2 and REPA-B/2 trained on ImageNet 256×256 for 450K iterations and compute CKA similarity between the hidden representations of different transformer blocks. CKA is a standard measure of how similar two sets of neural representations are, so a low value between two blocks means they encode different things. They visualize these similarities as heatmaps at checkpoints and repeat the analysis for models that align one block, several blocks, or several blocks with different pre-trained encoders (DINOv2-B, MAE-L, MoCov3).

The pattern they find — diversity across blocks grows during training, and alignment helps mainly by making the aligned block more distinct — leads to a two-part intervention.

First, instead of feeding each block only the output of the block immediately before it, they add long residual connections that concatenate the output of an earlier block with a later block's input, then normalize and linearly project the result. This makes each block's input less homogeneous.

Second, they add a representation diversity loss over pairs of blocks. It has three terms: an orthogonality term that penalizes high cosine similarity between per-block mean feature vectors; a mutual-information proxy that penalizes token-wise cosine similarity between normalized features of different blocks, avoiding the cost of full covariance computations; and a dispersion term that maximizes the variance of averaged per-channel activations so different channels get used. The three terms are weighted equally at 0.33 each by default, but a scalar weight on the total loss is scheduled by its own value — full weight above 0.5, ramped down between 0.1 and 0.5, and zero otherwise. This prevents the loss from being driven so low that the model diverges.

Evaluation follows the SiT, REPA, and MeanFlow training recipes with untouched other details, using the ImageNet dataset at 256×256 and 512×512, Stable Diffusion VAE features, ADM-style evaluation, and standard metrics (FID, sFID, IS, Precision, Recall) computed from 50K generated images. Training uses AdamW with a constant learning rate of 1e-4, betas (0.9, 0.999), no decay, and a batch size of 256, on 8× 80GB H800 GPUs. Sampling uses an Euler-Maruyama sampler with 250 steps; CFG is off unless stated.

Why This Matters

Impact on research. The paper reframes representation learning in diffusion transformers as a diversity problem rather than purely an alignment problem. It offers a mechanistic explanation for why REPA-style external alignment works, why aligning more blocks or bigger encoders yields only marginal gains, and it shows that competitive or better results can be obtained without any external foundation model. This gives the field a concrete, testable design principle — increase cross-block discrepancy — rather than a heuristic.

Real-world applications:

• High-resolution image synthesis systems where training cost is the bottleneck, since the method reaches competitive FID with substantially fewer epochs.
• One-step and few-step image generation, relevant to interactive and latency-sensitive creative tools, where the method sets a new reported one-step FID of 2.99.
• Generative pipelines that cannot afford large external encoders for licensing, memory, or deployment reasons, since the approach needs no external guidance model.
• Existing diffusion training stacks, because the method is complementary and was combined with REPA, DispLoss, and SRA to yield further gains.

Industry relevance. Training large diffusion transformers is expensive; a technique that both accelerates convergence and improves final quality across model scales, and that is shown to drop into three different backbones, is directly relevant to teams training image generation models. The one-step results matter for inference cost, and the method's independence from external encoders simplifies data and model licensing considerations.

Future Directions

• Applying the framework to larger and broader settings. The paper evaluates B/2, L/2, and XL/2 on ImageNet; the same diversity-driven design is untested here on text-to-image or text-to-video generation, or on substantially larger backbones.

• Understanding the diversity ceiling. The authors observe that pushing the diversity loss too low causes divergence, and that aligning too many blocks hurts performance. What determines the optimal amount of cross-block diversity, and how it should be scheduled, remains open.

• Reducing the cost of the diversity loss. The paper already limits computation by considering only a subset of block pairs and by using a mutual-information proxy instead of covariance matrices; further efficiency work is a natural extension.

• Comparing against stronger external-alignment regimes at equal compute. The method is shown to be complementary to REPA, DispLoss, and SRA, so the natural next question is how far the internal-only and hybrid variants can be pushed under matched training budgets.

• The paper states that limitations and potential future works are discussed in Appendix H; those details are not included in the provided content.

Target Audience

Researchers and engineers working on diffusion and flow-matching generative models, particularly those interested in representation learning, transformer scaling, and training efficiency. It is also valuable for practitioners who train image generation models at scale and need methods that fit existing pipelines. Readers should be comfortable with transformer architectures, diffusion/flow objectives, and representation similarity metrics; the paper's diagnostic analysis uses CKA and its baselines include REPA, SRA, and MeanFlow, so prior familiarity with that literature helps considerably.

Authors’ 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.