Research map / Architectures and efficiency / Attention and transformers
Exploring Diffusion Transformer Designs via Grafting
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 grafting, a two-stage method (activation distillation then lightweight finetuning) that edits pretrained diffusion transformers to materialize new operator and depth/width designs under small compute budgets.
Problem
Evaluating architecture design choices for diffusion transformers requires costly pretraining from scratch, which limits architectural investigation.
Approach
Edit pretrained DiTs by replacing operators (MHA with local gated convolutions, sliding window or linear attention; MLPs with variable expansion ratio or convolutional variants) and restructuring blocks (parallelizing sequential block pairs). Each new operator is initialized by regressing pretrained activations (L1 for MHA, L2 for MLP), then the whole model is lightly finetuned on a small data sub
Key results
- Many interleaved hybrid DiT-XL/2 designs reach FID 2.38-2.64 vs 2.27 baseline using under 2% pretraining compute.
- Grafting PixArt-Sigma with Hyena-X on 50% of layers gives 1.43x latency speedup with GenEval 47.78 vs 49.75.
- Depth-halved DiT-XL/2 (28 to 14 layers via parallel block pairs) reaches FID 2.77, better than DiT-L/2 (3.73) at depth 24.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace 50% of MHA with local alternatives using interleaved grafting | DiT-XL/2 · ImageNet-1K 256×256 | FID: 2.27 -> 2.65–2.74 | Mixed |
| Increase interleaved MHA replacement from 50% to 75% (ablation) | DiT-XL/2 with grafted MHA alternatives · ImageNet-1K 256×256 | FID: 2.74 -> 3.69 | Hurt |
| Increase interleaved MHA replacement to 100% (ablation) | DiT-XL/2 with grafted MHA alternatives · ImageNet-1K 256×256 | FID: all FIDs > 75 | Hurt |
| Increase fine-tuning data from 10% to 20% (ablation) | DiT-XL/2 with 50% MHA replacement · ImageNet-1K 256×256 | FID: 2.74 -> 2.61; 2.67 -> 2.62; 2.65 -> 2.55 | Helped |
| Use interleaved layer selection instead of targeted layer heuristics (ablation) | DiT-XL/2 with 50% Hyena-X MHA replacement · ImageNet-1K 256×256 | FID: 2.74; top-local 3.02; low-local 3.18; deep 4.00 | Helped |
| Increase MLP expansion ratio from 3 to 6 (ablation) | DiT-XL/2 with 50% MLP replacement · ImageNet-1K 256×256 | FID: 2.53 -> 2.38 | Helped |
| Increase MLP replacement ratio for Hyena-X beyond 50% (ablation) | DiT-XL/2 · ImageNet-1K 256×256 | FID: 2.64 -> 3.26 | Hurt |
| Use L1 rather than L2 regression for MHA activation distillation (ablation) | DiT-XL/2 MHA initialization · ImageNet-1K 256×256 | FID: 2.58 -> 2.51 | Helped |
| Use L2 rather than L1 regression for MLP activation distillation (ablation) | DiT-XL/2 MLP initialization · ImageNet-1K 256×256 | FID: 2.83 -> 2.33 | Helped |
| Increase full self-grafting fine-tuning data from 0.63% to 10.0% (ablation) | DiT-XL/2 MHA · ImageNet-1K 256×256 | FID: 16.78 -> 2.49 | Helped |
| Use modulation-aware rather than standard regression targets with L2 loss (ablation) | DiT-XL/2 MHA initialization · ImageNet-1K 256×256 | FID: 2.58 -> 3.00 | Hurt |
| Use modulation-aware rather than standard regression targets with L1 loss (ablation) | DiT-XL/2 MHA initialization · ImageNet-1K 256×256 | FID: 2.51 -> 2.51 | No clear effect |
Surprising findings
- Replacing all MHA layers with sliding window attention degrades FID from 2.27 to 53.9, while 50% interleaved replacement stays within 0.5 FID of baseline.
- MLP variable expansion ratio with r=6 at 50% replacement improves FID to 2.38 over the baseline 2.27-level design space.
Limitations
- Only DiT-based latent diffusion models studied; generalization to other families such as autoregressive models untested.
- Synthetic PixArt-Sigma data may propagate artifacts and biases.
- Unclear whether architectures good under grafting perform well when trained from scratch.
Experimental setup
Datasets: ImageNet-1K synthetic image-text pairs generated by PixArt-Sigma
Benchmarks: ImageNet-1K 256x256 GenEval
Metrics: FID IS sFID Precision Recall GenEval overall score latency FLOPs parameters
Baselines: DiT-XL/2 PixArt-Sigma DiT-L/2 U-ViT-L DiT-B/2 BK-SDM TinyDiT-D14 TinyDiT-D14 w/ MKD
Builds on: DiT Hyena Mamba-2 PixArt-Sigma
Closest papers by meaning
- FasterDiT: Towards Faster Diffusion Transformers Training without Architecture Modification (NeurIPS 2024 · similarity 0.89)
- EDT: An Efficient Diffusion Transformer Framework Inspired by Human-like Sketching (NeurIPS 2024 · similarity 0.89)
- TinyFusion: Diffusion Transformers Learned Shallow (CVPR 2025 · similarity 0.88)
- Rethinking Cross-Layer Information Routing in Diffusion Transformer (NeurIPS 2026 · similarity 0.88)
- U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers (NeurIPS 2024 · similarity 0.88)
- DiT4Edit: Diffusion Transformer for Image Editing (AAAI 2025 · 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.