Diffusion Transformers with Representation Autoencoders
核心思想
传统潜在扩散模型使用 VAE 将图像压缩到低维潜在空间。本文提出表征自编码器（RAE）：用冻结的预训练表征编码器（如 DINO、SigLIP）+ 训练的解码器替代 VAE。
核心发现：
- 预训练编码器的高维语义潜在空间更适合扩散训练
- 扩散模型在 RAE 空间收敛速度比 VAE 空间快 47 倍
- ImageNet 256×256 达到 FID 1.13（有引导）
背景知识
为什么需要潜在空间
直接在像素空间做扩散太贵（256×256×3 = 196,608 维）。VAE 把图像压缩到如 32×32×4 = 4,096 维的潜在空间 → 扩散在低维空间进行 → 速度快很多。
传统 VAE 的问题
| 问题 | 说明 |
|---|---|
| 语义贫乏 | VAE 潜在空间的线性探测准确率仅 8.0% |
| 信息丢失 | 压缩到 4 维/patch 丢失了大量语义信息 |
| 训练复杂 | 需要感知损失 + 对抗损失 + KL 正则 |
预训练编码器有什么好
| 编码器 | 线性探测准确率 | 类型 |
|---|---|---|
| SD-VAE | 8.0% | 压缩为主 |
| MAE-B | 68.0% | 自监督 |
| SigLIP2-B | 79.1% | 图文对比 |
| DINOv2-B | 84.5% | 自监督 |
预训练编码器的潜在空间语义极其丰富——这是 RAE 的根基。
方法详解
1. RAE 架构
编码器：冻结的预训练 ViT（如 DINOv2-B，patch size 14，隐藏维度 768）
解码器：可训练的 ViT，通过组合损失训练：
\[\mathcal{L}_\text{rec}(x) = \omega_L \cdot \text{LPIPS}(\hat{x}, x) + \text{L1}(\hat{x}, x) + \omega_G \cdot \lambda \cdot \text{GAN}(\hat{x}, x)\]- LPIPS：感知相似度
- L1：像素级精度
- GAN：对抗训练提升锐度
重建质量（rFID，越低越好）：
| 编码器 | rFID |
|---|---|
| MAE-B | 0.16 |
| DINOv2-B | 0.49 |
| SigLIP2-B | 0.53 |
| SD-VAE | 0.62 |
2. 宽度需求：理论分析
定理 1：设扩散模型的函数族为 $\mathcal{G}_d = {g(\mathbf{x}_t, t) = \mathbf{B} f(\mathbf{A}\mathbf{x}_t, t)}$，其中 $\mathbf{A} \in \mathbb{R}^{d \times n}$、$\mathbf{B} \in \mathbb{R}^{n \times d}$，$d$ 是模型宽度，$n$ 是 token 维度。
当 $d < n$ 时，训练损失存在不可约下界：
\[\mathcal{L}(g, \theta) \geq \sum_{i=d+1}^{n} \lambda_i\]其中 $\lambda_i$ 是协方差矩阵 $(\varepsilon - \mathbf{x})$ 的特征值。
直觉：如果模型宽度 $d$ 小于 token 维度 $n$，模型无法表示所有方向的去噪 → 存在不可消除的误差。
实验验证：
| 模型 \ 编码器 | DINOv2-S (384维) | DINOv2-B (768维) | DINOv2-L (1024维) |
|---|---|---|---|
| DiT-S (384维) | 3.6e-2 ✓ | 5.2e-1 ✗ | 6.5e-1 ✗ |
| DiT-B (768维) | 1.0e-3 ✓ | 2.4e-2 ✓ | 2.7e-1 ✗ |
| DiT-L (1024维) | 9.7e-4 ✓ | 1.3e-3 ✓ | 2.2e-2 ✓ |
模型宽度 ≥ token 维度才能收敛。
3. 维度相关的噪声调度偏移
token 维度越高，同样噪声水平下的信噪比越高（因为数据能量更大）。因此需要偏移噪声调度：
\[t_m = \frac{\alpha \cdot t_n}{1 + (\alpha - 1) \cdot t_n}\]其中 $\alpha = \sqrt{m/n}$，$n = 4096$（基准维度），$m$ = 有效数据维度（token 数 × 通道维度）。
| 配置 | gFID |
|---|---|
| 无偏移 | 23.075 |
| 有偏移 | 4.81 |
偏移后 FID 改善近 5 倍。
4. 噪声增强解码器训练
RAE 解码器在离散分布 $p(z)$（真实编码）上训练，但扩散模型的输出是连续的（含预测误差）。解决方案：训练时给编码加噪声：
\[p_n(z) = \int p(z - n) \cdot \mathcal{N}(0, \sigma^2 \mathbf{I})(n) \, dn\]| 其中 $\sigma \sim | \mathcal{N}(0, \tau^2) | $。 |
| 指标 | 无噪声增强 | 有噪声增强 |
|---|---|---|
| rFID（重建） | 0.49 | 0.57（略差） |
| gFID（生成） | 4.81 | 4.28（改善） |
5. DiT$^\text{DH}$ 架构：浅而宽的扩散头
高维潜在空间需要宽模型（定理 1），但标准 DiT 加宽会导致 FLOPs 二次增长。解决方案：
| $$z_t = M(\mathbf{x}_t | t, y) \quad \text{（DiT 主干提取特征）}$$ |
| $$v_t = H(\mathbf{x}_t | z_t, t) \quad \text{（扩散头预测速度）}$$ |
其中 $H$ 是仅 2 层、2048 维的浅宽网络。
为什么有效：主干 $M$ 处理语义理解（需要深度），扩散头 $H$ 处理维度匹配（需要宽度但不需要深度）→ 各取所长。
实验结果
ImageNet 256×256
| 方法 | Epochs | gFID（无引导） | gFID（有引导） | IS |
|---|---|---|---|---|
| DiT | 1400 | 9.62 | 2.27 | 121.5 |
| SiT | 1400 | 8.61 | 2.06 | 131.7 |
| REPA | 800 | 5.78 | 1.29 | 158.3 |
| DiT$^\text{DH}$（RAE） | 800 | 1.51 | 1.13 | 242.9 |
训练效率
| 对比 | 加速倍数 |
|---|---|
| vs 标准 DiT | 47× |
| vs REPA | 16× |
（在 80 epochs 内，同模型大小）
编码器选择消融（80 epochs）
| 编码器 | DiT-XL gFID | DiT$^\text{DH}$-XL gFID |
|---|---|---|
| DINOv2-S | 3.50 | — |
| DINOv2-B | 4.28 | 2.16 |
| DINOv2-L | 6.09 | — |
DINOv2-B 是最佳编码器选择——太小语义不足，太大维度过高。
高分辨率扩展（512×512）
| 方法 | gFID |
|---|---|
| EDM2（SOTA） | 1.25 |
| DiT$^\text{DH}$-XL | 1.13 |
个人思考
- “表征 > 压缩”的范式转换：传统 VAE 追求压缩效率，RAE 追求语义丰富性——重新定义了”好的潜在空间”是什么。
- 定理 1 的实践意义：模型宽度必须 ≥ token 维度——这解释了为什么 DiT 在高维 RAE 空间直接训练效果差（宽度不够）。
- DiT$^\text{DH}$ 的巧妙设计：深主干 + 浅宽头 → 分离语义理解和维度匹配 → 解决了高维空间的计算困境。
- 47× 训练加速说明语义潜在空间让扩散模型”更容易学”——因为去噪方向本身就有语义意义。
- 噪声调度偏移是容易被忽略但极其重要的细节——不做偏移 FID 差 5 倍。