深度研究：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 有三个根本缺陷：

缺陷	具体表现
架构过时	VAE 用 U-Net/CNN backbone，与现在主流的 Transformer 架构不一致
维度太低	典型 VAE 潜空间维度 < 64 channel，信息容量受限
表示能力弱	VAE 纯靠重建损失训练，学不到语义丰富的表示

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

---

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

RAE 的 radical insight 极其简洁：

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

RAE 架构

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

关键区别：

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

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

维度	VAE	RAE
Token 维度	~16-64 channel	768-1152 (匹配编码器输出)
信息容量	压缩损失	完整语义保留
语义结构	弱（纯像素级）	强（预训练模型已学）
与理解模型共享	否（需单独训练）	是（同编码器）

---

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

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

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

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

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

公式： $$ t_m = α × t_n / (1 + (α - 1) × t_n), 其中 α = √(m/n) $$

$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 性能结果

模型	分辨率	引导方式	FID
DiT^DH-XL (RAE)	256×256	无	1.51
DiT^DH-XL (RAE)	256×256	AutoGuidance	1.13
DiT^DH-XL (RAE)	512×512	AutoGuidance	1.13

→ 超越所有 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 图像 (FuseDiT)	39.3M	多样性
合成图像 (FLUX.1-schnell)	24.7M	风格一致性
文本渲染 (RenderedText)	9M	文本重建能力

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

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

重建性能对比 (rFID)：

模型	ImageNet	YFCC	Text
SDXL VAE	0.930	1.168	2.057
FLUX VAE	0.288	0.410	0.638
RAE (WebSSL)	0.388	0.558	1.372
RAE (SigLIP-2)	0.435	0.702	1.621

→ 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	加速比
GenEval	更快收敛	更慢	4.0×
DPG-Bench	更快收敛	更慢	4.6×

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

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

#### 微调结果

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

微调轮数	RAE GenEval	VAE GenEval
4 epochs	~52	~48
64 epochs	~53	~50
128 epochs	~52	~46 (过拟合)
256 epochs	~51	~42 (严重过拟合)

关键发现：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 目标
  - 维度感知噪声调度
  - 两阶段: 大规模预训练 → 高质量微调

模型配置

组件	选项
编码器	SigLIP-2 So400M (p14), DINOv2-B, WebSSL ViT-L
LLM	Qwen-2.5 (1.5B, 7B)
DiT	LightningDiT-based (0.5B, 2.4B, 5.5B, 9.8B)
图像分辨率	224×224 (RAE) / 256×256 (VAE baseline)
Visual tokens	256 (16×16)
Token 维度	1152 (SigLIP-2)

关键超参数

噪声调度 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 不再只是"看图说话"或"听写文字"，而是可以直接在 表示空间 中操作视觉内容。

---

八、资源汇总

资源	链接
v1 论文	https://arxiv.org/abs/2510.11690
v2 论文	https://arxiv.org/abs/2601.16208
GitHub 代码	https://github.com/bytetriper/RAE
v1 项目页	https://rae-dit.github.io/
v2 项目页	https://rae-dit.github.io/scale-rae/
模型权重	HuggingFace: `nyu-visionx/RAE-collections`
JAX/TPU 实现	https://github.com/bytetriper/diffuse_nnx

---

九、一句话总结

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

---

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

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