# REPA Works Until It Doesn't: Early-Stopped, Holistic Alignment Supercharges Diffusion Training

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

## Key facts

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

## Placed in

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

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [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)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [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)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Improved Techniques for Training GANs](https://mlsanity.com/paper/1606.03498/) (2016)
- [Video Diffusion Models](https://mlsanity.com/paper/2204.03458/) (2022)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learning](https://mlsanity.com/paper/1702.03118/) (2017)
- [Improved Precision and Recall Metric for Assessing Generative Models](https://mlsanity.com/paper/1904.06991/) (2019)
- [WorldSimBench: Towards Video Generation Models as World Simulators](https://mlsanity.com/paper/2410.18072/) (2024)
- [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)
- [All are Worth Words: A ViT Backbone for Diffusion Models](https://mlsanity.com/paper/2209.12152/) (2022)
- [Analyzing and Improving the Training Dynamics of Diffusion Models](https://mlsanity.com/paper/2312.02696/) (2023)
- [Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models](https://mlsanity.com/paper/2501.01423/) (2025)
- [SANA: Efficient High-Resolution Image Synthesis with Linear Diffusion Transformers](https://mlsanity.com/paper/2410.10629/) (2024)

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

- [What matters for Representation Alignment: Global Information or Spatial Structure?](https://mlsanity.com/paper/2512.10794/) (2025)
- [Latent Forcing: Reordering the Diffusion Trajectory for Pixel-Space Image Generation](https://mlsanity.com/paper/2602.11401/) (2026)
- [Representation Alignment for Just Image Transformers is not Easier than You Think](https://mlsanity.com/paper/2603.14366/) (2026)
- [Adapting Self-Supervised Representations as a Latent Space for Efficient Generation](https://mlsanity.com/paper/2510.14630/) (2025)
- [Stable Velocity: A Variance Perspective on Flow Matching](https://mlsanity.com/paper/2602.05435/) (2026)
- [SAMA: Factorized Semantic Anchoring and Motion Alignment for Instruction-Guided Video Editing](https://mlsanity.com/paper/2603.19228/) (2026)
- [Data Efficient Any Transformer-to-Mamba Distillation via Attention Bridge](https://mlsanity.com/paper/2510.19266/) (2025)
- [Sprint: Sparse-Dense Residual Fusion for Efficient Diffusion Transformers](https://mlsanity.com/paper/2510.21986/) (2025)
- [TORA: Topological Representation Alignment for 3D Shape Assembly](https://mlsanity.com/paper/2604.04050/) (2026)
- [LUMA: Low-Dimension Unified Motion Alignment with Dual-Path Anchoring for Text-to-Motion Diffusion Model](https://mlsanity.com/paper/2509.25304/) (2025)
- [Teacher-Feature Drifting: One-Step Diffusion Distillation with Pretrained Diffusion Representations](https://mlsanity.com/paper/2605.07327/) (2026)
- [R3DPA: Leveraging 3D Representation Alignment and RGB Pretrained Priors for LiDAR Scene Generation](https://mlsanity.com/paper/2601.07692/) (2026)
- [PILD: Physics-Informed Learning via Diffusion](https://mlsanity.com/paper/2601.21284/) (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)
- [MaskAlign: Token-Subset Representation Alignment for Efficient Diffusion Training](https://mlsanity.com/paper/2606.08788/) (2026)
- [Persistence Forcing: Exploiting Feature Specialization in Pixel-Space Diffusion](https://mlsanity.com/paper/2609.36014/) (2026)
- [Self-transcendence: Is External Feature Guidance Indispensable for Accelerating Diffusion Transformer Training?](https://mlsanity.com/paper/2601.07773/) (2026)
- [AG-REPA: Causal Layer Selection for Representation Alignment in Audio Flow Matching](https://mlsanity.com/paper/2603.01006/) (2026)
- [Generate in Reconstruction Space, Match in Semantic Space: Transport Geometry for One-Step Generation](https://mlsanity.com/paper/2606.00514/) (2026)
- [ELF-REG: Scaling Continuous Diffusion Language Models to Reasoning Tasks](https://mlsanity.com/paper/2609.29102/) (2026)
- [Align & Invert: Solving Inverse Problems with Diffusion and Flow-based Models via Representation Alignment](https://mlsanity.com/paper/2511.16870/) (2025)
- [Motif-Video 2B: Technical Report](https://mlsanity.com/paper/2604.16503/) (2026)
- [From SRA to Self-Flow: Data Augmentation or Self-Supervision?](https://mlsanity.com/paper/2607.02508/) (2026)
- [ReGen: Hierarchical Multi-Prompt Representation Generation for Efficient Waveform Diffusion Models](https://mlsanity.com/paper/2607.09134/) (2026)
- [SQuad: Sub-Quadratic Attention Distillation for Efficient Video Generation](https://mlsanity.com/paper/2608.16585/) (2026)
- [What Does Attention Transfer Transfer? Attention Structure and Robustness in Vision Transformers](https://mlsanity.com/paper/2608.18399/) (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)
- [On the Diffusibility of High-Dimensional Latents](https://mlsanity.com/paper/2609.28473/) (2026)

## Abstract (arXiv)

Diffusion Transformers (DiTs) deliver state-of-the-art image quality, yet their training remains notoriously slow. A recent remedy -- representation alignment (REPA) that matches DiT hidden features to those of a non-generative teacher (e.g. DINO) -- dramatically accelerates the early epochs but plateaus or even degrades performance later. We trace this failure to a capacity mismatch: once the generative student begins modelling the joint data distribution, the teacher's lower-dimensional embeddings and attention patterns become a straitjacket rather than a guide. We then introduce HASTE (Holistic Alignment with Stage-wise Termination for Efficient training), a two-phase schedule that keeps the help and drops the hindrance. Phase I applies a holistic alignment loss that simultaneously distills attention maps (relational priors) and feature projections (semantic anchors) from the teacher into mid-level layers of the DiT, yielding rapid convergence. Phase II then performs one-shot termination that deactivates the alignment loss, once a simple trigger such as a fixed iteration is hit, freeing the DiT to focus on denoising and exploit its generative capacity. HASTE speeds up training of diverse DiTs without architecture changes. On ImageNet 256X256, it reaches the vanilla SiT-XL/2 baseline FID in 50 epochs and matches REPA's best FID in 500 epochs, amounting to a 28X reduction in optimization steps. HASTE also improves text-to-image DiTs on MS-COCO, demonstrating to be a simple yet principled recipe for efficient diffusion training across various tasks. Our code is available at https://github.com/NUS-HPC-AI-Lab/HASTE .