Improved Baselines with Representation Autoencoders

Paper: arXiv:2605.18324 Code: nanovisionx/RAEv2 Code reference: main @ 6ea3798 (2026-05-18)

1. Motivation (研究动机)

1.1 当前生成模型中潜在空间（Latent Space）设计的瓶颈

在现代图像和视频生成（如 Diffusion Models 和 Diffusion Transformers）中，传统的两阶段生成范式依赖于一个预训练好的自编码器（通常是变分自编码器 VAE，如 SD-VAE、SDXL-VAE 或 Flux-VAE）将高维像素空间压缩到低维潜在空间，然后再在该潜在空间上训练扩散模型。然而，这种经典的 VAE 潜在空间存在以下核心限制：

重构细节丢失：VAE 为了追求极致的压缩率，通常会抛弃很多低维几何与高频纹理细节，导致生成的图像在字符渲染（Text Rendering）、细粒度线条以及复杂纹理（如手部、人脸细节、交通标志）上出现扭曲和模糊。
理解与生成维度的割裂：VAE 的潜在空间完全是像素级重构驱动的，缺乏语义结构，这使得下游的多模态理解模型（如 VLM）无法直接重用生成模型的中间特征。为了实现统一理解和生成（Unified Multimodal Understanding and Generation），学术界和工业界迫切需要一种共享的 tokenization/representation。

1.2 预训练视觉表征自编码器（RAE）的兴起与痛点

为了打破理解与生成的壁垒，表征自编码器（Representation Autoencoder, RAE）应运而生。RAE 的核心思想是直接使用预训练好的强大视觉表征模型（如 DINOv2、SigLIP 等 Frozen Vision Encoders）作为自编码器的 Encoder，仅训练一个轻量级的 Decoder 进行特征到像素的重建。这样，扩散模型就可以直接在具有丰富语义的表征空间中进行训练。

然而，第一代 RAE 在实际应用中存在三个极其严重的痛点，限制了其大规模推广：

重构性能显著落后于专用 VAE：由于预训练视觉编码器（如 DINOv2）是为高层语义任务（语义理解、分类等）设计的，其最顶层特征（Final-Layer Features）丢失了大量的空间位置和低级纹理信息。因此，RAE 的重建结果在精细细节和文字上非常模糊。
与传统无分类器引导（CFG）不兼容：实验表明，当 RAE 潜在空间与传统的 Classifier-Free Guidance (CFG) 结合时，生成效果不仅没有提升，反而会引起严重的训练不稳定和样本分布漂移（如 Inception Score 崩溃）。因此，初代 RAE 不得不依赖自引导（AutoGuidance, AG），即额外训练一个参数量较小、性能较弱的辅助扩散模型作为 weaker baseline。这不仅使系统的训练复杂度加倍，还额外消耗了大量计算资源。
潜在空间特征表达单一：初代 RAE 默认只采用 Frozen Encoder 的最后一层输出（Final-Layer Features）作为潜在表征，这极大地限制了对预训练模型多层级丰富特征的挖掘。

1.3 本文研究的目标与价值

针对上述痛点，本文提出了第二代表征自编码器 RAEv2。其目标是：

在重构性能上，无需微调 Frozen Encoder，也无需引入额外的人脸/文字专用训练集，仅靠修改特征融合策略，就达到甚至超越专用 VAE 的重构表现；
在训练效率和收敛速度上，将 RAE 的收敛速度提升 10 倍以上，并在极短的训练周期内达到 SOTA 生成质量；
在系统架构上，消除对辅助 weaker model 的依赖，在单次前向传播中以“零推理成本”实现自引导（Self-Guidance）。

解决这一问题将极大地加速统一多模态大模型（UMM）在理解与生成任务上的协同设计，使生成任务能够无缝利用顶尖视觉编码器（如 DINOv3、SigLIP2）的底层与高层特征。

2. Idea (核心思想)

本文的三个核心 Insight 分别解决了 RAE 的重构、收敛和引导痛点，使 RAE 在理论和工程上更加优雅和高效：

2.1 泛化表征编码器：多层简单相加（MLS）激活底层几何细节

预训练视觉编码器不同的层捕获了互补的信息。早期到中期层保留了更丰富的局部空间结构和精细纹理，而后期层则更关注全局语义（Semantic Context）。

Core Insight：与其仅使用最后一层特征，不如将最后 $K$ 层的表征直接进行简单相加（Multi-Layer Sum, MLS）：

x = ℓ = L - K + 1 \sum L z_{ℓ}

在数学上，高维空间中的简单相加可以在保持通道维度不变的前提下，完美保留各子空间的几何距离结构，将低级纹理与高级语义融合。这种纯粹免训练（Training-free）的改动在不增加任何可学习参数的前提下，以极高的保真度提升了 RAE 的重建表现。

2.2 RAE 与 REPA 的互补机制：语义表征与空间正则的双剑合璧

学术界此前普遍认为，既然 RAE 已经将预训练编码器的表征作为扩散模型的 Latent 目标，那么就不再需要表征对齐（Representation Alignment, REPA）。因为 REPA 的作用是将中间层特征对齐到同一个表征空间，在 RAE 中这似乎构成了一个冗余的 Skip Connection。

Core Insight：大范围实证研究表明，RAE 和 REPA 的作用机制是完全互补的。RAE 负责提供高语义的潜在空间（通过线性探测 Linear Probing 驱动生成语义），而 REPA 负责正则化扩散模型中间层特征的 Token-Token 空间自相似结构（Spatial Self-Similarity）。因此，将相同的视觉表征同时作为输入潜在空间（RAE）和中间层对齐目标（REPA），能够大幅加快生成模型的收敛速度并提升保真度，同时使得更强大的双能编码器（如 DINOv3-L）能在生成任务中发挥最大效能。

2.3 REPA 即潜在空间的 x-prediction：用轻量级投影头实现“零成本”自引导

初代 RAE 因为无法直接使用传统 CFG，需要训练额外的辅助模型。

Core Insight：在 RAE 潜在空间中，干净的 Latent 就是编码器表征 $x$ 。位于扩散模型中间层的 REPA 投影头 $h_{ϕ}$ 的任务是根据早期层特征预测这个干净的 $x$ 。这意味着，REPA 头部本质上在执行 $x$ -prediction。由于该头部极轻量且仅能访问早期特征，其预测结果自然比主模型弱，完美扮演了 AutoGuidance 中辅助 weaker model 的角色。
创新重新参数化：通过将扩散模型的输出也重新参数化为 $x$ -prediction 空间，使得主模型预测 $\overset{x}{^}_{full}$ 与 REPA 头预测 $\overset{x}{^}_{repa}$ 处于同一空间。我们可以直接利用它们做差分引导（Internal Guidance）：

\overset{x}{^}_{guided} = \overset{x}{^}_{full} + w \cdot (\overset{x}{^}_{full} - \overset{x}{^}_{repa})

这完全消除了对辅助模型的训练需求。同时，由于 REPA 在扩散模型主干的同一次前向传播（Forward Pass）中计算，无需像 CFG 那样执行额外的前向计算，使推理阶段的函数计算次数（NFE）直接减半。

3. Method (方法)

3.1 总体框架与多阶段流程

RAEv2 的整体训练与推理流程分为以下几个核心部分：

Figure 2 解读：图 2 展示了 RAEv2 在各个维度上的卓越表现。(a) 展示了生成性能：RAEv2 仅需训练 80 个 epoch 即可在 ImageNet-256 上达到 2.17 的 FDr6 分数，超越了初代 RAE 训练 800 个 epoch 的表现；(b) 展示了极快的收敛速度，将达到 unguided gFID $\leq 2$ 所需的训练 Epochs（即 $EP_{FID@2}$ ）从 177 大幅缩短至 35；(c) 展示了更优越的图像重构质量；(d) 展示了推理效率的飞跃：通过复用 REPA 头部进行自引导，彻底摒弃了传统 CFG 的双前向传播或 AG 的第二模型，将 NFE 直接减半。

Figure 1 解读：图 1 展示了 RAEv2 在重构-生成性能（Reconstruction-Generation Trade-off）上的帕累托最优（Pareto-optimal）曲线。相较于各种专为重构设计的微调编码器方案，RAEv2 无需微调视觉编码器，仅需调整 MLS 的聚合层数 $K$ ，即可在只使用 ImageNet 训练的条件下，实现与专有 VAE（如 Flux-VAE、SDXL-VAE）媲美的重建质量，并保持极高的生成质量。

流程步骤：

输入特征提取与 MLS 融合：输入图像 $I \in R^{H \times W \times 3}$ ，通过 Frozen Vision Encoder 得到各层特征 $z_{ℓ}$ 。采用多层简单相加（MLS）聚合最后 $K$ 层特征：

x = ℓ = L - K + 1 \sum L z_{ℓ}

此时，特征保持为 $N \times d$ 的空间图（如 DINOv3-L 为 $16 \times 16 \times 1024$ ）。 2. 第一阶段：自编码器重构训练（Frozen Encoder）：

提取的 MLS 特征 $x$ 被输入至轻量级的通用 Transformer 架构的 Decoder 中。
Decoder 的任务是将表征 $x$ 映射回 RGB 像素空间，训练时使用 L1 损失、感知损失和判别器对抗损失。
训练完成后，Decoder 被冻结，作为 RAE 的生成重构模块。

第二阶段：基于 $x$ -prediction 的扩散模型训练：
- 扩散模型使用 Flow Matching 框架在潜在空间 $x$ 上进行训练。
- 引入 REPA 表征对齐：在扩散 Transformer 的第 8 层，使用一个线性投影层 $h_{ϕ}$ 预测干净的潜在表征 $x$ （作为 REPA Loss 的监督目标）。
- 主模型输出通过重新参数化，同样预测干净的潜在表征 $\overset{x}{^}_{full}$ 。
- 扩散损失对主模型与 REPA 预测的弱模型同时进行监督。

3.2 关键组件与公式细节

3.2.1 泛化 RAE 的特征融合方案 (MLS vs MLR)

除了 MLS（Multi-Layer Summation）外，作者还探索了 MLR（Multi-Layer Random Projection）：将最后 $K$ 层特征在通道维度拼接，然后用一个在初始化时从高斯分布采样且保持冻结的随机矩阵 $R \in R^{K d \times d}$ 投影回 $d$ 维：

x = [z_{L - K + 1} ∥ \dots ∥ z_{L}] R

结论：MLS 和 MLR 在第一阶段重建上表现一致，但在第二阶段生成任务中，MLS 能够显著取得更好的 gFID。数学直觉上，简单相加自然地融合了表征流的各个语义层级，并没有像 MLR 那样引入随机特征混淆。

3.2.2 RAE 与 REPA 的互补性指标 (LP 与 LDS)

Linear Probing (LP)：通过在冻结特征上训练线性分类器，衡量全局语义的清晰度。
Local Distance Similarity (LDS)：用于衡量特征的空间自相似拓扑。作者通过对 27 种不同编码器的实证分析发现：
RAE alone 的性能高度取决于特征的全局语义品质（LP 与 gFID 的相关系数 $r = - 0.81$ ），而与 LDS 无关；
REPA alone (with VAE) 则极度依赖于空间拓扑保留能力（LDS 与 gFID 的相关系数 $r = - 0.89$ ）；
RAE + REPA 完美结合了两者的长处，当两者叠加时，综合评估分数 $Avg (LP^{'}, LDS^{'})$ 与最终生成质量呈现出高达 $- 0.83$ 的极强相关性。这解释了为什么将 DINOv3-L（在 LP 和 LDS 上都是顶尖表现）同时作为 RAE 潜空间和 REPA 目标时，RAEv2 能达到最佳的生成质量。

3.2.3 基于 $x$ -prediction 的 REPA 引导数学推导

传统的 Flow Matching 预测的是速度 $v_{t}$ ：

v_{t} = x_{0} - x_{1}

为了能够自然地进行自引导，我们需要主模型预测干净潜在空间的目标 $x_{1}$ （即 $x_{full}$ ）。在训练时，对于时间步 $t \sim Logit-Normal (0, 1)$ ：

混淆样本定义为： $x_{t} = (1 - t) \cdot x_{1} + t \cdot x_{0}$
扩散模型输入 $x_{t}$ ，直接预测 $x_{1}$ ，记为 $\overset{x}{^}_{full} = f_{θ} (x_{t}, t)$ ；
相应的速度预测通过差分关系隐式计算：

v_{pred} = \frac{x _{t} - x ^ _{full}}{t}

扩散损失定义在速度空间：

L_{diff} = ∥ v_{pred} - v_{t} ∥^{2}

同时，在中间层（深度 8），投影头 $h_{ϕ} (h_{8})$ 也预测 $x_{1}$ ，记为 $\overset{x}{^}_{repa}$ 。

REPA 对齐损失（对预测头施加梯度监督，使其逼近干净的真实表征）：

L_{repa} = λ \cdot ∥ \overset{x}{^}_{repa} - x_{1} ∥^{2}

基础弱模型损失（确保 REPA 头在预测整个潜在目标时是有意义的弱估计）：

L_{base} = γ \cdot ∥ \frac{x _{t} - x ^ _{repa}}{t} - v_{t} ∥^{2}

其中 $λ = 0.5$ ， $γ = 1.0$ 。

在**推理采样（Inference Sampling）**时，不需要 unconditional 条件前向传播。只进行一次前向传播，模型同时吐出 $\overset{x}{^}_{full}$ 和 $\overset{x}{^}_{repa}$ 。我们按照如下公式引导：

\overset{x}{^}_{guided} = \overset{x}{^}_{full} + w \cdot (\overset{x}{^}_{full} - \overset{x}{^}_{repa})

然后将引导后的特征转换为采样漂移向量：

v_{guided} = \frac{x _{t} - x ^ _{guided}}{t}

送入 ODE 积分器（如 Euler 求解器）进行下一步迭代。

3.3 论文公式与 Released Code 实现差异

IMPORTANT
论文描述与代码实现的细微 gap： 在论文第 2.1 节中，多层简单相加（MLS）定义为直接对最后 $K$ 层特征求和（Eq. 1）。然而，在官方开源代码 src/encoders/vision_encoder.py 的 DINOv3MultiLayerSimpleAddEncoder 中，具体的实现并不是对各层特征直接求和，而是**先求这 $K$ 层特征的平均值（mean），然后加上最后一层特征在空间维度上的广播平均值（final_mean）**作为全局偏置：
patch_tokens = torch.stack(outputs, dim=0).mean(dim=0)
final_mean = outputs[-1].mean(dim=1, keepdim=True)
patch_tokens = patch_tokens + final_mean
这一实现细节与公式 (1) 的直接求和（Sum）存在常数倍率与全局偏置的差异。这可能是为了保持特征尺度稳定在 $[- 1, 1]$ 附近，避免求和导致激活值过大而影响数值稳定性。在 EUPEMultiLayerSimpleAddEncoder 中则使用的是 .sum(dim=0)。写笔记时应明确指出这一工程实现上的差异。

3.4 代码-论文映射

Code reference: main @ 6ea3798 (2026-05-18) — 下表和伪代码基于此 commit 构建

论文概念 / 模块	对应源代码文件	关键类 / 函数	作用与实现细节
MLS 编码器特征融合	vision_encoder.py	`DINOv3MultiLayerSimpleAddEncoder`	提取 DINOv3 指定多层的 patch tokens，计算其均值并加上最后一层的全局均值。
RAE 阶段 1 模块	rae.py	`RAE`	封装了 Frozen Encoder 和 Decoder，包含 `encode`（含 MLS）、`decode`、和用于重构训练的 `noising`。
带有自引导的 DiT 架构	DDT.py	`DiTwDDTHeadIG`	实现了在第 8 层获取 `x_base` 并通过 `base_final_layer` 输出弱预测的 DiT 主干。
Flow Matching 传输损失	transport.py	`Transport.training_losses`	核心训练损失计算，支持 $x$ -prediction 重新参数化以及 REPA 损失与 Base 弱模型损失的加和。
自引导采样逻辑	guidance_utils.py	`forward_with_internalguidance`	推理采样时的自引导逻辑，根据主模型和 REPA 头输出做外推插值。
训练整体流程主控	train.py	`main`	初始化 RAE 潜在空间、可训练的 DDT 模型、优化器以及 Flow Matching 采样配置，控制多卡 DDP 训练。

3.5 核心模块 PyTorch 风格伪代码

3.5.1 多层简单相加（MLS）特征融合实现 (对应 `DINOv3MultiLayerSimpleAddEncoder`)

import torch
import torch.nn as nn
from typing import List, Dict
 
class MLSFeatureExtractor(nn.Module):
    """
    MLS 特征融合伪代码。
    对应代码中 DINOv3MultiLayerSimpleAddEncoder 的实现逻辑。
    """
    def __init__(self, vit_model: nn.Module, layer_indices: List[int]):
        super().__init__()
        self.model = vit_model
        self.layer_indices = layer_indices  # 例如 [11, 13, 15, 17, 19, 21, 23]
        
    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        # 从 Timm/DINOv3 模型中提取指定中间层的特征
        # outputs 的形状为: [num_layers, B, N, C]
        outputs = self.model.get_intermediate_layers(
            x, n=self.layer_indices, reshape=False, 
            return_class_token=False, norm=True
        )
        
        # 沿层维度堆叠
        stacked = torch.stack(outputs, dim=0) # [K, B, N, C]
        
        # 计算选定多层特征的平均值
        patch_tokens = stacked.mean(dim=0) # [B, N, C]
        
        # 计算最后一层特征在 token 维度上的均值，作为全局偏置
        final_mean = outputs[-1].mean(dim=1, keepdim=True) # [B, 1, C]
        
        # 融合全局偏置
        patch_tokens = patch_tokens + final_mean # [B, N, C]
        
        return {
            "x_norm_clstoken": final_mean.squeeze(1),
            "x_norm_patchtokens": patch_tokens
        }

3.5.2 阶段 2 带有内部弱预测分支的模型前向网络 (对应 `DiTwDDTHeadIG`)

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class DiTwDDTHeadIG(nn.Module):
    """
    带有内建 REPA 引导分支的 Diffusion Transformer 伪代码。
    """
    def __init__(self, in_channels=1024, enc_hidden_size=1440, dec_hidden_size=2048, 
                 depth_enc=28, depth_dec=2, base_model_depth=8, repa_layer_depth=8, z_dim=768):
        super().__init__()
        self.in_channels = in_channels
        self.enc_hidden_size = enc_hidden_size
        self.base_model_depth = base_model_depth     # 获取弱模型预测的中间层数 (8)
        self.repa_layer_depth = repa_layer_depth     # REPA 监督的目标层数 (8)
        
        # 各种 Embedder 略...
        self.s_embedder = nn.Identity()  # 仅为演示占位
        
        # Transformer 块列表 (28 个 Encoder + 2 个 Decoder)
        self.enc_blocks = nn.ModuleList([nn.Identity() for _ in range(depth_enc)])
        self.dec_blocks = nn.ModuleList([nn.Identity() for _ in range(depth_dec)])
        
        # 最终预测头
        self.final_layer = nn.Linear(dec_hidden_size, in_channels) # 预测 x_full
        self.base_final_layer = nn.Linear(enc_hidden_size, in_channels) # 预测 x_base
        self.repa_projector = nn.Linear(enc_hidden_size, z_dim) # REPA 对齐投影
 
    def forward(self, x, t, return_intermediate=False, **condition_kwargs):
        # 1. 构造输入 tokens (演示用，略去细节拼接)
        seq = x 
        
        zt_intermediate = None
        x_base = None
        
        # 2. 依次前向通过 28 层 Encoder blocks
        for i, block in enumerate(self.enc_blocks):
            seq = block(seq)
            
            # 在 REPA 监督深度（第 8 层）提取特征进行对齐投影
            if return_intermediate and (i + 1) == self.repa_layer_depth:
                zt_intermediate = self.repa_projector(seq)
                
            # 在指定弱分支深度（第 8 层）截断提取特征作为弱预测的输入
            if (i + 1) == self.base_model_depth:
                x_base = seq
                
        # 3. 主预测分支继续通过 Decoder blocks
        # x_full 的后续前向映射与 DDT 投影逻辑
        seq_dec = seq # ...
        for block in self.dec_blocks:
            seq_dec = block(seq_dec)
        x_full_out = self.final_layer(seq_dec)
        
        # 4. 弱预测分支：利用第 8 层提取特征，通过 base_final_layer 直接输出潜在空间预测
        x_base_out = self.base_final_layer(x_base)
        
        if return_intermediate:
            return (x_full_out, x_base_out), zt_intermediate
        return x_full_out, x_base_out

3.5.3 带有 x-prediction 和 REPA 的 Flow Matching 训练损失 (对应 `Transport.training_losses`)

import torch
import torch.nn.functional as F
 
class FlowMatchingTransportLoss:
    """
    Flow Matching 结合 x-prediction 重参数化和多任务损失的伪代码。
    """
    def __init__(self, prediction_type="x", t_eps=0.05):
        self.prediction_type = prediction_type
        self.t_eps = t_eps
 
    def sample_time(self, batch_size, device):
        # 从 Logit-Normal(0, 1) 分布中采样时间步
        t = torch.randn(batch_size, device=device).sigmoid()
        return t
 
    def training_losses(self, model, x1, model_kwargs, z_clean=None, repa_coeff=0.5, base_model_coeff=1.0):
        """
        Args:
            x1: 真实的干净潜在空间数据 (RAE 编码输出)
            z_clean: 用于 REPA 监督的真实视觉编码器第 8 层特征
        """
        device = x1.device
        batch_size = x1.size(0)
        
        # 1. 采样时间步并添加流混淆噪声
        t = self.sample_time(batch_size, device)
        x0 = torch.randn_like(x1) # 高斯先验
        
        # 线性插值 xt
        t_expand = t.view(batch_size, *([1] * (x1.ndim - 1)))
        xt = (1 - t_expand) * x1 + t_expand * x0
        
        # 目标速度 vt
        vt = (xt - x1) / t_expand.clamp_min(self.t_eps)
        
        # 2. 模型前向计算 (输出 full 与 base 预测，以及中间层 REPA 特征)
        enable_repa = (z_clean is not None) and (repa_coeff is not None)
        if enable_repa:
            outputs, zt_pred = model(xt, t, return_intermediate=True, **model_kwargs)
        else:
            outputs = model(xt, t, **model_kwargs)
            zt_pred = None
            
        model_output, base_output = outputs # 解包双分支预测
        
        # 3. 计算主模型的速度空间损失 (x-pred 转 v-pred)
        # v_pred = (xt - x_pred) / t
        v_pred_full = (xt - model_output) / t_expand.clamp_min(self.t_eps)
        loss_full = F.mse_loss(v_pred_full, vt)
        
        # 4. 计算弱预测模型分支损失
        v_pred_base = (xt - base_output) / t_expand.clamp_min(self.t_eps)
        loss_base = F.mse_loss(v_pred_base, vt)
        
        # 5. 总损失融合
        total_loss = loss_full + base_model_coeff * loss_base
        
        terms = {
            "loss_full": loss_full,
            "loss_base": loss_base
        }
        
        # 6. 融合 REPA 对齐损失
        if enable_repa and zt_pred is not None:
            loss_repa = F.mse_loss(zt_pred, z_clean)
            total_loss = total_loss + repa_coeff * loss_repa
            terms["loss_repa"] = loss_repa
            
        terms["loss"] = total_loss
        return terms

3.5.4 推理时的自引导采样逻辑 (对应 `forward_with_internalguidance`)

import torch
 
def forward_with_self_guidance(model, x, t, ig_scale=1.78, ig_interval=(0.10, 1.0), in_channels=1024, **condition_kwargs):
    """
    推理阶段的自引导逻辑伪代码。
     ig_output = base_out + ig_scale * (full_out - base_out)
    """
    # 1. 前向传播获取主模型预测与 REPA 弱模型预测
    # 注意此时不需要为 unconditional branch 进行二次前向
    full_out, base_out = model(x, t, **condition_kwargs)
    
    # 2. 截取有效的通道范围
    eps_full = full_out[:, :in_channels]
    eps_base = base_out[:, :in_channels]
    
    # 3. 判断是否在引导时间窗口内
    t_min, t_max = ig_interval
    t_expanded = t.view(-1, *[1] * (eps_full.ndim - 1))
    in_window = (t_expanded >= t_min) & (t_expanded <= t_max)
    
    # 4. 在有效窗口内进行外推差引导，否则保持主预测输出
    # 数学上这与 x_full + w * (x_full - x_base) 完全等价 (其中 ig_scale = 1 + w)
    guided_out = torch.where(
        in_window,
        eps_base + ig_scale * (eps_full - eps_base),
        eps_full
    )
    
    return guided_out

4. Experimental Setup (实验设置)

4.1 数据集与规模

本论文在三类截然不同的任务上对 RAEv2 进行了广泛的验证：

类条件图像生成 (Class-Conditional Image Generation)：
- ImageNet-1K：分辨率 $256 \times 256$ ，包含约 128 万张训练图片，1000 个类别。
文本生成图像 (Text-to-Image Generation)：
- JourneyDB + BLIP3o (包括长/短文本对)：用于预训练（Pre-training），总共训练 15 万个 Iteration（Batch Size 1024）。
- BLIP3o-60k 子集：用于文本对齐微调（Fine-tuning），训练 50 个 Epoch。
导航世界模型 (Navigation World Models)：
- RECON 户外导航视频数据集：在 4 FPS 下提取视频序列。模型输入 $N = 4$ 帧历史图像（ $256 \times 256$ ），回归自回归预测接下来的轨迹帧。

4.2 对比基线 (Baselines)

为了证明收敛效率和帕累托性能的优越性，论文对比了多类 SOTA 扩散模型和自编码器变体：

经典 Latent Diffusion 方案：SiT-XL/2 (使用传统 SD-VAE 潜在空间)、DDT-XL、LightningDiT。
表征对齐方案：SiT-XL/2-REPA (基于 VAE，但利用中间层对齐)、REG、REPA-E (端到端优化 VAE)。
第一代表征自编码器：RAE-XL (初代 RAE 方案，仅基于单层最后一级特征特征，并使用 AutoGuidance)。

4.3 评估指标 (Evaluation Metrics)

gFID (Generation FID)：无条件/有引导生成图像与 ImageNet 验证集之间的 Inception FID 分数，评估宏观生成保真度。
rFID (Reconstruction FID)：重建图像与原图的 FID，评估重建的真度。
IS (Inception Score)：评估生成类别的多样性与清晰度。
FDr (Representation Fréchet Distance) / FD6r：在 6 种不同的强表征空间（Inception, ConvNeXt, DINOv2, MAE, SigLIP, CLIP）中提取特征并计算 FID，比传统 FID 更能感知细节（如文字、细微线条）退化。FD6r 是这 6 个分数的几何平均值。
MDr (Representation Monge Distance)：基于 MIND 框架，计算在 6 个表征空间中的 Monge 距离最优传输开销，反映点对点的精密分布重合度。
EPFID@k (Epochs to reach unguided gFID ⇐ k)：衡量训练收敛效率的全新指标。记录不使用任何引导时，生成质量达到 gFID $\leq k$ 所需的最少训练 Epochs，论文中取 $k = 2$ 。
FVD (Fréchet Video Distance) / LPIPS：针对世界模型任务，评估连续生成帧的时序一致性与感知相似度。

4.4 训练超参数与硬件配置

所有实验（ImageNet, T2I, World Models）均基于统一的 DDT-XL 作为主干网络（Backbone）。

硬件开销

采用 4 x 8 块 NVIDIA H100 GPU 节点。
RAEv2 仅需训练 ~12 小时 即可达到 1.06 的 SOTA gFID。相比之下，初代 RAE 需要在该集群下训练超过 1 周以上（800 轮）。

核心超参数配置列表：

主干网络 (DDT-XL)：28 个 Encoder 块（隐藏维度 1152，20 头），2 个 DDT Decoder 块（隐藏维度 2048，16 头）。参数量约 8.39 亿 (0.9B)。
RAE 阶段 1 (自编码器)：
- 编码器：DINOv3-L ( $256 \times 256$ 输入，Patch Size 16，隐特征 1024 维)。
- 融合层数 $K$ ：ImageNet 图像生成默认选用 $K = 7$ ；图像重构任务选用 $K = 23$ 。
- 重建 Decoder：2 块 Transformer 架构，隐藏维度 2048，单独在 ImageNet 上预训练 16 轮并冻结。
- 潜在空间形状： $16 \times 16 \times 1024$ 。
优化器与调度器：
- 优化器：Gmuon，学习率 $2 \times 1 0^{- 4}$ ，权重衰减 0.0。
- 学习率调度：Linear Decay（前 25 轮 Warmup，在第 50 轮降至终点学习率 $2 \times 1 0^{- 5}$ ）。
- Batch Size：全局 Batch Size 1024，混合精度采用 bfloat16。
- 梯度裁剪（Max Norm）：1.0。
- EMA（指数滑动平均）衰减率：0.9995。
推理与引导设置：
- 采样求解器：ODE Euler Discretization，采样步数 50。
- 内部自引导配置 (REPA IG)：引导强度 $w = 0.78$ （即伪代码中 ig_scale = 1.78），激活时间段 $[t_{min}, t_{max}] = [0.10, 1.0]$ 。

5. Experimental Results (实验结果)

5.1 图像生成与收敛效率指标（ImageNet-256）

在 ImageNet 图像生成任务上，RAEv2 展现出了惊人的收敛优势和生成保真度。下表总结了 RAEv2 与传统扩散模型方案的对比：

方法 (Method)	训练轮数 (Epochs)	收敛效率 $EP_{FID@2}$ ↓	gFID (带引导) ↓	FD6r ↓
SiT-XL/2 (VAE)	800	$> 800$	8.44	17.69
DDT-XL (VAE)	800	-	5.70	13.86
SiT-XL/2-REPA	800	$> 800$	5.45	12.65
LightningDiT	800	$> 800$	4.57	10.21
REPA-E (End-to-End)	800	480	3.04	6.28
RAE-XL (初代 RAE)	800	177	3.26	7.87
RAEv2 (K=7, Ours)	80	35	1.06	2.17

结果分析：

生成保真度的绝对跃升：RAEv2 仅需训练 80 个 Epoch 即可在 ImageNet 上实现 1.06 的 gFID 以及 2.17 的 FD6r。这不仅远超训练了 800 个 Epoch 的传统 VAE 方案（如 DDT-XL 的 5.70），更将初代 RAE 的 SOTA 记录（3.26 gFID）缩短至其十分之一的训练时间。
收敛速度的十倍提升：引入 $EP_{FID@2}$ 指标后，传统的对齐方案和 VAE 方案需要 480 轮以上才能使 unguided gFID 小于 2。而 RAEv2 仅需 35 个 Epoch 即可实现这一目标，实现了名副其实的“生成模型速度跑（Speedrun）”。

5.2 核心消融实验 (Ablations)

消融 1：表征编码器的选择（不同 Frozen Encoder）

作者在统一使用 DDT-XL、训练 20 Epoch 且无引导的相同设置下，消融了 27 种预训练编码器。表 2 给出了代表性视觉编码器的消融数据：

视觉编码器 (Vision Encoder)	全局语义 (LP) ↑	空间拓扑 (LDS) ↑	综合指标 $Avg (LP^{'}, LDS^{'})$ ↑	RAE gFID ↓	RAEv2 gFID (Ours) ↓
MoCov3-B	76.4%	0.15	0.46	13.84	8.35
WebSSL-1B	84.1%	0.18	0.51	8.60	4.16
DINOv3-B	84.5%	0.38	0.61	4.25	2.76
DINOv2-B	83.9%	0.41	0.62	3.75	2.81
DINOv3-L	87.0%	0.42	0.65	3.30	2.61

消融结论：在初代 RAE 架构下，DINOv2-B 比 DINOv3-L 的表现更好（3.75 vs 3.30，DINOv3-L在 RAE 配方下甚至发生退化）。但引入 RAEv2 后，具有更强全局语义（LP 87.0%）和最强空间拓扑（LDS 0.42）的 DINOv3-L 展现了无可比拟的优势，取得了 2.61 的最佳无引导 gFID 分数。

消融 2：特征聚合层数 $K$ 的帕累托消融

在 DINOv3-L（24层）上，改变 MLS 融合的层数 $K$ （由 1 层即最后一层，到 23 层即全层平均）：

融合层数 $K$	阶段 1 重建 PSNR ↑	阶段 1 重建 rFID ↓	阶段 2 无引导 gFID ↓ (80轮)	阶段 2 自引导 gFID ↓ (80轮)
1 (Last layer)	18.93	0.60	1.50	2.23
7 (Ours Default)	22.57	0.29	1.65	1.06
23 (Full-Sum)	27.03	0.18	3.01	1.25

消融结论：随着聚合层数 $K$ 的增加，第一阶段重构表现（PSNR/rFID）呈现单调性的显著提升（ $K = 23$ 时 PSNR 提升至 27.03，重建 FID 降至 0.18）。但对于扩散模型生成质量而言，无引导下 $K = 1$ 表现最好，而在带自引导的采样下， $K = 7$ 取得了帕累托最优的图像重构与图像生成质量（gFID 1.06）。

消融 3：三种自引导机制对比消融

在不同 MLS 聚合深度下，对比不同的采样引导策略：

引导方案 (Guidance)	推理额外 forward 成本	额外辅助模型	K = 7 gFID ↓	K = 23 gFID ↓
无引导 (w/o Guidance)	无	否	1.65	3.01
CFG (传统分类器引导)	双倍（NFE x 2）	否	1.49	2.83
AutoGuidance (AG)	双倍（NFE x 2）	是 (训练 weaker 扩散器)	1.14	1.37
REPA IG (Ours)	无 (零成本自引导)	否	1.06	1.25

消融结论：传统的 CFG 无法显著改善 RAE 的 gFID；AG 虽有效果但需要多训练一个模型，且使采样推理计算加倍。相比之下，本论文提出的 REPA 自引导在不增加任何额外推理计算（零成本）的条件下，取得了最佳的生成保真度（gFID 1.06 / 1.25）。

5.3 跨任务泛化验证

1. 文本生成图像 (T2I) 表现

使用 SiGLIP2-B 作为 RAE 编码器，结合 Qwen3-0.6B 作为文本编码器。对比在 JournyDB 和 BLIP3o 上训练后的文本 Adherence 分数：

潜在空间类型	预训练 GenEval ↑	预训练 DPG-Bench ↑	微调后 GenEval ↑	微调后 DPG-Bench ↑
Flux-VAE	41.7	77.6	78.3	79.2
RAE	58.4	80.1	81.5	80.6
RAEv2 (Ours)	62.4	81.7	82.7	82.3

在预训练和微调两个维度上，基于表征自编码器的方案（RAE 与 RAEv2）均显著超过了像素级重建的 Flux-VAE（预训练 GenEval 从 41.7 飙升至 62.4），这表明在语义空间中学习扩散模型对文本 Adherence、多物体组合生成有降维打击式的优势。其中，RAEv2 比初代 RAE 在文本Adherence（GenEval 82.7）上表现更佳。

在 RECON 户外轨迹数据集上自回归推演 16 秒：

算法方案	主干模型参数量	视频质量指标 FVD ↓
DIAMOND	10 亿 (1B)	762.73
NWM	10 亿 (1B)	200.97
RAE	6.22 亿 (622M)	312.01
RAEv2-NWM (Ours)	6.22 亿 (622M)	105.61

世界模型高度依赖低级几何一致性以避免自回归过程中的漂移。通过 MLS 特征融合释放了 DINOv3-L 前几层对空间几何的表征能力，RAEv2-NWM 在参数量几乎减半的条件下取得了 105.61 的极低 FVD，大幅改善了时序连续性，并彻底解决了视频自回归中常见的闪烁（Flickering）问题。

5.4 方法的局限性 (Limitations)

作者在附录中总结了 RAEv2 目前的两个主要局限性：

多层表征融合同步的配方过于简单：目前 MLS 仅探索了简单无参数求和（MLS）或随机映射（MLR），虽然 MLS 效果优异，但理论上存在更优雅的多层非线性、注意力门控式融合策略以进一步榨取底层几何特征；
缺乏表征层面的端到端联合训练：表征的筛选仍局限于在离散的预训练模型库中进行经验搜索（如测试 DINOv2 vs DINOv3 表现）。若未来能配合下游扩散任务对表征提取层进行端到端微调，重构和生成质量可能会有更大提升。

Paper Notes

探索

Improved Baselines with Representation Autoencoders

Improved Baselines with Representation Autoencoders

1. Motivation (研究动机)

1.1 当前生成模型中潜在空间（Latent Space）设计的瓶颈

1.2 预训练视觉表征自编码器（RAE）的兴起与痛点

1.3 本文研究的目标与价值

2. Idea (核心思想)

2.1 泛化表征编码器：多层简单相加（MLS）激活底层几何细节

2.2 RAE 与 REPA 的互补机制：语义表征与空间正则的双剑合璧

2.3 REPA 即潜在空间的 x-prediction：用轻量级投影头实现“零成本”自引导

3. Method (方法)

3.1 总体框架与多阶段流程

流程步骤：

3.2 关键组件与公式细节

3.2.1 泛化 RAE 的特征融合方案 (MLS vs MLR)

3.2.2 RAE 与 REPA 的互补性指标 (LP 与 LDS)

3.2.3 基于 x-prediction 的 REPA 引导数学推导

3.3 论文公式与 Released Code 实现差异

3.4 代码-论文映射

3.5 核心模块 PyTorch 风格伪代码

3.5.1 多层简单相加（MLS）特征融合实现 (对应 DINOv3MultiLayerSimpleAddEncoder)

3.5.2 阶段 2 带有内部弱预测分支的模型前向网络 (对应 DiTwDDTHeadIG)

3.5.3 带有 x-prediction 和 REPA 的 Flow Matching 训练损失 (对应 Transport.training_losses)

3.5.4 推理时的自引导采样逻辑 (对应 forward_with_internalguidance)

4. Experimental Setup (实验设置)

4.1 数据集与规模

4.2 对比基线 (Baselines)

4.3 评估指标 (Evaluation Metrics)

4.4 训练超参数与硬件配置

硬件开销

核心超参数配置列表：

5. Experimental Results (实验结果)

5.1 图像生成与收敛效率指标（ImageNet-256）

结果分析：

5.2 核心消融实验 (Ablations)

消融 1：表征编码器的选择（不同 Frozen Encoder）

消融 2：特征聚合层数 K 的帕累托消融

消融 3：三种自引导机制对比消融

5.3 跨任务泛化验证

1. 文本生成图像 (T2I) 表现

2. 导航世界模型 (Navigation World Model) 表现

5.4 方法的局限性 (Limitations)

目录

3.2.3 基于 $x$ -prediction 的 REPA 引导数学推导

3.5.1 多层简单相加（MLS）特征融合实现 (对应 `DINOv3MultiLayerSimpleAddEncoder`)

3.5.2 阶段 2 带有内部弱预测分支的模型前向网络 (对应 `DiTwDDTHeadIG`)

3.5.3 带有 x-prediction 和 REPA 的 Flow Matching 训练损失 (对应 `Transport.training_losses`)

3.5.4 推理时的自引导采样逻辑 (对应 `forward_with_internalguidance`)

消融 2：特征聚合层数 $K$ 的帕累托消融