Research map / Image generation / Fast diffusion sampling
MC-DiT: Contextual Enhancement via Clean-to-Clean Reconstruction for Masked Diffusion Models
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 MC-DiT, a masked diffusion training paradigm that reconstructs unmasked clean patches from masked clean patches (plus two EMA decoder branches to avoid collapse), improving DiT image generation on ImageNet 256 and 512.
Problem
Masked-reconstruction training for diffusion transformers (DiT) reconstructs masked noisy patches from unmasked noisy patches, which under large noise weakens contextual information extraction and hurts semantic consistency and convergence.
Approach
Replace noisy-to-noisy mask reconstruction with clean-to-clean reconstruction (masked clean patches predict unmasked clean patches) at all noise levels, and stabilize training with two EMA decoder branches that enforce denoising and clean mask-reconstruction balance.
Key results
- MC-DiT-XL/2 reaches FID 4.14 on ImageNet 256x256 without guidance, versus 5.69 for MaskDiT-XL/2 and 9.62 for DiT-XL/2.
- With classifier-free guidance, MC-DiT-XL/2-G achieves FID 1.78 on ImageNet 256x256 and 2.03 on 512x512.
- Under matched 400K iterations for B/2, MC-DiT-B/2 reaches FID-50K 18.88 versus 28.62 for SD-DiT-B/2 and 43.47 for DiT-B/2.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace noisy-to-noisy masking with clean-to-clean contextual reconstruction | DiT-XL / 2 · ImageNet-256×256 class conditional | FID: 5.69 -> 4.14 | Helped |
| replace noisy-to-noisy masking with clean-to-clean contextual reconstruction | DiT-XL / 2 · ImageNet-512×512 class conditional | FID: 10.79 -> 9.30 | Helped |
| replace noisy-to-noisy masking with clean-to-clean contextual reconstruction | DiT-XL / 2 with classifier-free guidance · ImageNet-256×256 class conditional | FID: 1.79 -> 1.78 | Helped |
| replace noisy-to-noisy masking with clean-to-clean contextual reconstruction | DiT-XL / 2 with classifier-free guidance · ImageNet-512×512 class conditional | FID: 2.50 -> 2.03 | Helped |
| use MC-DiT training | DiT-S / 2 · ImageNet-256×256, 400K iterations | FID-50K: 48.39 -> 41.67 | Helped |
| use MC-DiT training | DiT-B / 2 · ImageNet-256×256, 400K iterations | FID-50K: 28.62 -> 18.88 | Helped |
| use MC-DiT training | DiT-B / 2 · CIFAR-10 | FID: 11.52 -> 9.28 | Helped |
| use MC-DiT training | DiT-B / 2 · CelebA | FID: 7.14 -> 5.36 | Helped |
| add the noisy EMA branch (ablation) | MC-DiT main branch · ImageNet-256×256, 400K iterations | FID: 22.10 -> 19.26 | Helped |
| add the clean EMA branch (ablation) | MC-DiT main branch · ImageNet-256×256, 400K iterations | FID: 22.10 -> 18.88 | Helped |
| include masked-patch noisy reconstruction alongside clean reconstruction (ablation) | MC-DiT main branch · ImageNet-256×256, 400K iterations | FID-50K: 25.88 -> 22.10 | Helped |
| remove unmasked tuning (ablation) | MC-DiT-XL / 2 · ImageNet-256×256, 1300K iterations | FID: 7.92 -> 8.33 | Hurt |
Surprising findings
- Removing unmasked finetuning increases FID only slightly (7.92 to 8.33), suggesting the finetuning stage is not essential.
Limitations
- Training and inference speed still need improvement, per the authors.
- No error bars or statistical significance reported.
- Code is not yet released; only a core file was provided in supplementary material.
Experimental setup
Datasets: ImageNet CIFAR-10 CelebA
Benchmarks: ImageNet 256x256 ImageNet 512x512 CIFAR-10 CelebA
Metrics: FID sFID IS Precision Recall
Baselines: DiT MaskDiT MDT SD-DiT ADM LDM U-ViT StyleGAN-XL MaskGIT
Builds on: DiT MaskDiT MDT SD-DiT EDM U-MAE
Closest papers by meaning
- SD-DiT: Unleashing the Power of Self-supervised Discrimination in Diffusion Transformer (CVPR 2024 · similarity 0.91)
- Not All Parameters Matter: Masking Diffusion Models for Enhancing Generation Ability (CVPR 2025 · similarity 0.89)
- Representation Entanglement for Generation: Training Diffusion Transformers Is Much Easier Than You Think (NeurIPS 2025 · similarity 0.89)
- Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think (ICLR 2025 · similarity 0.89)
- Simplified and Generalized Masked Diffusion for Discrete Data (NeurIPS 2024 · similarity 0.88)
- UDT: Reconciling U-Nets and Diffusion Transformers with Data-Adaptive Token Reduction (NeurIPS 2026 · similarity 0.88)
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.