# U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers

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

## Key facts

- Cited by 54 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 37 of its 46 references that are in the map.
- First posted on arXiv 4 May 2024.

## Placed in

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

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [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)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions](https://mlsanity.com/paper/2102.12122/) (2021)
- [SwinIR: Image Restoration Using Swin Transformer](https://mlsanity.com/paper/2108.10257/) (2021)
- [Restormer: Efficient Transformer for High-Resolution Image Restoration](https://mlsanity.com/paper/2111.09881/) (2021)
- [Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Transformers](https://mlsanity.com/paper/2012.15840/) (2020)
- [Swin Transformer V2: Scaling Up Capacity and Resolution](https://mlsanity.com/paper/2111.09883/) (2021)
- [Pre-Trained Image Processing Transformer](https://mlsanity.com/paper/2012.00364/) (2020)
- [Uformer: A General U-Shaped Transformer for Image Restoration](https://mlsanity.com/paper/2106.03106/) (2021)
- [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)
- [RepVGG: Making VGG-style ConvNets Great Again](https://mlsanity.com/paper/2101.03697/) (2021)
- [All are Worth Words: A ViT Backbone for Diffusion Models](https://mlsanity.com/paper/2209.12152/) (2022)
- [Simple diffusion: End-to-end diffusion for high resolution images](https://mlsanity.com/paper/2301.11093/) (2023)
- [Rethinking Spatial Dimensions of Vision Transformers](https://mlsanity.com/paper/2103.16302/) (2021)
- [PixArt-Sigma: Weak-to-Strong Training of Diffusion Transformer for 4K Text-to-Image Generation](https://mlsanity.com/paper/2403.04692/) (2024)
- [CMT: Convolutional Neural Networks Meet Vision Transformers](https://mlsanity.com/paper/2107.06263/) (2021)
- [FreeU: Free Lunch in Diffusion U-Net](https://mlsanity.com/paper/2309.11497/) (2023)
- [Efficient and Explicit Modelling of Image Hierarchies for Image Restoration](https://mlsanity.com/paper/2303.00748/) (2023)
- [Scalable High-Resolution Pixel-Space Image Synthesis with Hourglass Diffusion Transformers](https://mlsanity.com/paper/2401.11605/) (2024)
- [DiffiT: Diffusion Vision Transformers for Image Generation](https://mlsanity.com/paper/2312.02139/) (2023)

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

- [SANA: Efficient High-Resolution Image Synthesis with Linear Diffusion Transformers](https://mlsanity.com/paper/2410.10629/) (2024)
- [Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models](https://mlsanity.com/paper/2410.11081/) (2024)
- [Identity-Preserving Text-to-Video Generation by Frequency Decomposition](https://mlsanity.com/paper/2411.17440/) (2024)
- [Playground v3: Improving Text-to-Image Alignment with Deep-Fusion Large Language Models](https://mlsanity.com/paper/2409.10695/) (2024)
- [Inductive Moment Matching](https://mlsanity.com/paper/2503.07565/) (2025)
- [SANA-Sprint: One-Step Diffusion with Continuous-Time Consistency Distillation](https://mlsanity.com/paper/2503.09641/) (2025)
- [Align Your Flow: Scaling Continuous-Time Flow Map Distillation](https://mlsanity.com/paper/2506.14603/) (2025)
- [DeCo: Frequency-Decoupled Pixel Diffusion for End-to-End Image Generation](https://mlsanity.com/paper/2511.19365/) (2025)
- [Adversarial Score identity Distillation: Rapidly Surpassing the Teacher in One Step](https://mlsanity.com/paper/2410.14919/) (2024)
- [TinyFusion: Diffusion Transformers Learned Shallow](https://mlsanity.com/paper/2412.01199/) (2024)
- [U-REPA: Aligning Diffusion U-Nets to ViTs](https://mlsanity.com/paper/2503.18414/) (2025)
- [PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations](https://mlsanity.com/paper/2505.24717/) (2025)
- [ZipVoice: Fast and High-Quality Zero-Shot Text-to-Speech with Flow Matching](https://mlsanity.com/paper/2506.13053/) (2025)
- [LaVin-DiT: Large Vision Diffusion Transformer](https://mlsanity.com/paper/2411.11505/) (2024)
- [Selective Underfitting in Diffusion Models](https://mlsanity.com/paper/2510.01378/) (2025)
- [DiC: Rethinking Conv3x3 Designs in Diffusion Models](https://mlsanity.com/paper/2501.00603/) (2024)
- [Ingredients: Blending Custom Photos with Video Diffusion Transformers](https://mlsanity.com/paper/2501.01790/) (2025)
- [Towards Stabilized and Efficient Diffusion Transformers through Long-Skip-Connections with Spectral Constraints](https://mlsanity.com/paper/2411.17616/) (2024)
- [Learning to Integrate Diffusion ODEs by Averaging the Derivatives](https://mlsanity.com/paper/2505.14502/) (2025)
- [PixelU: A U-Shaped Transformer for Efficient End-to-End Pixel Diffusion](https://mlsanity.com/paper/2606.27760/) (2026)
- [RainFusion: Adaptive Video Generation Acceleration via Multi-Dimensional Visual Redundancy](https://mlsanity.com/paper/2505.21036/) (2025)
- [Spiral RoPE: Rotate Your Rotary Positional Embeddings in the 2D Plane](https://mlsanity.com/paper/2602.03227/) (2026)
- [DiCo: Revitalizing ConvNets for Scalable and Efficient Diffusion Modeling](https://mlsanity.com/paper/2505.11196/) (2025)
- [Scaling Diffusion Transformers Efficiently via μP](https://mlsanity.com/paper/2505.15270/) (2025)
- [Improved Training Technique for Shortcut Models](https://mlsanity.com/paper/2510.21250/) (2025)
- [FSVideo: Fast Speed Video Diffusion Model in a Highly-Compressed Latent Space](https://mlsanity.com/paper/2602.02092/) (2026)
- [E-MMDiT: Revisiting Multimodal Diffusion Transformer Design for Fast Image Synthesis under Limited Resources](https://mlsanity.com/paper/2510.27135/) (2025)
- [DDiT: Dynamic Patch Scheduling for Efficient Diffusion Transformers](https://mlsanity.com/paper/2602.16968/) (2026)
- [FrequencyBooster: Full-Frequency Modeling for High-Fidelity Pixel Diffusion](https://mlsanity.com/paper/2605.17759/) (2026)
- [On the Generalization in Topology Optimization via Sensitivity-Conditioned Bernoulli Flow Matching](https://mlsanity.com/paper/2606.02179/) (2026)

## Abstract (arXiv)

Diffusion Transformers (DiTs) introduce the transformer architecture to diffusion tasks for latent-space image generation. With an isotropic architecture that chains a series of transformer blocks, DiTs demonstrate competitive performance and good scalability; but meanwhile, the abandonment of U-Net by DiTs and their following improvements is worth rethinking. To this end, we conduct a simple toy experiment by comparing a U-Net architectured DiT with an isotropic one. It turns out that the U-Net architecture only gain a slight advantage amid the U-Net inductive bias, indicating potential redundancies within the U-Net-style DiT. Inspired by the discovery that U-Net backbone features are low-frequency-dominated, we perform token downsampling on the query-key-value tuple for self-attention that bring further improvements despite a considerable amount of reduction in computation. Based on self-attention with downsampled tokens, we propose a series of U-shaped DiTs (U-DiTs) in the paper and conduct extensive experiments to demonstrate the extraordinary performance of U-DiT models. The proposed U-DiT could outperform DiT-XL/2 with only 1/6 of its computation cost. Codes are available at https://github.com/YuchuanTian/U-DiT.