Research map / Image generation / Autoregressive and latent generation
U-REPA: Aligning Diffusion U-Nets to ViTs
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.
Adapts representation alignment (REPA) to diffusion U-Nets (SiT-downsampled) by aligning mid-stage features after MLP and upsampling, plus a manifold loss on sample-pair similarities, reaching FID 1.41 on ImageNet 256 at 400 epochs.
Problem
REPA accelerates DiT training but has not been validated on diffusion U-Nets, which face block-functionality changes from skip connections, spatial mismatch from downsampling, and feature-space gaps versus ViT encoders.
Approach
Align the middle U-Net stage (best due to skip connections), apply MLP before upsampling to match the ViT token grid, and add a manifold loss that matches pairwise cosine similarities between samples instead of rigid tokenwise alignment.
Key results
- SiT↓-XL/2+U-REPA reaches FID 1.41 with CFG on ImageNet 256 at 400 epochs (2M iterations), versus 1.42 for SiT-XL/2+REPA at 800 epochs.
- U-REPA reaches FID 1.48 at 200 epochs with CFG on ImageNet 256.
- Without CFG at 400K iterations, SiT↓-XL/2+U-REPA reaches FID 5.4 versus 5.9 for SiT-XL/2+REPA at 4M iterations.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply U-REPA instead of REPA to SiT↓-XL/2 | SiT↓-XL/2 · ImageNet 256×256, w/o cfg | FID: 5.9 -> 5.4 | Helped |
| Use U-REPA instead of REPA | SiT-XL/2 · ImageNet 256×256, w/ cfg | FID: 1.42 -> 1.41 | Helped |
| Use U-REPA instead of REPA | SiT-XL/2 · ImageNet 512×512, w/ cfg | FID: 2.44 -> 2.21 | Helped |
| Align at the centering layer (depth 18) rather than depth 8 (ablation) | SiT↓-XL/2+U-REPA · ImageNet 256×256, w/ cfg | FID: 9.35 -> 6.25 | Helped |
| Upsample U-Net features rather than downsample DINOv2 features (ablation) | SiT↓-XL/2+U-REPA · ImageNet 256×256, w/ cfg | FID: 5.99 -> 5.72 | Helped |
| Apply MLP before upsampling (ablation) | SiT↓-XL/2+U-REPA · ImageNet 256×256, w/ cfg | FID: 5.84 -> 5.72 | Helped |
| Apply MLP before upsampling rather than upsampling within MLP (ablation) | SiT↓-XL/2+U-REPA · ImageNet 256×256, w/ cfg | FID: 6.36 -> 5.72 | Helped |
| Add manifold loss with weight w=3 (ablation) | SiT↓-XL/2+U-REPA · ImageNet 256×256, w/ cfg | FID: 6.25 -> 5.72 | Helped |
| Add time-aware shift and scale to the MLP output (ablation) | SiT↓-XL/2+U-REPA · ImageNet 256×256, w/ cfg | FID: 5.72 -> 5.63 | Helped |
| Use the triangular loss-weight schedule instead of a constant weight (ablation) | SiT↓-XL/2+U-REPA · ImageNet 256×256, w/ cfg | FID: 5.72 -> 5.58 | Helped |
| Use schedule max(1,t+0.5) rather than a constant loss weight (ablation) | SiT↓-XL/2+U-REPA · ImageNet 256×256, w/ cfg | FID: 5.72 -> 5.85 | Hurt |
| Use schedule max(1,-t+1.5) rather than a constant loss weight (ablation) | SiT↓-XL/2+U-REPA · ImageNet 256×256, w/ cfg | FID: 5.72 -> 5.72 | No clear effect |
Surprising findings
- Downsampling, not skip connections, accounts for most of the U-Net convergence advantage in the toy DiT experiments (FID 19.86 to 13.78 from downsampling vs 20.05 to 19.86 from skips).
Limitations
- The U-Net is simple with only one intermediate stage and is not further refined.
- Applicability to downstream U-Net tasks such as low-level vision is not investigated.
- Comparisons of SiT↓ against SiT are argued to be less direct because aligning a U-Net to a ViT is harder.
Experimental setup
Datasets: ImageNet
Benchmarks: ImageNet 256x256 ImageNet 512x512
Metrics: FID IS training energy (J)
Baselines: REPA SiT-XL/2 DiT-XL/2 MDTv2-XL/2 U-ViT-H/2 ADM-U VDM++ Simple diffusion DiffiT MaskDiT
Builds on: REPA SiT DiT manifold knowledge distillation classifier-free guidance interval
Closest papers by meaning
- REPA Works Until It Doesn’t: Early-Stopped, Holistic Alignment Supercharges Diffusion Training (NeurIPS 2025 · similarity 0.92)
- UDT: Reconciling U-Nets and Diffusion Transformers with Data-Adaptive Token Reduction (NeurIPS 2026 · similarity 0.90)
- Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think (ICLR 2025 · similarity 0.88)
- AHPA: Adaptive Hierarchical Prior Alignment for Diffusion Transformers (NeurIPS 2026 · similarity 0.88)
- U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers (NeurIPS 2024 · similarity 0.87)
- Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models (CVPR 2025 · similarity 0.86)
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.