Research map / Image generation / Autoregressive and latent generation

UDT: Reconciling U-Nets and Diffusion Transformers with Data-Adaptive Token Reduction

NeurIPS 2026 · Image generation › Autoregressive and latent generation · approach: High-resolution diffusion synthesis · 0 citations (OpenAlex)

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 UDT, a U-Net diffusion transformer that uses parameter-free, data-adaptive token merging/unmerging for down/upsampling while keeping the DiT token dimension, achieving faster convergence on ImageNet generation.

Problem

Isotropic diffusion transformers learn representations that degrade toward the output, and existing U-Net DiTs use fixed 2x2 spatial downsampling with learnable operators that ignore token similarity and complicate representation alignment.

Approach

Replace learnable fixed-neighborhood downsampling in U-shaped DiTs with training-free bipartite token merging (ToMe) and exact unmerging using recorded merge indices, preserving token channel dimension so that REPA and cross-attention remain compatible.

Key results

• UDT-XL/2+ without CFG reaches 7.7 FID (UDT) and 6.1 (UDT+) at 80 epochs, versus SiT-XL/2's 17.2 at 80 epochs and 7.9-level performance at 1400 epochs.
• With CFG and SD-VAE, UDT-XL/2+ + REPA reaches FID 1.38 at 320 epochs; UDT-XL/2+ with VA-VAE reaches FID 1.35 at 500 epochs.
• UDT uses about 75% of SiT GFLOPs at the same model size (e.g., UDT-B/2 17.7 vs SiT-B/2 23.0 GFLOPs) while achieving lower FID (24.0 vs 33.0 at 80 epochs).

Surprising findings

• Learnable token-reduction operators degrade FID sharply (45.5 vs 24.0 for data-adaptive merging at the same 256->112 token count), indicating token mixing hurts diffusion transformers.
• Out-of-the-box UDT without REPA or architectural tweaks already matches or beats several REPA-augmented baselines in the no-CFG setting.

Limitations

• Scope excludes video generation and 2K-resolution images, per the authors.
• Gains at the T2I setting are evaluated only on MSCOCO with a small 24-layer model and 150K iterations.

Experimental setup

Datasets: ImageNet MSCOCO

Benchmarks: ImageNet 256x256 ImageNet 512x512 MSCOCO

Metrics: FID IS sFID precision recall linear probing accuracy GFLOPs

Baselines: SiT U-DiT SiT↓ (UREPA) REPA DiT LightningDiT JiT MMDiT SRA LSEP DDT REPA-E

Builds on: ToMe SiT REPA DiT U-DiT

Closest papers by meaning

• U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers (NeurIPS 2024 · similarity 0.94)
• Layer- and Timestep-Adaptive Differentiable Token Compression Ratios for Efficient Diffusion Transformers (CVPR 2025 · similarity 0.90)
• U-REPA: Aligning Diffusion U-Nets to ViTs (NeurIPS 2025 · similarity 0.90)
• SD-DiT: Unleashing the Power of Self-supervised Discrimination in Diffusion Transformer (CVPR 2024 · similarity 0.89)
• Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think (ICLR 2025 · similarity 0.89)
• SparseDiT: Token Sparsification for Efficient Diffusion Transformer (NeurIPS 2025 · similarity 0.89)