REPA Works Until It Doesn't: Early-Stopped, Holistic Alignment Supercharges Diffusion Training
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866915298633842688 |
|---|---|
| author | Wang, Ziqiao Zhao, Wangbo Zhou, Yuhao Li, Zekai Liang, Zhiyuan Shi, Mingjia Zhao, Xuanlei Zhou, Pengfei Zhang, Kaipeng Wang, Zhangyang Wang, Kai You, Yang |
| author_facet | Wang, Ziqiao Zhao, Wangbo Zhou, Yuhao Li, Zekai Liang, Zhiyuan Shi, Mingjia Zhao, Xuanlei Zhou, Pengfei Zhang, Kaipeng Wang, Zhangyang Wang, Kai You, Yang |
| contents | 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 . |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2505_16792 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | REPA Works Until It Doesn't: Early-Stopped, Holistic Alignment Supercharges Diffusion Training Wang, Ziqiao Zhao, Wangbo Zhou, Yuhao Li, Zekai Liang, Zhiyuan Shi, Mingjia Zhao, Xuanlei Zhou, Pengfei Zhang, Kaipeng Wang, Zhangyang Wang, Kai You, Yang Computer Vision and Pattern Recognition Artificial Intelligence 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 . |
| title | REPA Works Until It Doesn't: Early-Stopped, Holistic Alignment Supercharges Diffusion Training |
| topic | Computer Vision and Pattern Recognition Artificial Intelligence |
| url | https://arxiv.org/abs/2505.16792 |