# TinyFusion: Diffusion Transformers Learned Shallow

arXiv: https://arxiv.org/abs/2412.01199 · HTML: https://mlsanity.com/paper/2412.01199/

## Key facts

- Cited by 45 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 43 of its 55 references that are in the map.
- First posted on arXiv 2 Dec 2024.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 43)

- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Learning both Weights and Connections for Efficient Neural Networks](https://mlsanity.com/paper/1506.02626/) (2015)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Diffusion Models: A Comprehensive Survey of Methods and Applications](https://mlsanity.com/paper/2209.00796/) (2022)
- [MaskGIT: Masked Generative Image Transformer](https://mlsanity.com/paper/2202.04200/) (2022)
- [Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction](https://mlsanity.com/paper/2404.02905/) (2024)
- [PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis](https://mlsanity.com/paper/2310.00426/) (2023)
- [SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers](https://mlsanity.com/paper/2401.08740/) (2024)
- [Autoregressive Image Generation without Vector Quantization](https://mlsanity.com/paper/2406.11838/) (2024)
- [All are Worth Words: A ViT Backbone for Diffusion Models](https://mlsanity.com/paper/2209.12152/) (2022)
- [ShortGPT: Layers in Large Language Models are More Redundant Than You Expect](https://mlsanity.com/paper/2403.03853/) (2024)
- [SANA: Efficient High-Resolution Image Synthesis with Linear Diffusion Transformers](https://mlsanity.com/paper/2410.10629/) (2024)
- [Massive Activations in Large Language Models](https://mlsanity.com/paper/2402.17762/) (2024)
- [Q-Diffusion: Quantizing Diffusion Models](https://mlsanity.com/paper/2302.04304/) (2023)
- [Post-training Quantization on Diffusion Models](https://mlsanity.com/paper/2211.15736/) (2022)
- [SDXL-Lightning: Progressive Adversarial Diffusion Distillation](https://mlsanity.com/paper/2402.13929/) (2024)
- [Mixture-of-Depths: Dynamically allocating compute in transformer-based language models](https://mlsanity.com/paper/2404.02258/) (2024)
- [Depth-Adaptive Transformer](https://mlsanity.com/paper/1910.10073/) (2019)
- [Inpaint Anything: Segment Anything Meets Image Inpainting](https://mlsanity.com/paper/2304.06790/) (2023)
- [SnapFusion: Text-to-Image Diffusion Model on Mobile Devices within Two Seconds](https://mlsanity.com/paper/2306.00980/) (2023)
- [Structural Pruning for Diffusion Models](https://mlsanity.com/paper/2305.10924/) (2023)
- [PTQD: Accurate Post-Training Quantization for Diffusion Models](https://mlsanity.com/paper/2305.10657/) (2023)

## Built on by (most cited first, 30 of 45)

- [No Other Representation Component Is Needed: Diffusion Transformers Can Provide Representation Guidance by Themselves](https://mlsanity.com/paper/2505.02831/) (2025)
- [SepPrune: Structured Pruning for Efficient Deep Speech Separation](https://mlsanity.com/paper/2505.12079/) (2025)
- [Memory-Efficient Visual Autoregressive Modeling with Scale-Aware KV Cache Compression](https://mlsanity.com/paper/2505.19602/) (2025)
- [D-OPSD: On-Policy Self-Distillation for Continuously Tuning Step-Distilled Diffusion Models](https://mlsanity.com/paper/2605.05204/) (2026)
- [OBS-Diff: Accurate Pruning For Diffusion Models in One-Shot](https://mlsanity.com/paper/2510.06751/) (2025)
- [Dense2MoE: Restructuring Diffusion Transformer to MoE for Efficient Text-to-Image Generation](https://mlsanity.com/paper/2510.09094/) (2025)
- [Video-As-Prompt: Unified Semantic Control for Video Generation](https://mlsanity.com/paper/2510.20888/) (2025)
- [On-Device Diffusion Transformer Policy for Efficient Robot Manipulation](https://mlsanity.com/paper/2508.00697/) (2025)
- [Exploring Diffusion Transformer Designs via Grafting](https://mlsanity.com/paper/2506.05340/) (2025)
- [Neodragon: Mobile Video Generation using Diffusion Transformer](https://mlsanity.com/paper/2511.06055/) (2025)
- [Training-free Mixed-Resolution Latent Upsampling for Spatially Accelerated Diffusion Transformers](https://mlsanity.com/paper/2507.08422/) (2025)
- [Taming Diffusion Transformer for Efficient Mobile Video Generation in Seconds](https://mlsanity.com/paper/2507.13343/) (2025)
- [Shallow-π: Knowledge Distillation for Flow-based VLAs](https://mlsanity.com/paper/2601.20262/) (2026)
- [Elastic Diffusion Transformer](https://mlsanity.com/paper/2602.13993/) (2026)
- [TinySR: Pruning Diffusion for Real-World Image Super-Resolution](https://mlsanity.com/paper/2508.17434/) (2025)
- [Massive Activations are the Key to Local Detail Synthesis in Diffusion Transformers](https://mlsanity.com/paper/2510.11538/) (2025)
- [Pluggable Pruning with Contiguous Layer Distillation for Diffusion Transformers](https://mlsanity.com/paper/2511.16156/) (2025)
- [S2DiT: Sandwich Diffusion Transformer for Mobile Streaming Video Generation](https://mlsanity.com/paper/2601.12719/) (2026)
- [ALTER: All-in-One Layer Pruning and Temporal Expert Routing for Efficient Diffusion Generation](https://mlsanity.com/paper/2505.21817/) (2025)
- [Test-Time Iterative Error Correction for Efficient Diffusion Models](https://mlsanity.com/paper/2511.06250/) (2025)
- [DDiT: Dynamic Patch Scheduling for Efficient Diffusion Transformers](https://mlsanity.com/paper/2602.16968/) (2026)
- [MobileWan: Closing the Quality Gap for Mobile Video Diffusion](https://mlsanity.com/paper/2607.06173/) (2026)
- [Efficient High-Resolution Image Editing with Hallucination-Aware Loss and Adaptive Tiling](https://mlsanity.com/paper/2510.06295/) (2025)
- [Phase-Aligned RoPE for Mixed-Resolution Diffusion Transformer](https://mlsanity.com/paper/2511.19778/) (2025)
- [Steering Video Diffusion Transformers with Massive Activations](https://mlsanity.com/paper/2603.17825/) (2026)
- [FREE: Uncertainty-Aware Autoregression for Parallel Diffusion Transformers](https://mlsanity.com/paper/2511.20390/) (2025)
- [Diff-ES: Stage-wise Structural Diffusion Pruning via Evolutionary Search](https://mlsanity.com/paper/2603.05105/) (2026)
- [TRaM-VSR: Importance-Aware Token Routing and Merging for One-Step Diffusion Video Super-Resolution](https://mlsanity.com/paper/2607.22231/) (2026)
- [HPSD: Hybrid-Policy Self-Distillation for Text-Image-to-Video Diffusion Models](https://mlsanity.com/paper/2608.13205/) (2026)
- [Designing Parameter and Compute Efficient Diffusion Transformers using Distillation](https://mlsanity.com/paper/2502.14226/) (2025)

## Abstract (arXiv)

Diffusion Transformers have demonstrated remarkable capabilities in image generation but often come with excessive parameterization, resulting in considerable inference overhead in real-world applications. In this work, we present TinyFusion, a depth pruning method designed to remove redundant layers from diffusion transformers via end-to-end learning. The core principle of our approach is to create a pruned model with high recoverability, allowing it to regain strong performance after fine-tuning. To accomplish this, we introduce a differentiable sampling technique to make pruning learnable, paired with a co-optimized parameter to simulate future fine-tuning. While prior works focus on minimizing loss or error after pruning, our method explicitly models and optimizes the post-fine-tuning performance of pruned models. Experimental results indicate that this learnable paradigm offers substantial benefits for layer pruning of diffusion transformers, surpassing existing importance-based and error-based methods. Additionally, TinyFusion exhibits strong generalization across diverse architectures, such as DiTs, MARs, and SiTs. Experiments with DiT-XL show that TinyFusion can craft a shallow diffusion transformer at less than 7% of the pre-training cost, achieving a 2× speedup with an FID score of 2.86, outperforming competitors with comparable efficiency. Code is available at https://github.com/VainF/TinyFusion.