扩散Transformer训练优化：HASTE框架解析与实践

1. 项目概述

今天要和大家分享的是2025年NIPS会议上的一篇重磅论文《REPA Works Until It Doesn't》，这篇论文针对扩散Transformer（DiTs）训练过程中的一个关键瓶颈问题提出了创新性解决方案。作为一名长期关注生成模型优化的研究者，我发现这个问题其实困扰着很多团队——为什么明明初期训练效果很好，到后期却突然"卡住"甚至性能下降？这篇论文不仅精准定位了问题根源，还给出了一个简单却有效的框架HASTE，实测能达到28倍的训练加速，这对我们实际项目中的模型训练效率提升有着重要参考价值。

扩散Transformer作为当前图像生成领域的SOTA架构，其训练过程却异常耗时。以标准的ImageNet 256×256训练为例，通常需要1400个epoch才能达到理想效果。虽然REPA方法通过特征对齐在初期显著加速了收敛，但论文发现当DiT开始学习复杂的联合数据分布时，教师模型的低维嵌入反而会成为性能瓶颈。这个发现解开了我一直以来的疑惑——为什么我们团队之前使用REPA时，总在训练中期遇到难以突破的性能天花板。

2. 核心问题诊断

2.1 REPA方法的局限性

REPA（Representation Alignment）的基本思路很直观：使用预训练好的非生成式教师模型（如DINOv2）作为指导，通过匹配学生模型（DiT）的隐藏层特征来加速训练。在初期这确实很有效，因为教师模型已经学习到了良好的视觉表征。但论文通过详实的实验揭示了三个关键现象：

1. 梯度方向变化：训练初期，学生模型和教师模型的梯度方向呈锐角（高度一致），但随训练进行逐渐变为直角（正交）甚至钝角（冲突）
2. 容量不匹配：教师模型（如ViT-L）通常只有24层，而DiT学生模型可能有40+层，深层网络的表征能力远超教师模型
3. 任务差异：教师模型学习的是判别式特征，而DiT需要建模生成式任务的联合分布

重要发现：当学生模型的FID降到约15时，REPA的正面效果会急剧下降，继续使用反而会阻碍性能提升

2.2 问题本质分析

通过大量的消融实验，作者团队发现核心矛盾在于：

1. 维度限制：教师模型的CLS token仅包含768维特征，而生成任务需要更丰富的空间信息
2. 注意力模式固化：