用分阶段对齐加速扩散模型训练,前快后稳不降质。
REPA Works Until It Doesn't: Early-Stopped, Holistic Alignment Supercharges Diffusion Training
- 前阶段联合对齐教师注意力与特征,快速收敛;后阶段一键关闭对齐,释放生成能力。
- 在ImageNet上仅需50轮达基线FID,优化步数减少28倍。
- 无需改架构,通用性强,适合图像生成与文生图任务。
扩散Transformer(DiT)虽能生成高质量图像,但训练速度慢。近期提出的表示对齐(REPA)虽加速初期训练,但后期性能停滞甚至下降。我们发现根源在于容量不匹配:当学生模型开始建模联合数据分布后,教师低维嵌入与注意力模式反而成为束缚。为此提出HASTE(分阶段终止的全面对齐),分两阶段进行:第一阶段在中层同时蒸馏教师注意力图(关系先验)和特征投影(语义锚点),实现快速收敛;第二阶段在固定迭代后一键关闭对齐损失,让DiT专注去噪并发挥生成潜力。HASTE无需修改架构即可加速多种DiT训练。在ImageNet 256×256上,50轮即达原版SiT-XL/2基线FID,500轮时媲美REPA最优FID,优化步数减少28倍。在MS-COCO上也提升文生图性能,证明其为跨任务高效的扩散训练通用方案。代码已开源。
原文摘要 · Abstract (English)
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 .
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。