Research map / Image generation / Fast diffusion sampling
REPA Works Until It Doesn’t: Early-Stopped, Holistic Alignment Supercharges Diffusion Training
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.
Proposes HASTE, a two-phase training schedule that applies holistic (attention + feature) alignment from a frozen DINOv2 teacher to mid-layers of diffusion transformers, then terminates the alignment loss to speed up DiT/SiT training.
Problem
Representation alignment (REPA) accelerates early diffusion transformer training but plateaus or degrades later, due to a capacity mismatch between the frozen non-generative teacher and the generative student.
Approach
Combine REPA's feature alignment with attention-map distillation from DINOv2 into mid-level DiT blocks (holistic alignment), then switch off all alignment losses at a fixed iteration (or gradient-angle trigger) so training reverts to pure denoising.
Key results
- On ImageNet 256x256, SiT-XL/2+HASTE reaches FID 8.39 at 50 epochs without CFG, matching vanilla SiT-XL/2 at 1400 epochs (28x fewer steps).
- With CFG, SiT-XL/2+HASTE at 500K iterations reaches FID 1.42, matching SiT+REPA at 800 epochs.
- On MS-COCO text-to-image, MM-DiT+HASTE improves FID over MM-DiT+REPA (6.81 vs 7.33 with SDE, no CFG).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add holistic alignment and terminate it at 50 epochs rather than leaving it active (ablation) | SiT-XL/2 · ImageNet 256×256, 100 epochs | FID: 8.1 -> 5.3 | Helped |
| Add attention alignment to feature alignment (ablation) | SiT-XL/2 · ImageNet 256×256, 40 epochs | FID: 10.7 -> 9.9 | Helped |
| Add attention alignment to feature alignment and keep both active (ablation) | SiT-XL/2 · ImageNet 256×256, 100 epochs | FID: 7.5 -> 8.1 | Hurt |
| Use attention alignment only instead of feature alignment only (ablation) | SiT-XL/2 · ImageNet 256×256, 40 epochs | FID: 10.7 -> 13.6 | Hurt |
| Terminate alignment at 100K rather than continuing it (ablation) | SiT-B/2 · ImageNet 256×256, 400K iterations | FID: 21.3 -> 19.6 | Helped |
| Terminate alignment at 250K rather than leaving it active (ablation) | SiT-XL/2 · ImageNet 256×256, 500K iterations | FID: 8.1 -> 5.3 | Helped |
| Terminate alignment at 400K rather than 250K (ablation) | SiT-XL/2 · ImageNet 256×256, 500K iterations | FID: 5.3 -> 7.4 | Hurt |
| Increase attention alignment weight from 0.5 to 3.0 (ablation) | SiT-L/2 with HASTE · ImageNet 256×256, 400K iterations | FID: 7.9 -> 9.0 | Hurt |
| Use only two aligned layer pairs instead of four (ablation) | SiT-L/2 with HASTE · ImageNet 256×256, 400K iterations | FID: 7.9 -> 8.9 | Hurt |
| Align additional shallower DINO and SiT layers (ablation) | SiT-L/2 with HASTE · ImageNet 256×256, 400K iterations | FID: 7.9 -> 8.3 | Hurt |
| Apply HASTE rather than MM-DiT+REPA | MM-DiT · MS-COCO 2014, SDE (NFEs = 250), without CFG | FID: 7.33 -> 6.81 | Helped |
| Apply HASTE rather than vanilla SiT | SiT-XL/2 · ImageNet 256×256, without CFG | FID: 8.61 -> 8.39 | Helped |
Surprising findings
- Stopping REPA late (400K) improves FID over always-on REPA, while stopping early (100K) hurts.
- Attention alignment alone achieves a speed-up comparable to REPA.
Limitations
- Focused only on latent diffusion transformers for image generation; pixel-level diffusion and video not explored.
- Termination point tau is chosen by hand/iteration per model size; gradient-angle trigger is described but not fully validated.
- Text-to-image experiments do not use termination due to limited iterations.
Experimental setup
Datasets: ImageNet MS-COCO 2014
Benchmarks: ImageNet 256x256 MS-COCO 2014
Metrics: FID sFID IS Precision Recall
Baselines: REPA DiT SiT MaskDiT FasterDiT MDT MDTv2 DiT+TREAD MM-DiT+REPA
Builds on: REPA attention transfer SiT DiT
Closest papers by meaning
- Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think (ICLR 2025 · similarity 0.92)
- U-REPA: Aligning Diffusion U-Nets to ViTs (NeurIPS 2025 · similarity 0.92)
- AHPA: Adaptive Hierarchical Prior Alignment for Diffusion Transformers (NeurIPS 2026 · similarity 0.90)
- Representation Entanglement for Generation: Training Diffusion Transformers Is Much Easier Than You Think (NeurIPS 2025 · similarity 0.90)
- Learning Diffusion Models with Flexible Representation Guidance (NeurIPS 2025 · similarity 0.89)
- Why DiT Models Underperform as Representation Learners without Long Skip Connections (NeurIPS 2026 · similarity 0.89)
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.