REPA：让 DiT 训得快一点

§1 REPA 是什么

一句话：正常 DiT 训练只让模型学会预测噪声 / velocity；REPA 额外要求 DiT 的某一层中间特征，长得像 DINOv2 的特征。

注意，这里指的不是前面 VAE 的表征，而是后面 Transformer Block 的 representation。

§2 为什么要做

直观感受是：DiT 训练，收敛都比较慢。那么这种情况下应该怎么办？

作者做了一些 DiT representation 的探索，发现：

• DiT 中间层的 linear probing 表现远不如 DINOv2；
• 就算想把 DiT 的表征和 DINOv2 对齐去测一下，也对不齐。

由此自然冒出一个想法：能不能让 DiT 的 representation 显式地和 DINOv2 对齐？这种情况是否能提升能力？—— 实验答案是能，而且收敛速度提升一档（SiT-XL 大概 17.5× iter speedup）。

§3 怎么做

两条支路

训练时同时跑两条支路：

第一条：原来的 diffusion 训练支路。图像 x 先经 VAE encoder 变成 latent z = E(x)，加噪得到 z_t，送进 DiT，DiT 输出 velocity 预测目标 —— 这部分就是普通 diffusion loss。

第二条：REPA 表征对齐支路。同一张原图 x 送进一个冻结的预训练视觉 encoder（比如 DINOv2），拿到高质量语义特征 y* = f(x)。同时 DiT 处理 z_t 时，从某一层取出 hidden state h_t。h_t 的维度通常和 DINOv2 不一样，所以加一个可训练投影层 MLP：y = h_φ(h_t)。最后让 y 和 y* 尽量相似（cosine similarity）。

总 loss

两条支路的 loss 加权相加：

论文里 λ = 0.5，效果对这个值不敏感（0.25–1.0 都差不多）。

REPA 主要的贡献，就是发现"representation alignment 是 DiT 训练加速的有效杠杆"，并且证明这个简单 regularization 很有效。

§4 两个容易忽略的细节

① 加在前几层（不是最后一层）

DiT 的各层有自然的功能分工：

• 浅层 / 中间层：更像在建立图像结构、局部到全局的语义表示。
• 深层 / 最后层：更接近 diffusion 任务本身，比如预测 noise / velocity。

所以实验发现 REPA 加在前几层效果最好，论文最终选择前 8 层都加一个 regularization。直觉是：用 REPA 帮前面层把语义这一摊子事接管掉，让后面层专心做 denoising。

② Patch-level 对齐，不是 image-level

REPA 不是只拿一整张图的 global embedding 去对齐，而是更细一点：按 patch / token 对齐。

DiT 里 latent image 会被切成很多 patch tokens，某一层 hidden state 可以看成 h_t ∈ ℝ^(N×D)，N 是 patch 数、D 是 hidden dim。DINOv2 也能输出 patch-level features y* ∈ ℝ^(N×D')。MLP 投影统一维度之后，对每个 patch 单独算相似度。

所以它不是说："这张图整体像 DINO 的语义。"而是说：

DiT 的每个 patch token，都应该在语义空间里更接近 DINOv2 对应 patch 的表示。

这个更合理，因为图像生成需要空间结构。如果只对齐 global feature，模型可能只学到"这是一只狗 / 一辆车"的整体语义，但不一定学到局部区域怎么组织。