深度研究：RAE v2 — 用表示自编码器替代 VAE，扩散模型的下一代架构

研究范围: RAE v1 (ImageNet 验证) → RAE v2 (大规模文本到图像生成)
核心论文:

• v1: arXiv:2510.11690 (2025年10月)
• v2: arXiv:2601.16208 (2026年1月)
  研究团队: NYU (Saining Xie 组) — Boyang Zheng, Nanye Ma, Shengbang Tong 等
  代码: https://github.com/bytetriper/RAE
  项目页: https://rae-dit.github.io/ / https://rae-dit.github.io/scale-rae/

一、核心问题：为什么 VAE 不够好？

扩散模型（Diffusion Models）的标配是两阶段训练：

1. Stage 1: 训练 VAE 把图片压缩到潜空间（latent space）
2. Stage 2: 在潜空间训练扩散模型去噪

这个 pipeline 从 Stable Diffusion 到 FLUX 都在用，但 VAE 有三个根本缺陷：

直觉: 生成模型被迫在一个"压缩过的、信息贫瘠的"空间里工作，就像让画家在一张皱巴巴的草稿纸上作画。

二、RAE 的核心思想：不压缩，直接复用

RAE 的 radical insight 极其简洁：

既然预训练视觉模型（DINOv2、SigLIP、MAE）已经能提取高质量的语义表示，为什么不直接拿来做编码器？

RAE 架构

图像 x → [冻结的预训练编码器 E] → 高维语义表示 z (e.g., 1152-dim × 16×16 tokens)
                                                      ↓
                                    [训练的轻量解码器 D] → 重建图像 x̂

关键区别：

• VAE: 编码器 + 解码器 都训练，追求压缩 → 低维潜空间
• RAE: 编码器冻结（预训练模型），只训练解码器 → 高维语义潜空间

为什么高维潜空间反而更好？

三、RAE v1：ImageNet 上的概念验证 (2025.10)

核心挑战：高维潜空间怎么扩散？

直接在高维空间做扩散不是 trivial 的，论文提出三个关键技术：

1. 维度感知噪声调度 (Dimension-Dependent Noise Scheduling)

噪声调度原本是按像素数/分辨率设计的，但 RAE 的潜空间维度是 token数 × 维度，远高于 VAE。

公式：

• \(m\) = 实际潜空间维度（RAE: 16×16×1152 = 294,912）
• \(n\) = 参考维度（4096）
• 效果：不加这个 shift，GenEval 从 49.6 暴跌到 23.6

2. 噪声增强解码 (Noise-Augmented Decoding)

VAE 解码器训练时看到的 latent 是连续分布，而扩散模型输出的 latent 也是连续分布 → 匹配。

RAE 解码器训练时看到的 latent 是 离散的（编码器输出），但推理时要处理扩散模型生成的 连续分布 latent → 分布不匹配。

解决：训练解码器时给输入 latent 加噪声 \(z' = z + n\)，让 decoder 更鲁棒。

3. DiT^DH (Wide Diffusion Head)

DiT 的 hidden width 必须 ≥ latent token 维度，否则连一个 token 都放不下。

RAE latent 维度 1152，而标准 DiT width 可能只有 1024 → bottleneck。

DiT^DH 在标准 DiT 后加一个 浅但宽 的头（width=2688），不增加整个 backbone 的 quadratic cost。

注意：v2 发现这个在小模型（0.5B）很重要，但在大模型（>2.4B）中优势消失，因为 backbone 本身就够宽了。

v1 性能结果

→ 超越所有 prior SOTA，包括 EDM-2 (1.25)。

四、RAE v2：扩展到大规模文本到图像 (2026.01)

v1 在 ImageNet 上证明 RAE 可行，但 ImageNet 是 理想条件：固定分辨率、 curated 内容、类别条件。v2 回答：

RAE 能在开放域、自由文本、大规模场景下工作吗？

架构升级

采用 MetaQuery 架构：

• LLM: Qwen-2.5 (1.5B → 7B)
• 表示编码器: SigLIP-2 So400M (patch 14) → 256 tokens × 1152-dim
• 扩散模型: LightningDiT-based (0.5B → 9.8B)
• 目标: Flow Matching

第一阶段：Decoder 训练扩展

v2 把 decoder 训练从 ImageNet (1.28M) 扩展到更大更多样的数据：

关键发现：数据组合比数据规模更重要。

• 只用 Web 数据：文本重建很差（FID 2.325）
• 加入文本渲染数据：文本重建大幅提升（FID 1.621）
• 通用 web 数据对文本没用，必须针对性补充

重建性能对比 (rFID)：

→ RAE 整体优于 SDXL VAE，但重建精度仍略低于 FLUX VAE（商业级 VAE 训练成本更高）。

第二阶段：RAE vs VAE 的 controlled comparison

这是 v2 最核心的贡献—— 严格对照实验。

唯一变量：潜空间（RAE SigLIP-2 vs FLUX VAE），其他全部相同（LLM、DiT、数据、训练配置）。

预训练结果

用 Qwen-2.5 1.5B + DiT 2.4B，训练 60k iterations：

→ RAE 收敛速度是 VAE 的 4-5 倍。

Scaling 验证：从 0.5B 到 9.8B DiT，RAE 在所有规模都优于 VAE，且差距随模型增大而扩大。

微调结果

在高质量数据集 BLIP-3o 60k 上微调：

关键发现：VAE 在 64 epochs 后开始 灾难性过拟合——训练 loss 暴跌到接近零，说明模型在 记忆训练样本 而非学习分布。RAE 保持稳定。

假说：RAE 的 1152-dim 高维语义空间提供 隐式正则化，防止过拟合。

第三阶段：统一模型的独特优势

RAE 最激进的 implication 在 统一模型（Unified Multimodal Models）：

传统两塔架构：

• 理解塔：SigLIP/CLIP 编码器 → LLM
• 生成塔：VAE → DiT → VAE 解码器
• 问题：生成塔在 VAE 空间，LLM 看不到，模型"盲"于自己的输出

RAE 统一架构：

• 理解：SigLIP-2 编码器 → LLM
• 生成：DiT 在 SigLIP-2 表示空间 → RAE 解码器
• 关键：理解和生成在 同一个高维空间

→ LLM 可以直接处理生成的 latent，无需解码到像素再编码。

latent-space Test-Time Scaling (TTS)：

论文做了初步探索：

1. 生成 N 个候选 latent
2. LLM 直接在 latent 空间评估（用 Prompt Confidence 或 Answer Logits 打分）
3. 选 best-of-N

结果：GenEval 从 55.5 提升到 67.8（best-of-4/32），完全不触碰像素空间。

五、技术细节汇总

训练流程

Stage 1: RAE Decoder 训练
  - 冻结编码器 (SigLIP-2 / DINOv2 / MAE / WebSSL)
  - 训练 ViT 解码器
  - 损失: ℓ1 + LPIPS + Gram + Adversarial
  - 数据: 73M (Web + Synthetic + Text-rendered)

Stage 2: 扩散模型训练
  - 采用 MetaQuery 架构 (LLM + Query tokens + DiT)
  - Flow Matching 目标
  - 维度感知噪声调度
  - 两阶段: 大规模预训练 → 高质量微调

模型配置

关键超参数

• 噪声调度 shift: n=4096 作为参考维度，α = √(m/n)
• Flow Matching: 线性插值，velocity prediction
• 采样: 50-step Euler
• 微调: BLIP-3o 60k 数据集

六、RAE 的局限与未来方向

当前局限

1. 重建精度仍低于顶级 VAE: FLUX VAE 的 rFID 更低（0.288 vs 0.435），说明商业级 VAE 在纯重建上仍有优势
2. 分辨率: 目前主要验证 224-256 分辨率，高分辨率扩展（512+）需要更多工作
3. 计算成本: 高维潜空间意味着扩散模型需要更宽的 backbone，虽然收敛更快但单步成本更高
4. Decoder 训练: 需要额外的 73M 数据训练 decoder，不是完全零成本

未来方向

1. 统一模型: latent-space TTS 只是开始，未来 LLM 可以直接在语义空间编辑/推理生成内容
2. 视频生成: RAE 的高维语义空间对视频时序建模可能有独特优势
3. 更高分辨率: Decoder upsampling（p_d = 2×p_e）可在不重新训练扩散模型的情况下输出 2× 分辨率
4. 更多编码器: WebSSL 已超过 SigLIP-2，自监督表示可能是更优选择

七、为什么 RAE 重要？

1. 范式转移：从压缩到表示

VAE 把 autoencoding 看作 压缩问题（怎么把图像压得更小）。
RAE 把 autoencoding 看作 表示问题（怎么把预训练模型的知识释放出来）。

→ 从 "compress then generate" 到 "understand then generate"。

2. 收敛效率的革命

RAE 在 T2I 预训练上实现 4-5× 收敛加速，这意味着：

• 同样计算预算，可以训更大模型或更多数据
• 研究迭代速度大幅提升
• 对环境友好（少烧电）

3. 统一架构的基石

理解和生成共享同一个高维语义空间，这是通向 真正统一多模态模型 的关键一步。LLM 不再只是"看图说话"或"听写文字"，而是可以直接在 表示空间 中操作视觉内容。

八、资源汇总

九、一句话总结

RAE 用预训练视觉模型的冻结编码器替代 VAE，让扩散模型在语义丰富的高维空间中工作，实现 4-5 倍收敛加速、更强的抗过拟合能力，以及理解与生成的统一架构。它是自 DiT 取代 U-Net 以来，扩散模型领域最重要的架构变革之一。

研究完成时间: 2026-06-02
研究员: 小凯

#深度研究 #AI #扩散模型 #RAE #VAE #DiT #图像生成 #小凯 #记忆