Research map / Architectures and efficiency / Pruning and efficient tuning
TinyFusion: Diffusion Transformers Learned Shallow
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 TinyFusion, a learnable depth-pruning method for diffusion transformers that uses differentiable N:M layer-mask sampling jointly optimized with LoRA weight updates to select layers with high post-fine-tuning recoverability.
Problem
Depth pruning of pre-trained diffusion transformers usually relies on heuristic or loss-minimizing criteria, which do not guarantee strong performance after fine-tuning, and the pruning-plus-finetuning objective is non-differentiable and expensive to search.
Approach
Model pruning masks as learnable categorical distributions over local N:M blocks, sampled with Gumbel-Softmax/STE, and jointly optimize them with a co-optimized LoRA update that simulates fine-tuning, so the search favors masks with high recoverability rather than low immediate loss.
Key results
- TinyDiT-D14 (28 to 14 layers, distilled with Masked KD, 500K steps) reaches FID 2.86 versus 2.27 for DiT-XL/2 with 2x speedup, at under 7% of pre-training cost.
- With 100K fine-tuning steps and no KD, the learnable method gets FID 5.73, versus 22.28 for ShortGPT and 21.15 for Sensitivity Analysis at 50% depth.
- TinyMAR-D16 achieves FID 2.28 with 40 epochs, beating MAR-Base (24 blocks, FID 2.31 at 400 epochs).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Prune DiT-XL/2 from 28 to 14 layers with TinyFusion rather than using the reported pruning baselines | DiT-XL/2 · ImageNet 256×256 | FID: 22.28 -> 5.73 | Helped |
| Add MaskedKD during TinyDiT-D14 recovery (ablation) | TinyDiT-D14 · ImageNet 256×256 | FID: 5.73 -> 3.73 | Helped |
| Continue MaskedKD fine-tuning from 100K to 500K steps (ablation) | TinyDiT-D14 (KD) · ImageNet 256×256 | FID: 3.73 -> 2.86 | Helped |
| Prune MAR-Large to 16 blocks using TinyFusion | MAR-Large · ImageNet 256×256 | FID: 1.78 -> 2.28 | No clear effect |
| Prune SiT-XL/2 from 28 to 14 blocks | SiT-XL/2 · ImageNet 256×256 | FID: 2.06 -> 3.02 | Hurt |
| Use a 2:4 local pruning pattern instead of 1:2 with LoRA recoverability updates (ablation) | TinyFusion · ImageNet 256×256, 10K fine-tuning iterations | FID: 33.39 -> 34.21 | Hurt |
| Use a 7:14 local pruning pattern instead of 1:2 with LoRA recoverability updates (ablation) | TinyFusion · ImageNet 256×256, 10K fine-tuning iterations | FID: 33.39 -> 49.41 | Hurt |
| Use full fine-tuning rather than LoRA for 1:2 recoverability estimation (ablation) | TinyFusion · ImageNet 256×256, 10K fine-tuning iterations | FID: 33.39 -> 35.77 | Hurt |
| Freeze weights rather than use LoRA for 1:2 recoverability estimation (ablation) | TinyFusion · ImageNet 256×256, 10K fine-tuning iterations | FID: 33.39 -> 39.56 | Hurt |
| Add vanilla logits knowledge distillation during recovery (ablation) | TinyDiT-D14 · ImageNet 256×256, 100K steps | FID: 5.79 -> 4.66 | Helped |
| Distill representations without masking massive activations (ablation) | TinyDiT-D14 · ImageNet 256×256, 100K steps | FID | Hurt |
| Use a 0.1σ activation threshold for masked representation distillation (ablation) | TinyDiT-D14 · ImageNet 256×256, 100K steps | FID | Hurt |
Surprising findings
- Pruned models with the lowest calibration loss yield poor post-fine-tuning FID (20.69 or NaN), so loss minimization is not the primary determinant of recoverability for diffusion transformers.
- Naive representation distillation on DiT hidden states diverges (NaN) due to massive activations; masking them fixes training.
Limitations
- Relatively high initial calibration loss of the learned pruned model compared with other baselines.
- Layer-pruning search requires LoRA-based recoverability estimation over the dataset, which adds training cost versus purely metric-based pruning.
Experimental setup
Datasets: ImageNet
Benchmarks: ImageNet 256x256 class-conditional generation
Metrics: FID sFID IS Precision Recall sampling speed (it/s)
Baselines: ShortGPT Flux-Lite Sensitivity Analysis BK-SDM (Oracle) Diff-Pruning U-ViT U-DiT DTR MAR-Base
Builds on: Gumbel-Softmax LoRA ShortGPT BK-SDM Diff-Pruning
Closest papers by meaning
- FasterDiT: Towards Faster Diffusion Transformers Training without Architecture Modification (NeurIPS 2024 · similarity 0.90)
- SparseDiT: Token Sparsification for Efficient Diffusion Transformer (NeurIPS 2025 · similarity 0.90)
- LazyDiT: Lazy Learning for the Acceleration of Diffusion Transformers (AAAI 2025 · similarity 0.90)
- HarmoniCa: Harmonizing Training and Inference for Better Feature Caching in Diffusion Transformer Acceleration (ICML 2025 · similarity 0.90)
- Layer- and Timestep-Adaptive Differentiable Token Compression Ratios for Efficient Diffusion Transformers (CVPR 2025 · similarity 0.89)
- DiP-GO: A Diffusion Pruner via Few-step Gradient Optimization (NeurIPS 2024 · 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.