REPA Works Until It Doesn't: Early-Stopped, Holistic Alignment Supercharges Diffusion Training
Ziqiao Wang, Wangbo Zhao, Yuhao Zhou, Zekai Li, Zhiyuan Liang, Mingjia Shi, Xuanlei Zhao, Pengfei Zhou, Kaipeng Zhang, Zhangyang (Atlas) Wang, Kai Wang, Yang You
摘要
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 .
问问这篇 Paper
智能体会读完全文。
Lune 把这篇 Paper 索引到了最后一个公式,引用它的顶会 Paper 也一样。你提问,回答直接引用原文。
引用它的顶会 Paper11
- Representation Alignment for Diffusion Transformers without External ComponentsDengyang Jiang, Mengmeng Wang, Liuzhuozheng Li, Lei Zhang 等ICLR 2026 · 被引用 532 次
- What matters for Representation Alignment: Global Information or Spatial Structure?Jaskirat Singh, Xingjian Leng, Zongze Wu, Liang Zheng 等ICLR 2026 · 被引用 84 次
- Latent Forcing: Reordering the Diffusion Trajectory for Pixel-Space Image GenerationAlan Baade, Eric Chan, Kyle Sargent, Changan Chen 等ICML 2026 · 被引用 25 次
- Adapting Self-Supervised Representations as a Latent Space for Efficient GenerationMing Gui, Johannes Schusterbauer, Timy Phan, Felix Krause 等ICLR 2026 · 被引用 13 次
- SPRINT: Sparse-Dense Residual Fusion for Efficient Diffusion TransformersDogyun Park, Moayed Haji-Ali, Yanyu Li, Willi Menapace 等ICLR 2026 · 被引用 6 次
它引用的顶会 Paper34
- Learning Transferable Visual Models From Natural Language SupervisionAlec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh 等ICML 2021 · 被引用 47,906 次
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 被引用 35,902 次
- An Image is Worth 16x16 Words: Transformers for Image Recognition at ScaleAlexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn 等ICLR 2021 · 被引用 21,477 次
- Diffusion Models Beat GANs on Image SynthesisPrafulla Dhariwal, Alexander Quinn NicholNeurIPS 2021 · 被引用 13,211 次
- High-Resolution Image Synthesis with Latent Diffusion ModelsRobin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser 等CVPR 2022 · 被引用 13,123 次
相关 Paper
- Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You ThinkSihyun Yu, Sangkyung Kwak, Huiwon Jang, Jongheon Jeong 等ICLR 2025
- U-REPA: Aligning Diffusion U-Nets to ViTsYuchuan Tian, Hanting Chen, Mengyu Zheng, Yuchen Liang 等NeurIPS 2025 · 被引用 30 次
- SD-DiT: Unleashing the Power of Self-Supervised Discrimination in Diffusion Transformer*Rui Zhu, Yingwei Pan, Yehao Li, Ting Yao 等CVPR 2024 · 被引用 15 次
- Diffusion Transformers with Representation AutoencodersBoyang Zheng, Nanye Ma, Shengbang Tong, Saining XieICLR 2026 · 被引用 288 次
- REPA-E: Unlocking VAE for End-to-End Tuning with Latent Diffusion TransformersXingjian Leng, Jaskirat Singh, Yunzhong Hou, Zhenchang Xing 等ICCV 2025 · 被引用 15 次
