← 返回 2026-01-23

用表示自编码器扩展文本到图像的扩散 Transformer Scaling Text-to-Image Diffusion Transformers with Representation Autoencoders

Shengbang Tong, Boyang Zheng, Ziteng Wang, Bingda Tang, Nanye Ma, Ellis Brown, Jihan Yang, Rob Fergus, Yann LeCun, Saining Xie 📅 2026-01-22 👍 55 2026-07-13 08:35
Transformer 扩散模型 文生图 自编码器 表示学习

RAE 用高维语义空间替代 VAE 低维压缩,让扩散模型收敛更快、生成更好且更抗过拟合

前置知识

变分自编码器 (VAE)

VAE 是一种生成模型,由编码器和解码器组成。编码器将高维输入(如图像)压缩到低维潜在空间(通常 4-8 通道,分辨率下采样 8 倍),解码器从潜在表示重建图像。在 Stable Diffusion 等文生图模型中,VAE 负责将像素空间压缩为紧凑的潜在表示,扩散模型在这个低维空间中工作,大幅降低计算成本。但压缩不可避免地丢失信息,尤其在文字、细节纹理等高频信息上。

本文的核心论点是 VAE 的低维压缩空间是瓶颈,RAE 用高维语义空间替代之,理解 VAE 的工作原理和局限是理解本文创新点的前提。

扩散 Transformer (DiT)

DiT 将 Transformer 架构引入扩散模型,替代传统的 U-Net。它将图像的潜在表示分割为 patch 序列,通过 Transformer 的自注意力机制进行去噪。DiT 在扩展性上优于 U-Net:模型参数从 0.5B 到 10B+ 都能稳定提升。Flow matching 是 DiT 常用的训练目标,通过学习从噪声到数据的向量场来生成图像。

本文使用 DiT 作为生成骨干网络,研究 RAE 潜在空间下 DiT 的扩展行为,理解 DiT 的架构特性对分析实验结果至关重要。

表示学习 / 视觉编码器

表示学习旨在训练编码器将图像映射到语义丰富的高维特征空间。代表模型包括 CLIP(语言监督)、DINO/DINOv2(自监督)、SigLIP(两者结合)。这些编码器输出的特征维度通常为 768-1152,每个 patch 一个向量,16x16 的 patch 网格产生 256 个 token。这些特征擅长语义理解任务(分类、检索、VQA),但此前被认为过于抽象而不适合生成。

RAE 的核心是用这种语义编码器(如 SigLIP-2)替代 VAE 编码器作为扩散模型的潜在空间,理解表示编码器的特性是理解 RAE 工作原理的关键。

表示自编码器 (RAE)

RAE 由 Zheng et al. (2025) 提出,核心思路是:冻结一个预训练的视觉表示编码器(如 SigLIP-2),只训练一个轻量解码器从高维语义特征重建像素。这产生了一个特殊的自编码器——编码器是强大的语义模型,解码器学习从语义到像素的映射。扩散模型在这个语义潜在空间中工作,而非 VAE 的压缩空间。在 ImageNet 上,RAE 已展示出比 VAE 更快的收敛和更好的生成质量。

本文是 RAE 从 ImageNet 类条件生成扩展到大规模文生图的关键工作,理解 RAE 的基本框架是理解本文所有实验设计的基础。

维度感知噪声调度

标准扩散模型的噪声调度在低维潜在空间设计,当潜在空间维度大幅增加时(如 RAE 的 256x1152 vs VAE 的 256x4),需要调整噪声调度。RAE 提出根据有效数据维度 m = N x d(token 数 x token 维度)对扩散时间步进行缩放。这确保模型在高维空间中仍能有效学习去噪。

实验表明这是 RAE 在 T2I 场景中唯一不可省略的设计选择,理解其数学原理对复现和改进 RAE 至关重要。

研究动机

当前主流文生图模型(如 Stable Diffusion、FLUX)依赖 VAE 将图像压缩到低维潜在空间(通常 4-16 通道,空间下采样 8 倍),然后在这个紧凑空间中训练扩散模型。这种设计虽然计算高效,但存在根本性瓶颈:压缩过程不可避免地丢失高频信息,特别是文字字形、精细纹理和空间细节。VAE 编码器本身只是一个简单的卷积网络,其学到的表示缺乏语义结构,这意味着扩散模型需要在语义贫乏的空间中学习复杂的生成分布。此外,在统一多模态模型的语境下,理解模块使用高维 CLIP/SigLIP 特征,生成模块使用低维 VAE 潜在空间,两个空间完全割裂——模型无法直接理解自己生成的内容,必须先解码回像素再重新编码,这在架构上是不优雅的,在效率上也是浪费的。

本文的目标是本文的目标是将 RAE 框架从 ImageNet 类条件生成扩展到大规模自由形式文生图(T2I)场景,验证高维语义潜在空间是否能在更复杂的生成任务中保持优势。具体而言,作者希望回答三个核心问题:(1)RAE 解码器能否在比 ImageNet 更多样化的数据上有效训练,特别是处理文字渲染等高频细节?(2)RAE 在 ImageNet 上提出的设计选择(维度感知噪声调度、噪声增强解码、宽 DDT 头)在大规模 T2I 训练中是否仍然必要?(3)RAE 在与 SOTA VAE(FLUX-VAE)的公平对比中,能否在收敛速度和最终生成质量上全面胜出?此外,作者还探索了一个更长远的目标:利用 RAE 的共享潜在空间实现统一多模态模型中的潜在空间测试时缩放(Test-Time Scaling)。

与已有工作不同的是,此前的 RAE 工作仅在 ImageNet 上验证——这是固定分辨率、精心策划内容、类条件生成的理想场景。而文生图涉及更广泛的视觉多样性、开放式组合、更大的模型和计算量,高维潜在扩散在这种场景下的有效性完全未经验证。更重要的是,此前没有人系统研究过 RAE 设计选择在大规模场景中的必要性——原论文的所有消融都在 ImageNet 小模型上进行,无法区分哪些是高维扩散的根本需求,哪些只是小规模训练的权宜之计。本文的独特切入角度是:通过大规模 T2I 实验,揭示 scale 是一个简化器——它淘汰了不必要的复杂性,只保留真正本质的设计。此外,本文首次展示了 RAE 的共享潜在空间如何在统一模型中实现无需解码的测试时缩放,这是一个此前被忽视的应用方向。

核心方法

本文的方法可以用一个简单的类比理解:传统 VAE 就像把一幅画压缩成一张模糊的缩略图,然后在这个缩略图上作画;RAE 则是把画转换成一组语义描述(比如'左上角有一只红色的鸟,背景是蓝色天空'),然后在这个语义描述空间中作画。后者虽然维度更高,但语义更丰富,模型学习起来更容易。技术路线分为三个阶段:(1)解码器扩展训练——在冻结的 SigLIP-2 编码器上训练解码器,从 256 个 1152 维的语义 token 重建像素,训练数据从 ImageNet 扩展到 Web + 合成 + 文字渲染的 73M 混合数据;(2)T2I 扩散模型训练——使用 MetaQuery 架构,将预训练 LLM(Qwen-2.5)与 DiT 通过可学习的 query token 连接,在 RAE 潜在空间中进行 flow matching 训练;(3)统一模型探索——利用 RAE 共享的语义空间,让 LLM 直接理解和评估扩散模型生成的潜在表示,无需解码回像素。

本文的核心创新点不是提出新的架构或训练技巧,而是通过大规模实验揭示了一个关键洞察:RAE 的高维语义潜在空间本身就是最好的设计,不需要额外的复杂性。具体而言,原 RAE 论文提出了三个专门设计来应对高维扩散的挑战:维度感知噪声调度、噪声增强解码、宽 DDT 头。本文发现,在 T2I 规模下,只有维度感知噪声调度是不可或缺的(去掉它 GenEval 从 49.6 暴跌到 23.6),而另外两个设计的收益随模型规模增大而消失——噪声增强解码只在训练早期有效,DiTDH 只在 0.5B 小模型上有明显优势(+11.2 GenEval),到 2.4B 以上就无关紧要了。这意味着 scale 本身就是一个简化器:当模型容量足够大时,标准 DiT 架构配合适当的噪声调度就够了。这比原 RAE 论文的配方更简洁,也更适合大规模训练。

方法步骤详情

第一步:RAE 解码器扩展训练。输入是 SigLIP-2 ViT-So400M 编码器产生的 256 个 1152 维 token(来自 224x224 图像的 16x16 patch 网格),使用 ViT-XL 解码器(28 层,隐藏维度 1152,16 头)重建像素。训练目标包含 L1 损失、LPIPS 感知损失、Gram 损失和对抗损失,其中对抗损失使用冻结的 DINO-S/16 作为判别器。训练数据为 73M 的混合数据(Web 39M + FLUX 合成 24.7M + RenderedText 文字数据)。第二步:T2I 预训练。使用 MetaQuery 架构,在文本 prompt 前拼接 256 个可学习 query token,由 Qwen-2.5 LLM 联合处理文本和 query,产出的 query 表示通过 2 层 MLP 投影到 DiT 的输入空间。DiT 使用 flow matching 目标训练,预测速度场,推理使用 50 步 Euler 采样器。预训练数据约 39.3M Web 图像 + 24.7M 合成图像。第三步:高质量数据微调。在 BLIP-3o 60k 数据集上微调,同时更新 LLM 和 DiT。第四步(可选):潜在空间测试时缩放。生成多个候选潜在表示,用 LLM 作为验证器直接在特征空间打分,选择最佳结果。

技术新颖性

本文的技术新颖性体现在三个层面。首先,在框架层面,本文首次将 RAE 从 ImageNet 类条件生成扩展到大规模自由形式 T2I,证明了高维语义潜在空间在复杂生成场景中的有效性,这打破了此前高维潜在被认为过于抽象而无法有效生成的共识。其次,在方法论层面,本文通过系统消融揭示了 scale 作为简化器的规律——维度感知噪声调度是唯一不可省略的组件,而噪声增强解码和 DiTDH 都是小模型的补丁而非本质需求。这种通过大规模实验来验证设计必要性的方法论本身就有价值。第三,在应用层面,本文首次展示了 RAE 共享潜在空间在统一多模态模型中的独特优势:LLM 可以直接理解扩散模型生成的潜在表示,实现无需解码-重编码的测试时缩放。这在 VAE 体系下是不可能的,因为 VAE 潜在对 LLM 来说是不透明的。

在更多数据上训练的 RAE 解码器能跨域泛化
Figure 2: 在更多数据上训练的 RAE 解码器能跨域泛化
训练流程概览
Figure 3: 训练流程概览
在 T2I 规模下饱和的设计选择
Figure 4: 在 T2I 规模下饱和的设计选择
潜在空间中的测试时缩放
Figure 8: 潜在空间中的测试时缩放

实验结果

本文的实验结果系统地验证了 RAE 在大规模 T2I 中的优势,核心发现如下:(1)收敛速度优势:在 Qwen-2.5 1.5B + DiT 2.4B 的设置下,RAE 在 GenEval 上实现 4.0x 加速,在 DPG-Bench 上实现 4.6x 加速——这意味着达到相同生成质量,RAE 只需要 VAE 约 1/4 的训练迭代。(2)跨规模一致性优势:从 0.5B 到 9.8B 的 DiT 规模,从 1.5B 到 7B 的 LLM 规模,RAE 始终优于 VAE。即使在最小的 0.5B DiT(网络宽度仅略超过 RAE 潜在维度),RAE 仍有明显优势。(3)微调稳定性优势:VAE 模型在 64 个 epoch 后开始灾难性过拟合——训练损失骤降到接近零,GenEval 性能急剧下降;RAE 模型在 256 个 epoch 仍保持稳定,甚至延长到 512 个 epoch 也只有轻微下降。训练损失分析显示 VAE 损失迅速降到极低值(约 0.06),而 RAE 损失更缓慢下降并稳定在较高值(约 0.16),表明高维语义空间提供了隐式正则化。(4)数据组合效应:合成数据(45.1 GenEval)显著优于 Web 数据(25.9),两者的组合(49.5 GenEval)甚至超过两倍合成数据(48.0),说明数据的互补性比数量更重要。(5)潜在空间测试时缩放:在 1.5B LLM + 5.5B DiT 设置下,从 8 个候选中选 4 个最佳(best-of-4/8),使用 Answer Logits 验证器可将 GenEval 从 53.2 提升到 59.6(+6.4);在 7B LLM 下进一步提升到 67.8(+12.3)。这证明了潜在空间 TTS 的可行性和有效性。(6)理解性能不受影响:添加生成训练不会损害视觉理解性能,RAE 和 VAE 在 MME、TextVQA、AI2D 等基准上表现相当。

数据对 RAE 重建保真度很重要
Table 1: 数据对 RAE 重建保真度很重要
数据组成比规模更重要
Table 3: 数据组成比规模更重要
不同 LLM-DiT 配置下的测试时缩放结果
Table 5: 不同 LLM-DiT 配置下的测试时缩放结果
生成训练不影响理解能力;RAE 和 VAE 表现相当
Table 6: 生成训练不影响理解能力;RAE 和 VAE 表现相当
RAE 在文生图预训练中比 VAE 收敛更快
Figure 1: RAE 在文生图预训练中比 VAE 收敛更快
RAE 在不同 LLM 和 DiT 规模下均优于 VAE
Figure 5: RAE 在不同 LLM 和 DiT 规模下均优于 VAE
RAE 模型优于 VAE 且更抗过拟合
Figure 6: RAE 模型优于 VAE 且更抗过拟合
查看结构化数据
任务指标本文基线提升
文生图预训练收敛速度 GenEval RAE: 49.6(30k 步) VAE (FLUX): ~12.4(30k 步) 4.0x 收敛加速
文生图预训练收敛速度 DPG-Bench RAE: 76.8(30k 步) VAE (FLUX): ~16.7(30k 步) 4.6x 收敛加速
文生图生成质量(1.5B LLM + 2.4B DiT) GenEval RAE: 49.5 VAE (FLUX): 39.6 +9.9 绝对提升
文生图生成质量(1.5B LLM + 2.4B DiT) DPG-Bench RAE: 76.9 VAE (FLUX): 70.5 +6.4 绝对提升
图像重建质量 rFID-50k (ImageNet) RAE (SigLIP-2): 0.435 SDXL-VAE: 0.930; FLUX-VAE: 0.288 优于 SDXL-VAE 53%,略逊 FLUX-VAE
潜在空间测试时缩放 GenEval (7B LLM + 5.5B DiT) best-of-4/32 + Answer Logits: 67.8 无 TTS: 55.5 +12.3 绝对提升

局限与改进

尽管本文的结果令人印象深刻,但仍存在几个重要局限。首先,RAE 的重建质量仍不及 FLUX-VAE:在 ImageNet 上 rFID 为 0.435 vs FLUX-VAE 的 0.288,在 YFCC 上为 0.702 vs 0.410。这意味着 RAE 解码器在像素级保真度上仍有差距,特别是在高频细节的重建上。其次,文字重建高度依赖特定训练数据:没有文字渲染数据时,解码器完全无法生成清晰的文字字形,这说明 RAE 解码器的泛化能力对数据分布非常敏感。第三,实验在 224x224 分辨率下进行,而实际文生图应用通常需要 512x1024 或更高分辨率,RAE 在高分辨率下的表现未经验证。第四,RAE 的推理成本可能更高:256 个 1152 维 token 的计算量远大于 VAE 的 256x4 维潜在,虽然本文没有报告推理时间对比,但这在实际部署中可能是问题。第五,DiT 规模超过 6B 后性能趋于平稳,说明简单增加模型规模而不提升数据质量和多样性会导致容量利用不足。最后,统一模型的探索(潜在空间 TTS)只是初步尝试,缺乏深入分析和与像素空间 TTS 的公平对比。

独立分析的弱点

从独立分析的角度,本文存在几个值得关注的弱点。第一,重建质量差距被低估:虽然 RAE 在生成任务上优于 VAE,但其重建 rFID 仍明显落后于 FLUX-VAE(0.435 vs 0.288)。这意味着 RAE 解码器本身是一个瓶颈——如果重建质量能提升到 FLUX-VAE 水平,生成质量可能进一步提高。改进方向包括使用更强的解码器架构(如更大的 ViT 或 U-Net 混合架构)、更先进的重建损失函数、以及渐进式训练策略。第二,文字重建的数据依赖性是一个工程隐患:论文承认没有文字数据就无法重建文字,但没有量化需要多少文字数据、什么类型的文字数据最有效。实际应用中,文字场景极其多样(不同字体、大小、语言、排版),当前的 RenderedText 数据集可能不足以覆盖真实分布。改进方向是设计更系统的文字数据采集和增强策略。第三,缺乏推理效率对比:RAE 的潜在维度是 VAE 的约 288 倍(1152 vs 4),这在理论上会增加 DiT 的计算量和内存占用,但论文完全没有讨论推理效率。这是一个重要的实际考量,需要补充推理延迟、内存占用、吞吐量等指标。第四,评估指标的局限性:GenEval 和 DPG-Bench 主要评估语义对齐,但对图像质量、多样性、细节保真度的评估不足。需要补充 FID、IS、CLIP Score 等更全面的指标。

未来方向

基于本文的成果,有几个值得探索的未来方向。首先,更高分辨率的 RAE 训练:当前实验在 224x224 分辨率下进行,扩展到 512x512 或 1024x1024 是必要的实际需求。这可能需要设计多尺度解码器或使用分层潜在空间。其次,探索更强的表示编码器:本文主要使用 SigLIP-2 和 WebSSL,但最新的 DINOv3、InternVL 等编码器可能提供更好的语义表示。特别是,专门针对生成任务优化的表示编码器(而不仅仅是理解任务)是一个有前景的方向。第三,RAE 与离散 token 的结合:将 RAE 的连续高维潜在空间量化为离散 token,可以与自回归生成范式(如 Chameleon、Emu3)结合,实现真正的统一多模态建模。第四,潜在空间 TTS 的深入探索:本文只是初步展示了 LLM 可以在潜在空间验证生成结果,但更系统的研究——包括不同验证策略、不同选择算法、计算-质量权衡分析——仍有很大空间。第五,RAE 在视频生成中的应用:视频生成的潜在空间设计面临类似的信息保留 vs 计算效率的权衡,RAE 的高维语义潜在可能在时序建模中也有优势。

复现评估

本文在可复现性方面做出了很好的承诺:作者明确表示将发布所有代码、数据和模型检查点。这包括(1)预训练数据:来自 FuseDiT 的 Web 数据(约 39.3M 图像)、FLUX.1-schnell 生成的合成数据(24.7M)、RenderedText 文字数据;(2)模型权重:SigLIP-2 编码器、RAE 解码器、各规模 DiT 和 LLM 的检查点;(3)训练代码:基于 TorchXLA 的 TPU 训练代码。然而,复现成本是一个重要考量:实验使用 TPU v4、v5p 和 v6e 集群,最大规模的 9.8B DiT 训练需要大量 TPU 资源。对于大多数研究团队来说,完整复现所有实验可能不现实,但复现核心结果(如 2.4B DiT 的 RAE vs VAE 对比)在 8xA100 或类似配置下应该是可行的。解码器训练(ViT-XL,80 个 epoch)在单机多卡下可在数天内完成。数据方面,Web 数据和合成数据需要自行准备,但 RenderedText 是公开数据集。总体而言,复现难度为中等偏高——核心方法不复杂,但完整实验需要显著的计算资源。