Research map / Architectures and efficiency / Attention and transformers
U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers
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 U-DiT, a U-Net-style diffusion transformer that performs self-attention on 2x-downsampled query-key-value tokens (pixel-unshuffle, lossless, parallel over four sub-grids), reaching DiT-XL/2 quality with about 1/6 the FLOPs on ImageNet 256.
Problem
Diffusion Transformers abandoned the U-Net architecture in favor of isotropic transformer stacks for latent-space generation, although the U-Net inductive bias is believed to help denoising; it is unclear whether U-Net-style DiTs can be made competitive and efficient.
Approach
Downsample queries, keys and values together (via pixel-unshuffle into four 2x-downsampled sub-grids) for self-attention within a U-Net DiT, keeping tensor size constant so the attention cost drops to 1/4 with no information loss, motivated by the low-frequency dominance of U-Net features.
Key results
- U-DiT-B (22.2 GFLOPs) reaches FID 16.64 on ImageNet 256 at 400K steps, beating DiT-XL/2 (118.6 GFLOPs, FID 19.47 reported / 20.05 reproduced).
- U-DiT-L reaches FID 10.08 at 400K steps (vs. 20.05 for reproduced DiT-XL/2) and 7.54 at 1M steps, versus 9.62 for DiT-XL/2 at 7M steps.
- Downsampled self-attention on the DiT-UNet toy model lowers FID from 93.48 to 89.43 while reducing GFLOPs from 1.40 to 0.90.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace full-scale self-attention with query-key-value token downsampling (ablation) | DiT-UNet · ImageNet 256×256 | FID: 93.48 -> 89.43 | Helped |
| use key-value downsampling only instead of query-key-value token downsampling (ablation) | DiT-UNet · ImageNet 256×256 | FID: 89.43 -> 94.38 | Hurt |
| add a parallel shortcut to depthwise convolution in the downsampler (ablation) | U-DiT-T · ImageNet 256×256 | FID: 89.87 -> 89.43 | Helped |
| use Pixel Shuffle downsampling instead of the depthwise-convolution-plus-shortcut design (ablation) | U-DiT-T · ImageNet 256×256 | FID: 89.43 -> 96.15 | Hurt |
| add cosine-similarity attention (ablation) | U-DiT-T · ImageNet 256×256 | FID: 89.43 -> 86.96 | Helped |
| add RoPE2D (ablation) | U-DiT-T · ImageNet 256×256 | FID: 89.43 -> 84.64 | Helped |
| add depthwise convolution to the FFN (ablation) | U-DiT-T · ImageNet 256×256 | FID: 89.43 -> 79.30 | Helped |
| add FFN re-parameterization (ablation) | U-DiT-T · ImageNet 256×256 | FID: 89.43 -> 75.71 | Helped |
| add all reported modifications (ablation) | U-DiT-S · ImageNet 256×256 | FID: 41.01 -> 31.51 | Helped |
| add all reported modifications (ablation) | U-DiT-B · ImageNet 256×256 | FID: 20.89 -> 16.64 | Helped |
| add all reported modifications (ablation) | U-DiT-L · ImageNet 256×256 | FID: 12.04 -> 10.08 | Helped |
Surprising findings
- Downsampling queries as well as keys and values improves FID while cutting attention cost, despite losing spatial resolution inside attention.
- A plain DiT-UNet is only marginally better than an isotropic DiT of similar size, suggesting U-Net inductive bias is under-used without modification.
Limitations
- Authors lacked compute to extend training further or scale model size fully.
- Toy experiment and ablations are on small models (DiT-S-scale) trained for 400K steps.
Experimental setup
Datasets: ImageNet 2012
Benchmarks: ImageNet 256x256 ImageNet 512x512
Metrics: FID sFID IS Precision Recall FLOPs
Baselines: DiT-XL/2 DiT-L/2 DiT-B/2 DiT-S/2 U-ViT-L U-ViT-XL PixArt-alpha-XL/2 DiffiT-XL/2 SiT SiT-LLaMA
Builds on: DiT U-ViT Swin Transformer V2 RoPE Stable Diffusion
Closest papers by meaning
- UDT: Reconciling U-Nets and Diffusion Transformers with Data-Adaptive Token Reduction (NeurIPS 2026 · similarity 0.94)
- Layer- and Timestep-Adaptive Differentiable Token Compression Ratios for Efficient Diffusion Transformers (CVPR 2025 · similarity 0.89)
- Exploring Diffusion Transformer Designs via Grafting (NeurIPS 2025 · similarity 0.88)
- SparseDiT: Token Sparsification for Efficient Diffusion Transformer (NeurIPS 2025 · similarity 0.88)
- Dynamic Diffusion Transformer (ICLR 2025 · similarity 0.88)
- On Inductive Biases That Enable Generalization in Diffusion Transformers (NeurIPS 2025 · similarity 0.87)
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.