# U-REPA: Aligning Diffusion U-Nets to ViTs

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

## Key facts

- Cited by 32 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 38 of its 48 references that are in the map.
- First posted on arXiv 24 Mar 2025.

## Placed in

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

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

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [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)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [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)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (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)
- [Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think](https://mlsanity.com/paper/2410.06940/) (2024)
- [Open-Sora: Democratizing Efficient Video Production for All](https://mlsanity.com/paper/2412.20404/) (2024)
- [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)
- [Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models](https://mlsanity.com/paper/2501.01423/) (2025)
- [Pyramidal Flow Matching for Efficient Video Generative Modeling](https://mlsanity.com/paper/2410.05954/) (2024)
- [PixArt-Sigma: Weak-to-Strong Training of Diffusion Transformer for 4K Text-to-Image Generation](https://mlsanity.com/paper/2403.04692/) (2024)
- [Understanding Diffusion Objectives as the ELBO with Simple Data Augmentation](https://mlsanity.com/paper/2303.00848/) (2023)
- [Applying Guidance in a Limited Interval Improves Sample and Distribution Quality in Diffusion Models](https://mlsanity.com/paper/2404.07724/) (2024)
- [Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps](https://mlsanity.com/paper/2501.09732/) (2025)
- [FreeU: Free Lunch in Diffusion U-Net](https://mlsanity.com/paper/2309.11497/) (2023)
- [Structural Pruning for Diffusion Models](https://mlsanity.com/paper/2305.10924/) (2023)
- [Fast Training of Diffusion Models with Masked Transformers](https://mlsanity.com/paper/2306.09305/) (2023)
- [Delta-DiT: A Training-Free Acceleration Method Tailored for Diffusion Transformers](https://mlsanity.com/paper/2406.01125/) (2024)

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

- [VideoREPA: Learning Physics for Video Generation through Relational Alignment with Foundation Models](https://mlsanity.com/paper/2505.23656/) (2025)
- [DeCo: Frequency-Decoupled Pixel Diffusion for End-to-End Image Generation](https://mlsanity.com/paper/2511.19365/) (2025)
- [REPA Works Until It Doesn't: Early-Stopped, Holistic Alignment Supercharges Diffusion Training](https://mlsanity.com/paper/2505.16792/) (2025)
- [PixelU: A U-Shaped Transformer for Efficient End-to-End Pixel Diffusion](https://mlsanity.com/paper/2606.27760/) (2026)
- [DivControl: Knowledge Diversion for Controllable Image Generation](https://mlsanity.com/paper/2507.23620/) (2025)
- [PhysAlign: Physics-Coherent Image-to-Video Generation through Feature and 3D Representation Alignment](https://mlsanity.com/paper/2603.13770/) (2026)
- [LUMA: Low-Dimension Unified Motion Alignment with Dual-Path Anchoring for Text-to-Motion Diffusion Model](https://mlsanity.com/paper/2509.25304/) (2025)
- [R3DPA: Leveraging 3D Representation Alignment and RGB Pretrained Priors for LiDAR Scene Generation](https://mlsanity.com/paper/2601.07692/) (2026)
- [Test-Time Conditioning with Representation-Aligned Visual Features](https://mlsanity.com/paper/2602.03753/) (2026)
- [V-Co: A Closer Look at Visual Representation Alignment via Co-Denoising](https://mlsanity.com/paper/2603.16792/) (2026)
- [MPDiT: Multi-Patch Global-to-Local Transformer Architecture For Efficient Flow Matching and Diffusion Model](https://mlsanity.com/paper/2603.26357/) (2026)
- [ReWorld: Representation Learning for World Action Models](https://mlsanity.com/paper/2606.27504/) (2026)
- [Repulsor: Accelerating Generative Modeling with a Contrastive Memory Bank](https://mlsanity.com/paper/2512.08648/) (2025)
- [RePack then Refine: Efficient Diffusion Transformer with Vision Foundation Model](https://mlsanity.com/paper/2512.12083/) (2025)
- [Self-transcendence: Is External Feature Guidance Indispensable for Accelerating Diffusion Transformer Training?](https://mlsanity.com/paper/2601.07773/) (2026)
- [Empty SPACE: Cross-Attention Sparsity for Concept Erasure in Diffusion Models](https://mlsanity.com/paper/2605.10198/) (2026)
- [Pixel-Space Diffusion Transformers](https://mlsanity.com/paper/2607.17585/) (2026)
- [UDT: Reconciling U-Nets and Diffusion Transformers with Data-Adaptive Token Reduction](https://mlsanity.com/paper/2608.01298/) (2026)
- [Align & Invert: Solving Inverse Problems with Diffusion and Flow-based Models via Representation Alignment](https://mlsanity.com/paper/2511.16870/) (2025)
- [Divide and Conquer: Decoupled Representation Alignment for Multimodal World Models](https://mlsanity.com/paper/2605.01896/) (2026)
- [Resolving Representation Ambiguity in Feedforward Novel View Synthesis Transformer via Semantic-Spatial Decoupling](https://mlsanity.com/paper/2605.18599/) (2026)
- [ROAD: Reciprocal-Objective Alignment of Discriminative Semantics for 3D Shape Generation](https://mlsanity.com/paper/2607.28581/) (2026)
- [VoxStruct3D: Structure-Leading Flow Matching for Voxel-Space 3D MRI Synthesis](https://mlsanity.com/paper/2608.04557/) (2026)
- [Representation Learning in Diffusion and Flow-based Model: An Application Aspect](https://mlsanity.com/paper/2608.24068/) (2026)
- [Revisiting Spectral Representations in Generative Diffusion Models](https://mlsanity.com/paper/2609.08253/) (2026)
- [ComplexSync: High-Fidelity and Real-Time Lip Sync in Complex Scenarios](https://mlsanity.com/paper/2609.29225/) (2026)
- [What Visual Generators Need from Teachers: Rethinking Representation Alignment](https://mlsanity.com/paper/2609.34732/) (2026)
- [Scaffold Then Internalize: Representation Injection for Diffusion Transformers](https://mlsanity.com/paper/2609.35292/) (2026)
- [Embedding Prediction Helps Image Generation](https://mlsanity.com/paper/2610.02203/) (2026)
- [ProAR: Learning Prospective Reasoning with Autoregressive Video Models](https://mlsanity.com/paper/2610.03664/) (2026)

## Abstract (arXiv)

Representation Alignment (REPA) that aligns Diffusion Transformer (DiT) hidden-states with ViT visual encoders has proven highly effective in DiT training, demonstrating superior convergence properties, but it has not been validated on the canonical diffusion U-Net architecture that shows faster convergence compared to DiTs. However, adapting REPA to U-Net architectures presents unique challenges: (1) different block functionalities necessitate revised alignment strategies; (2) spatial-dimension inconsistencies emerge from U-Net's spatial downsampling operations; (3) space gaps between U-Net and ViT hinder the effectiveness of tokenwise alignment. To encounter these challenges, we propose U-REPA, a representation alignment paradigm that bridges U-Net hidden states and ViT features as follows: Firstly, we propose via observation that due to skip connection, the middle stage of U-Net is the best alignment option. Secondly, we propose upsampling of U-Net features after passing them through MLPs. Thirdly, we observe difficulty when performing tokenwise similarity alignment, and further introduces a manifold loss that regularizes the relative similarity between samples. Experiments indicate that the resulting U-REPA could achieve excellent generation quality and greatly accelerates the convergence speed. With CFG guidance interval, U-REPA could reach FID<1.5 in 200 epochs or 1M iterations on ImageNet 256 × 256, and needs only half the total epochs to perform better than REPA under sd-vae-ft-ema. Codes: https://github.com/YuchuanTian/U-REPA