Improved Baselines with Representation Autoencoders

会议: NeurIPS 2026
arXiv: 2605.18324
领域: 图像生成
关键词: 表征自编码器, 扩散Transformer, 多层特征聚合, 表征对齐, 内置自导引

一句话总结

RAEv2 针对表征自编码器（RAE）重建质量受限、不兼容 CFG 需额外模型导引以及与 REPA 关系不明等瓶颈，提出无参数多层求和特征聚合、揭示 RAE 与 REPA 的互补正则机制，并将 REPA 头重参数化为零开销的潜空间自导引，在 ImageNet-256 上仅用 80 轮训练便达到 1.06 gFID 与 2.17 \(\text{FD}_r^6\)，收敛速度提升逾 10 倍。

研究背景与动机

近年来，表征自编码器（Representation Autoencoder, RAE）通过直接采用预训练视觉编码器（如 DINOv2）替代传统的变分自编码器（VAE），构建起连续语义潜空间，在将图像理解与图像生成统一到同一表征空间的探索上迈出了重要一步。然而，RAE 在走向实际应用时面临三大根本性瓶颈：首先，其图像重建质量显著落后于专有 VAE，而以往工作多将视觉编码器的最后一层特征作为唯一表征输入，忽略了深层网络中空间几何细节向顶层高阶语义退化的问题，通常需要引入复杂的微调或人脸/文本等专项数据才能弥补重建缺陷。

与此同时，社区对 RAE 与表征对齐技术（Representation Alignment, REPA）的结合机制存在长期误解，普遍假设既然 RAE 已将预训练编码器作为输入端，再向 Diffusion Transformer 的中间层蒸馏相同表征就是冗余的“跳跃连接”。此外，RAE 与传统的无分类器导引（Classifier-Free Guidance, CFG）高度不相容，导致其不得不借助 AutoGuidance（AG）额外训练一个较弱的从属扩散模型来充当负向基线，这不仅加剧了训练计算负担与管线复杂度，还使得推理阶段的函数评估次数（NFE）难以压降。

本文直面上述设计假设中的盲区，从特征层级抽象、表征对齐的深层作用机理与导引范式重构三个维度切入。核心 idea：将编码器表征泛化为深部多层特征的逐元素求和以无损提升重建质量，利用 REPA 对中间层空间注意力结构的互补正则释放更强编码器的生成潜力，并将 REPA 投影头重参数化为潜空间 \(x\)-prediction 以实现无需二次前向与额外模型的单次推理内置自导引。

方法详解

整体框架

RAEv2 构建了一个从图像表征提取、流匹配扩散建模到高效解码生成的统一管线。首先，输入图像经冻结的预训练视觉编码器提取特征，通过无参数的多层求和（Multi-Layer Sum, MLS）聚合后 \(K\) 层特征，形成保持原始时空尺寸与通道维度的紧凑潜空间表征；其次，在基于 Flow Matching 的 DiTDH-XL 扩散骨干网络中，浅层特征（如第 8 层）同时接入 REPA 线性投影头，该头被显式约束为在潜空间中执行直接去噪（\(x\)-prediction）；最后，在推理阶段，全模型输出与浅层 REPA 头的 \(x\)-prediction 预测值直接做差值外推，形成完全内生于同一前向计算的自导引信号，再映射回速度场以求解 ODE 轨迹。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    IMG["输入图像 I"] --> ENC["预训练视觉编码器<br/>(DINOv3-L, 24层)"]
    ENC --> D1["多层特征求和聚合<br/>提取后 K 层特征并逐元素求和 x=∑z_l"]
    D1 --> LAT["紧凑潜空间表征 x ∈ R^{N×d}"]
    LAT --> DEC["RAE 解码器<br/>重建高质量图像"]
    LAT --> FLOW["Flow Matching 加噪潜变量 x_t"]
    FLOW --> DIT_EARLY["DiTDH-XL 浅层 Transformer<br/>提取中间特征 h (第8层)"]
    DIT_EARLY --> D2["RAE 与 REPA 互补协同<br/>REPA 线性头 h_phi 预测 clean latent x_hat_repa"]
    DIT_EARLY --> DIT_LATE["DiTDH-XL 深层 Transformer<br/>全网络预测 x_hat_full"]
    D2 --> D3["基于 x-prediction 的零开销内置自导引<br/>x_hat_guided = x_hat_full + w·(x_hat_full - x_hat_repa)"]
    DIT_LATE --> D3
    D3 --> OUT["转换为速度 v_t 单步更新 / 生成结果"]

关键设计

1. 多层特征求和聚合：突破顶层表征的重建瓶颈

针对原有 RAE 仅抽取最后一层特征导致底层纹理与几何轮廓丢失、重建质量受限的问题，本设计探究如何在不膨胀潜变量体积、不增加训练参数的前提下吸收深浅层互补信号。若沿序列长度方向拼接 \(K\) 层特征，潜变量序列长度将暴增为 \(KN\)，大幅增加 DiT 自注意力算力开销；若沿通道维度拼接，潜变量通道数增至 \(Kd\)，严重恶化扩散模型优化难度。

RAEv2 采用极其简洁的逐元素求和策略（Multi-Layer Sum, MLS），将编码器最后 \(K\) 层的 token 特征直接相加：

\[\boldsymbol{x} = \sum_{\ell=L-K+1}^{L} \boldsymbol{z}_{\ell} \in \mathbb{R}^{N \times d}\]

在高维表征空间中，线性相加能够高度保持底层子空间的几何结构。相较于采用固定随机矩阵降维的多层随机投影（MLR），MLS 完全没有随机投影带来的子空间畸变与方差扩散。实验表明，随着 \(K\) 增大（例如 DINOv3-L 取 \(K=23\)），重建指标单调跃升（rFID 由 0.60 降至 0.18，PSNR 提升至 27.03 dB），且在 ImageNet 线性探测任务上依然维持约 85.2% 的顶层语义分类精度，真正实现了理解与生成两用表征的无损聚合。

2. RAE 与 REPA 互补协同：解耦语义潜空间与空间注意力正则

针对“RAE 输入端已有预训练表征，中间层再做 REPA 蒸馏纯属浪费跳跃连接”的流行假设，作者在 27 种预训练编码器上系统分析了生成质量（gFID）与表征特性的皮尔逊相关性。分析证实 RAE 与 REPA 的底层驱动机制完全正交：RAE 的性能增益高度由表征的全局语义能力决定（线性探测精度 LP 与 gFID 的相关系数高达 \(|r|=0.81\)，而局部距离相似度 LDS 相关系数仅 \(|r|=0.13\)）；相反，单独的 REPA（配合传统 VAE）则几乎全依赖空间结构正则（LDS 相关系数 \(|r|=0.89\)），对 LP 甚至呈微弱负相关。

将两者叠加时，RAE 负责提供强语义潜空间，而施加于 DiT 第 8 层的 REPA 投影损失（权重 \(\lambda=0.5\)）则强制约束生成网络内部 token 之间的空间自相似性结构。由于两机制互补，RAE + REPA 与 LP 与 LDS 综合分数的负相关性高达 \(|r|=0.83\)。这一发现彻底打破了旧有认知，不仅在全量模型族中普适提升生成质量，更使得像 DINOv3-L 这样在全局语义与空间结构上均属顶级的大模型得以在生成任务中完全释放潜力，显著优于原先在 RAE 中偏向选择的 DINOv2-B。

3. 基于 x-prediction 的零开销内置自导引：替代二次模型与额外前向

传统 RAE 无法直接受益于标准 CFG，以往必须借助 AutoGuidance 独立训练一个较弱的从属扩散模型提供负向梯度，显著抬高了显存占用与训练维护成本。RAEv2 敏锐地抓住了 RAE 潜空间本身的独特性质：由于潜空间直接对应视觉表征本身，即真实目标就是 \(\boldsymbol{x} = E(\boldsymbol{I})\)，挂载在 DiT 浅层的 REPA 线性头 \(h_\phi(\boldsymbol{h})\) 其本质就是在执行直接重构干净潜变量的预测任务（即 \(x\)-prediction）：

\[\hat{\boldsymbol{x}}_{\text{repa}} = h_{\phi}(\boldsymbol{h})\]

因为投影头仅依赖网络第 8 层的早熟特征，其重构表征的能力天然劣于主网络全层输出。作者将主干网络的输出头也重参数化为直接预测干净潜变量 \(\hat{\boldsymbol{x}}_{\text{full}}\)，从而在完全相同的表征维度下将 REPA 预测作为天然的弱基线，实施内置自导引（Internal Guidance）：

\[\hat{\boldsymbol{x}}_{\text{guided}} = \hat{\boldsymbol{x}}_{\text{full}} + w \cdot (\hat{\boldsymbol{x}}_{\text{full}} - \hat{\boldsymbol{x}}_{\text{repa}})\]

最终只需通过公式 \(\boldsymbol{v} = (\boldsymbol{x}_t - \hat{\boldsymbol{x}}_{\text{guided}})/t\) 转换回流匹配所需的速度向量，即可驱动标准 ODE 采样器。这一设计将弱模型的输出计算无缝整合进主网络单次前向传播中，既彻底免去了训练二级弱模型的必要，又避免了标准 CFG 必须跑两次前向（有条件与无条件）带来的 NFE 翻倍问题，实现了真正“免费”的高质量导引。

一个完整示例

以 ImageNet-256 条件图像生成为例，推理流程如下： 1. 采样初始高斯噪声潜变量 \(\boldsymbol{x}_1 \sim \mathcal{N}(\mathbf{0}, \boldsymbol{I})\)，形状为 \(16 \times 16 \times 1024\)。 2. 类别标签嵌入为 8 个 token，时间步 \(t\) 通过高斯傅里叶变换映射为 4 个 token，与当前潜变量拼接后输入 DiTDH-XL。 3. 数据流经前 8 个 Transformer 块后，中间隐藏特征 \(\boldsymbol{h}\)（维度 1152）经线性投影头直接输出早熟去噪估计 \(\hat{\boldsymbol{x}}_{\text{repa}}\)。 4. 特征继续流经剩余 20 个 Transformer 块及 2 层 DDT 解码器，输出主网络全量去噪估计 \(\hat{\boldsymbol{x}}_{\text{full}}\)。 5. 在同一时刻执行外推导引计算 \(\hat{\boldsymbol{x}}_{\text{guided}} = \hat{\boldsymbol{x}}_{\text{full}} + w \cdot (\hat{\boldsymbol{x}}_{\text{full}} - \hat{\boldsymbol{x}}_{\text{repa}})\)，并换算为瞬时速度 \(\boldsymbol{v}\)。 6. 使用 Euler 求解器走 50 步 ODE 积分生成最终表征 \(\boldsymbol{x}_0\)，送入冻结的 RAE 解码器，还原出高质量 \(256 \times 256\) RGB 图像。

损失函数 / 训练策略

模型采用连续时间 Flow Matching 目标，时间步采样遵循 Logit-Normal 分布。主损失为速度场匹配的均方误差，浅层挂载 REPA 损失：

\[\mathcal{L} = \mathcal{L}_{\text{FM}}(\boldsymbol{v}_{\theta}, \boldsymbol{v}_{t}) + \lambda \cdot \mathcal{L}_{\text{REPA}}(\hat{\boldsymbol{x}}_{\text{repa}}, \boldsymbol{x})\]

其中 \(\lambda=0.5\)，REPA 目标与主干网络均以 \(\boldsymbol{x}\)-prediction 形式对齐。优化采用 AdamW，全局批大小 1024，基础学习率 \(2 \times 10^{-4}\)，在 25 个 epoch 预热后线性衰减至 \(2 \times 10^{-5}\)（第 50 轮完成衰减），无权重衰减，采用 bfloat16 混合精度与梯度截断（范数 1.0），EMA 衰减率设定为 0.9995。

实验关键数据

主实验

在 ImageNet-256 类别条件生成基准上，作者引入了训练效率新指标 \(\text{EP}_{\text{FID@2}}\)（无导引 gFID 达到 \(\le 2\) 所需训练轮数），并在六个独立特征空间（Inception, ConvNeXt, DINOv2, MAE, SigLIP, CLIP）评估表征 Fréchet 距离（\(\text{FD}_r\)）及其几何平均 \(\text{FD}_r^6\)。

方法	训练轮数	\(\text{EP}_{\text{FID@2}} \downarrow\)	导引 gFID \(\downarrow\)	Incep. \(\downarrow\)	DINOv2 \(\downarrow\)	CLIP \(\downarrow\)	\(\text{FD}_r^6 \downarrow\)
SiT-XL/2	800	\(\gg 800\)	2.12	1.26	7.89	17.69	8.44
DDT-XL	800	–	1.26	0.75	4.26	13.86	5.70
SiT-XL/2-REPA	800	\(\gg 800\)	1.42	0.85	4.27	12.65	5.45
LightningDiT	800	\(\gg 800\)	1.42	0.85	3.76	10.21	4.57
REG	800	560	1.54	0.92	3.45	10.86	4.64
REPA-E	800	480	1.12	0.70	2.44	6.28	3.04
RAE-XL	800	177	1.13	0.69	2.11	7.87	3.26
RAEv2 (\(K=7\), 本文)	80	35	1.06	0.64	1.15	5.21	2.17

消融实验

1. 导引机制消融（DiTDH-XL, 80 epochs, DINOv3-L）

导引方式	\(K=7\) gFID \(\downarrow\)	\(K=7\) IS \(\uparrow\)	\(K=23\) gFID \(\downarrow\)	\(K=23\) IS \(\uparrow\)	额外推理开销
无导引	1.65	228.0	3.01	206.0	无
标准 CFG	1.49	242.1	2.83	220.1	2× NFE
AutoGuidance (AG)	1.14	255.3	1.37	252.0	需额外从属模型
REPA Guidance (本文)	1.06	255.3	1.25	256.8	无 (1× NFE, 零额外模型)

2. 编码器选择与 RAE vs RAEv2 对比（DiTDH-XL @ 20 epochs, 无导引）

编码器	LP Top-1 (%) \(\uparrow\)	LDS \(\uparrow\)	复合分 \(\text{Avg} \uparrow\)	RAE gFID \(\downarrow\)	RAEv2 (\(K=1\)) gFID \(\downarrow\)
MoCov3-B	76.4	0.15	0.46	13.84	8.35
WebSSL-1B	84.1	0.18	0.51	8.60	4.16
DINOv2-B	83.9	0.41	0.62	3.75	2.81
DINOv3-B	84.5	0.38	0.61	4.25	2.76
DINOv3-L	87.0	0.42	0.65	3.30	2.61

关键发现

• 收敛速度实现数量级跃升：RAEv2 仅需 35 个 epoch 即可实现无导引 gFID \(\le 2\)，而前代 RAE 需要 177 个 epoch，传统扩散模型更需要 480 乃至 800+ epoch，整体训练耗时由一周以上缩短至仅约 10.5 小时。
• 重建与生成的帕累托最优权衡：多层特征相加中，增大 \(K\) 会单调优化图像重建（\(K=23\) 时 rFID 达到 0.18，PSNR 达到 27.03 dB）；而在导引生成中，\(K=7\) 表现最优（gFID 达到 1.06），用户可通过调节 \(K\) 在极少牺牲甚至提升生成性能的同时自由适配下游重建需求。
• 下游跨任务泛化力强：在文本生成图像任务（T2I）中，RAEv2 在微调后 GenEval 达到 82.7（优于 Flux-VAE 的 78.3 与原版 RAE 的 81.5）；在基于 RECON 的导航世界模型长程视频预测中，RAEv2 获得 105.61 的 FVD，彻底消除了帧间结构闪烁，显著优于 DIAMOND（762.73）与 NWM（200.97）。

亮点与洞察

• 多层求和（MLS）以极简实现高维子空间结构保持：无需复杂微调与跨模态特化数据，仅通过深浅层向量逐元素相加，就兼顾了底层空间细节与高阶全局语义。
• 澄清 RAE 与 REPA 的机制正交性：打破了“RAE 必然取代 REPA”的思维惯性，证明了潜空间语义与骨干网络空间注意力的解耦互补，使得全局和空间均优的编码器（如 DINOv3-L）真正成为生成的最优解。
• 将辅助对齐头转化为潜空间自导引源：发现 REPA 头天然构成一种局部的弱 \(x\)-prediction，借由流匹配重参数化将其无缝复用为内生导引，抹平了传统 CFG 与 AutoGuidance 的双倍前向或双重模型成本。

局限与展望

• 特征聚合方式仍较朴素：目前仅探索了线性等权求和与随机投影，未来针对各层级注意力权重自适应学习或可进一步提升生成与重建的帕累托前沿。
• 最佳表征仍依赖离散网格搜索：当前针对 27 种编码器的筛选是经验性的，尚未实现编码器与扩散生成的端到端联合梯度反传与表征演化。

相关工作与启发

• vs 原版 RAE (Zheng et al., 2025)：原版仅取末层特征且依赖 AutoGuidance 额外模型；本文通过多层求和（MLS）大幅强化重建，并借助 REPA 头的 \(x\)-prediction 重构实现单模型免 CFG 内置导引，80 轮指标即可大幅超越其 800 轮效果。
• vs REPA / iREPA (Yu et al., 2024; Singh et al., 2025)：传统 REPA 作用于标准 VAE 潜空间；本文证明其在 RAE 语义潜空间下依然不可替代，二者通过语义丰富度与空间自相似性达成互补正则。
• vs AutoGuidance / Internal Guidance (Karras et al., 2024; Zhou et al., 2025)：本文免去了 AutoGuidance 训练第二模型的繁琐开销，通过潜空间中特有的 \(x\)-prediction 物理意义直接借用 REPA 线性头做弱基线，实现了更纯粹高效的无损导引。

评分

• 新颖性: ⭐⭐⭐⭐⭐ 颠覆了 RAE 淘汰 REPA 的共识，提出首个将表征蒸馏头无损复用为自导引基线的极简优雅范式。
• 实验充分度: ⭐⭐⭐⭐⭐ 横跨 27 种编码器、ImageNet-256、T2I 以及具身导航世界模型，多指标全方位评测。
• 写作质量: ⭐⭐⭐⭐⭐ 逻辑层层推进，机理解析深入透彻，图表呈现严谨直观。
• 价值: ⭐⭐⭐⭐⭐ 将 RAE 训练成本压缩 10 倍以上，有效推动了多模态表征统一与生成大模型的基础设施落地。

相关论文

• [ICLR 2026] Diffusion Transformers with Representation Autoencoders
• [CVPR 2026] MeanFlow Transformers with Representation Autoencoders
• [NeurIPS 2026] Align & Invert: Solving Inverse Problems with Diffusion and Flow-based Models via Representation Alignment
• [NeurIPS 2026] D2D: Detector-to-Differentiable Critic for Improved Numeracy in Text-to-Image Generation
• [CVPR 2026] SRA 2: Variational Autoencoder Self-Representation Alignment for Efficient Diffusion Training