Research map / Image generation / Autoregressive and latent generation
SD-DiT: Unleashing the Power of Self-supervised Discrimination in Diffusion Transformer
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes SD-DiT, a teacher-student DiT training scheme that adds a DINO-style inter-image discriminative loss on PF-ODE noise pairs, decouples encoder and decoder, and reaches better FID than DiT and MaskDiT on ImageNet 256 with faster convergence.
Problem
DiT training converges slowly and mask-based DiT training suffers from a training-inference discrepancy and fuzzy coupling between mask reconstruction and the generative diffusion objective.
Approach
Replace the mask reconstruction loss with a self-distillation discriminative loss between a student view (regular EDM noise, masked) and a teacher view (EMA encoder, fixed minimum noise), while decoupling the encoder (discriminative + generative) from a decoder trained only on the generative loss.
Key results
- SD-DiT-XL/2 reaches FID 9.01 at 1300k steps, versus 9.62 for DiT-XL/2 at 7000k steps.
- SD-DiT-S/2 achieves FID 48.39 at 400k steps, compared with 68.40 for DiT-S/2 and 53.46 for MDT-S/2.
- Removing the discriminative objective raises FID from 53.72 to 62.84 in the S/2 ablation.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use SD-DiT instead of MDT | DiT-XL/2 · ImageNet 256×256 | FID-50K: 7.41 -> 7.21 | Helped |
| Use SD-DiT instead of DiT | DiT-XL/2 · ImageNet 256×256 | FID-50K: 9.62 -> 7.21 | Helped |
| Use SD-DiT instead of MDT | DiT-S/2 · ImageNet 256×256, 400k steps | FID-50K: 53.46 -> 48.39 | Helped |
| Use SD-DiT instead of MDT | DiT-B/2 · ImageNet 256×256, 400k steps | FID-50K: 34.33 -> 28.62 | Helped |
| Add the discriminative objective and teacher branch (ablation) | SD-DiT-S/2 · ImageNet 256×256, 400k steps | FID: 62.84 -> 53.72 | Helped |
| Add the mask strategy (ablation) | SD-DiT-S/2 · ImageNet 256×256, 400k steps | FID: 58.92 -> 53.72 | Helped |
| Vary the mask ratio from 0 to 1; 20% gives the best reported result (ablation) | SD-DiT-S/2 · ImageNet 256×256, 400k steps | FID: 48.4 (20% mask ratio); 61.0 (90% mask ratio) | Mixed |
| Use fixed minimum teacher noise instead of teacher noise from the student-noise distribution (ablation) | SD-DiT-S/2 · ImageNet 256×256, 400k steps | FID: 63.4 -> 53.7 | Helped |
| Set teacher noise to 0.5 rather than the fixed minimum noise (ablation) | SD-DiT-S/2 · ImageNet 256×256, 400k steps | FID: 53.7 -> 66.2 | Hurt |
Surprising findings
- Teacher noise drawn from the same distribution as the student noise yields worse FID (63.4) than the fixed minimum-noise teacher view, consistent with the InfoMin principle.
- MaskDiT with a 75% mask ratio degrades severely (FID 121.16 for MaskDiT-XL/2), while SD-DiT remains stable at high mask ratios.
Limitations
- Mixed precision led to NaN losses, so full precision is used for larger backbones, adding cost.
- Teacher noise scale is sensitive; noise values near the student noise distribution degrade FID sharply.
- Main results for XL/2 use a batch size of 256 and do not match MaskDiT's batch size of 1024 in the reported table.
Experimental setup
Datasets: ImageNet-1K
Benchmarks: ImageNet 256x256
Metrics: FID-50K sFID Inception Score Precision Recall
Baselines: DiT MDT MaskDiT ADM LDM MaskGIT BigGAN-deep StyleGAN
Builds on: DiT MaskDiT MDT DINO iBOT EDM MAE
Closest papers by meaning
- MC-DiT: Contextual Enhancement via Clean-to-Clean Reconstruction for Masked Diffusion Models (NeurIPS 2024 · similarity 0.91)
- Faster Diffusion: Rethinking the Role of the Encoder for Diffusion Model Inference (NeurIPS 2024 · similarity 0.90)
- Layer- and Timestep-Adaptive Differentiable Token Compression Ratios for Efficient Diffusion Transformers (CVPR 2025 · similarity 0.90)
- DiCo: Revitalizing ConvNets for Scalable and Efficient Diffusion Modeling (NeurIPS 2025 · similarity 0.90)
- Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think (ICLR 2025 · similarity 0.90)
- UDT: Reconciling U-Nets and Diffusion Transformers with Data-Adaptive Token Reduction (NeurIPS 2026 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.