RiT：原生扩散Transformer在表征空间中已足够

证明原生DiT加x-prediction在DINOv2表征空间即可达图像生成SOTA

前置知识

流匹配 / Flow Matching

一种生成建模框架，学习一个速度场把高斯噪声沿直线路径 $z_t = t z_0 + (1-t)\epsilon$ 传输到数据，$t\in[0,1]$，路径时间导数即速度目标 $v = z_0 - \epsilon$。相比传统扩散，路径为直线、ODE更易高效求解。

本文整体方法建立在 flow matching 之上，理解 x-prediction 与 v-prediction 的区别是核心。

Diffusion Transformer (DiT)

用 Transformer 替换 U-Net 作为扩散去噪骨干，把潜变量 token 化后用自注意力建模。本文的 vanilla DiT 含 SwiGLU、RMSNorm、2D-RoPE、QK 归一化等标准组件，参数 676M。

本文核心主张正是 vanilla DiT 在表征空间已足够，需理解其基本结构。

x-prediction 与 v-prediction

v-prediction 让网络回归速度 $v=(z_0-z_t)/(1-t)$，在 $t\to1$ 时发散且依赖离流形的 $z_t$；x-prediction 直接回归干净数据 $\hat{z}_0$，目标落在数据流形上、不显式依赖 $t$。

本文最关键的创新就是改用 x-prediction 消解几何干扰，必须理解二者差异。

DINOv2

Meta 的自监督视觉 Transformer，通过学生-教师蒸馏学习强语义表征，其特征对每个 token 施加 LayerNorm，把范数钉在 $\|z\|\approx\sqrt{d}$，形成范数集中、近各向同性的分布。

RiT 直接在冻结 DINOv2 特征空间建模，其几何性质是全文分析的对象。

本征维数 (Intrinsic Dimension)

数据流形的真实自由度，即剥离环境冗余后描述数据所需的独立方向数。本文用 TwoNN 估计器从最近邻距离比的最大似然恢复 $\hat{d}$，测得像素与 DINOv2 均约 33。

它是四条几何分析轴之一，用于证明像素与 DINOv2 底层复杂度相同。

分类器自由引导 (CFG)

推理时把条件与无条件预测按 $\hat{\epsilon}=\epsilon_{uncond}+w(\epsilon_{cond}-\epsilon_{uncond})$ 外推，提升样本质量。本文引导强度 CFG=3.7，把 FID 从 1.45 降到 1.14。

主结果区分有/无引导两种设定，理解 CFG 才能读懂 FID 对比。

研究动机

在基于表征空间的扩散生成领域，研究者面临被称为「几何干扰（geometric interference）」的核心难题。DINOv2 这类自监督编码器对每个 token 施加 LayerNorm，会把特征范数钉死在 $\|z\|\approx\sqrt{d}$ 附近，导致线性 flow matching 路径 $z_t = t z_0 + (1-t)\epsilon$ 穿过编码器从未输出过的「环境区域」；在这些离流形的中间态 $z_t$ 处，v-prediction 的速度目标会获得一个大的径向分量。为应对这种径向模糊，先前工作付出了相当代价：RAE 在 v-prediction 之上叠加专门的宽预测头 DDT，还需一个 ViT 解码器把 DINOv2 特征映射回像素；并发工作 RJF 把欧氏传输替换为范数集中球面上的黎曼流匹配（SLERP 路径）。这些改动要么增加架构复杂度，要么改动传输路径，且都未从统计几何角度系统解释为何 DINOv2 空间对 flow matching 困难还是有利。

本文的目标是本文要回答一个基本问题：预训练表征空间相比像素空间是否提供了对 flow matching 学习更有利的分布？若表征空间的统计几何本身足够友好，就能用最朴素的原生 DiT 配合 x-prediction 达到竞争力，从而避免 DDT 预测头、黎曼传输、表示对齐损失等一切额外组件。具体目标包括：在 ImageNet $256\times256$ 类条件生成上，用 19% 更少参数（676M vs 839M）和最小 DINOv2 变体（DINOv2-S，$d=384$）取得优于 DiTDH-XL 的 FID；并证明不依赖任何蒸馏或一致性训练即可实现少步生成（5-10 步 Heun）。

与已有工作不同的是，本文的独特切入是「目标侧重构」而非「架构侧」或「传输侧」修补。作者先沿四条互补几何轴——本征维数、有效秩、边缘高斯性、流形上线性插值——系统刻画像素、SD-VAE、DINOv2 三种空间差异，发现像素与 DINOv2 本征维数几乎相同（$\hat{d}\approx33$），但 DINOv2 在另三轴上全面占优；SD-VAE 居中，说明优势来自表征学习目标而非单纯压缩。基于此，他们把问题归因到 v-prediction 的预测目标，改用 x-prediction 直接回归流形上的 $z_0$，在输出端而非输入端消解径向模糊。这种「先理解几何、再最小化架构」的思路是本文区别于所有并发工作的本质。

核心方法

整体直觉是：只要表征分布相对高斯源 $\mathcal{N}(0,I)$ 足够友好（传输路径短而良态），那么把预测目标放在数据流形上的 x-prediction 就能让最朴素的 DiT 胜任生成。技术路线：复用 RAE 冻结的 DINOv2-with-Registers 编码器和 ViT 解码器在像素与特征间转换；编码器输出 patch token $z\in\mathbb{R}^{d\times h\times w}$ 与 [CLS] token $z_{cls}\in\mathbb{R}^d$，两者投影后拼接、联合参与自注意力，分别用独立线性头预测 $\hat{z}_0$ 和 $\hat{z}_{cls,0}$。主干是标准 DiT（SwiGLU、RMSNorm、2D-RoPE、QK 归一化、in-context class token）。整套方法没有任何架构创新，创新全在训练配方与对几何的洞察上。

核心创新点是「目标侧重构消解几何干扰」。v-prediction 下网络须拟合 $v=(z_0-z_t)/(1-t)$，该目标依赖离流形的 $z_t$、在 $t\to1$ 时按 $(1-t)^{-1}$ 发散、横跨整个环境空间，迫使网络把容量花在径向（范数）方向而非切向（沿流形）方向。x-prediction 让网络直接输出 $\hat{z}_0=f_\theta(z_t,t,c)$，预测速度 $\hat{v}_\theta=(\hat{z}_0-z_t)/(1-t)$，目标 $z_0$ 按定义落在低维数据流形上、不显式依赖 $t$。输入端穿越离流形 $z_t$ 的弦仍存在，但输出端目标被限制在流形上。该性质本身并非 DINOv2 专属（JiT 在像素空间已用），本文独特之处在于把它与 DINOv2 各向同性逐坐标方差、近高斯边缘分布组合，二者共同界定目标并平滑其对 $z_t$ 的依赖，使原生 DiT 足够。表2实证：同架构同编码器下 x-pred 在 80/200/400 epoch 的 FID（2.63/1.89/1.70）始终优于 v-pred（3.17/2.11/1.86）。

方法步骤详情

一、冻结 DINOv2-S 编码器与 RAE ViT 解码器，得 patch 与 [CLS] token。二、逐元素标准化 $\tilde{z}=(z-\mu)/\sqrt{\sigma^2+\epsilon}$ 以降协方差条件数、放松 LayerNorm 范数约束；此步为先决条件，缺之训练完全发散（FID>300）。三、用 x-prediction flow matching 训 vanilla DiT，损失 $\mathcal{L}_{fm}=\mathbb{E}\|\hat{v}_\theta-v\|^2$，等价于 $(1-t)^{-2}$ 加权的 $\|\hat{z}_0-z_0\|^2$。四、联合 [CLS]-patch 建模，二者噪声独立采样并以耦合噪声初始化，总损失 $\mathcal{L}=\mathcal{L}_{fm}+\lambda\mathcal{L}_{cls}$（$\lambda=0.2$），推理只解码 patch。五、维度感知调度：对 logit-正态 $X$ 做时间平移 $X'=Xs/(1+(s-1)X)$，中位 $t$ 由 0.31 推到 0.17。六、推理用 Heun + CFG（=3.7）积分 ODE。

技术新颖性

技术新颖性分三层。其一，首次沿四条互补几何轴（本征维数/有效秩/边缘高斯性/流形插值）系统比较像素、SD-VAE、DINOv2，并用 SD-VAE 居中这一事实把「表征学习目标」与「单纯压缩」分离——像素与 DINOv2 环境维数相同（$D=196608$）但几何性质迥异，证明优势源于 SSL 训练而非维度压缩。其二，提出「目标侧」解决几何干扰的范式：相比 RAE 的架构侧 DDT 头、RJF 的传输侧黎曼匹配，RiT 仅在输出参数化上做改动，最小化侵入。其三，揭示表征空间生成的独特红利——[CLS] token 是表征本身固有的全局语义摘要，可被联合建模并作为上下文引导，这是 SD-VAE 潜空间不具备的。整套方法无新结构、新损失、新传输，所有组件均为标准件，新颖性在于「组合与配方」及对为何有效的几何解释。

**Figure 4**: RiT Arch. Frozen RAE encoder/decoder bracket a vanilla DiT trained by x-prediction.

**Figure 5**: Time sampling distribution p(t) and per-token SNR.

实验结果

几何刻画：像素与 DINOv2 本征维数相近（$\hat{d}\approx33$），但 DINOv2 有效秩高 7.3×、$t=0.9$ 处协方差条件数好 35×（$\kappa\approx56$ vs $2000$）、中位超额峰度低 11.5×、流形插值误差低 1.7×。收敛：RiT-XL 800 ep 达 FID 1.45，比同编码器同参数 RAE-XL(DINOv2-S) 的 1.87 好23%；100 ep 匹配 RAE-XL(DINOv2-B) 720 ep。消融：去标准化训练直接发散（FID 344.8）；去时间平移退到 3.17；去 CLS 退到 1.63；DINOv2-S(1.44) 反优于 DINOv2-B(1.56)。少步生成：像素截断误差仅 3.6× 衰减而 RiT 衰减 12.9×；10 NFE 下 JiT-H 26.2、DiTDH-XL 3.29、RiT 2.38；引导 5 步 FID 2.0、10 步 1.25。对比：RiT 无引导 1.45、有引导 1.14，均优于 DiTDH-XL(1.51/1.28)、FAE(1.48/1.29) 等表征方法。

**Table 1**: Marginal Gaussianity. Excess kurtosis across three representation spaces.

**Table 2**: x-pred vs v-pred. FID-50K, ImageNet 256², Heun 50 steps, w/o guidance.

**Table 3**: Ablation study. Default: x-prediction, element-wise standardization, time-shift schedule, λ=0.2, DINOv2-S.

**Table 4**: Sampling schedule ablation. FID-50K across six ODE time-discretization schedules and Heun step counts.

**Table 5**: Class-conditional image generation on ImageNet 256×256.

**Figure 6**: Convergence comparison on ImageNet 256². FID-50K vs training epochs.

**Figure 7**: Few-step FID at matched NFE.

**Figure 8**: DINOv2 ODE converges in few Heun steps. Pixel-space truncation error vs step count K.

**Figure 9**: Curated RiT-XL samples on ImageNet 256².

查看结构化数据

任务	指标	本文	基线	提升
ImageNet 256×256 类条件图像生成（无引导 FID-50K）	FID↓	1.45（RiT-XL, 676M, DINOv2-S, 800 ep）	RAE-XL(DINOv2-S) 1.87；DiTDH-XL 1.51；FAE-DINOv2-G 1.48	相比 RAE-XL(同编码器)好 23%；取得所有表征方法最佳无引导 FID
ImageNet 256×256 类条件图像生成（有引导 FID-50K）	FID↓	1.14（CFG=3.7, 25 Heun 步）	DiTDH-XL 1.28（839M）；FAE 1.29；REPA-XL 1.29；REG-XL 1.36	用 19% 更少参数（676M vs 839M）和最小 DINOv2 变体取得 SOTA
训练收敛速度	达到目标 FID 所需 epoch	RiT-XL 100 ep 达 RAE-XL(DINOv2-B) 720 ep 水平	RAE-XL(DINOv2-S/B)、REPA-XL、REG-XL	匹配编码器下 7× 加速；200 ep 匹配 RAE-XL(DINOv2-S) 800 ep（4× 加速）
少步生成（匹配 NFE）	FID-50K @ 10 NFE	RiT-XL 2.38（20 NFE 1.58）	JiT-H(像素) 26.2；DiTDH-XL(DINOv2-B) 3.29	10 NFE 较像素空间数量级改善；5 引导步即 FID 2.0、10 步 1.25，无需蒸馏
ODE 截断误差衰减	像素空间 $\\|x^{(K)}-x^{(ref)}\\|_F$ 衰减倍数（K=2→50）	RiT-XL 12.9×（113.4→8.8）	JiT-H(像素) 3.6×（87.1→23.9）	晚期斜率 -1.33 vs -0.91，对应更低有效曲率，解释少步可解性

局限与改进

作者承认的局限主要在评估范围：仅限 ImageNet $256\times256$ 类条件生成，未涉及文本到图像或更高分辨率；生成依赖外部 RAE 冻结解码器，最终像素质量受其重建能力上界约束；几何分析仅在 ImageNet 上完成，对其他数据分布的推广未验证。我观察到的额外局限：CLS 耦合噪声改善小于 0.1 FID，说明该机制边际作用小；论文未给出推理延迟/吞吐的实际测量，少步 FID 虽好但每步在 $D=196608$ 高维空间上的计算成本未必低于 VAE 潜空间少步；x-prediction 在高噪声（小 $t$）区目标的 $(1-t)^{-2}$ 重加权可能影响稳定性，论文以经验结果带过、缺理论界；无引导精度 0.81 但召回率仅 0.62，多样性偏弱，CFG 仍带来明显增益（1.45→1.14），与「CFG 变得不那么必要」的论述存在张力。

独立分析的弱点

弱点一：泛化性证据不足。所有结论基于 ImageNet 单一数据集，未在 COCO、LAION 或类条件之外的任务上验证，「四个几何轴」是否在自然文本到图像分布上同样成立未知。改进方向：在多样化数据上重复 Section 2 的几何测量。弱点二：架构红利被低估地绑定到 DINOv2+RAE 生态。RiT 的少步优势部分来自 RAE 解码器，换用其他解码器是否保持未测。改进：消融解码器选择。弱点三：与 FAE 的比较虽强调「正交」，但没有实际组合实验证明二者叠加有效，论断偏推测。改进：实现 FAE+RiT 组合验证。弱点四：高维计算开销未充分讨论——$D=196608$ 远大于 SD-VAE 的 4096，每步 NFE 成本更高，少步 FID 的「高效」需结合实际墙钟时间评估。改进：报告端到端采样延迟与显存。弱点五：召回率偏低（0.62），多样性受限。改进：研究调度或目标重加权以提升召回。

未来方向

作者明确提出的方向有二：一是 RiT 可作为蒸馏/一致性模型的天然教师——因基模型本身已具竞争力少步 FID，与渐进蒸馏、一致性模型、整流流等正交可叠加；二是推动「统一理解-生成」，指出当前统一视觉模型通常维护 CLIP/DINOv2（感知）与 SD-VAE（合成）两套编码器，RiT 能在单一 DINOv2 语义表征 + 单一 vanilla Transformer 上同时服务分类、检索与合成，降低附加生成头成本。基于成果可延伸的方向：把 Section 2 的四轴几何分析框架推广到文本到图像、视频、3D 等其他模态的表征空间，建立「什么样的潜空间对 diffusion 友好」的一般判据；探索 x-prediction 在更大骨干或更长训练下的缩放律；将 [CLS] 联合建模思想迁移到多模态条件（如文本 embedding）的扩散中。

复现评估

复现性总体良好。代码已开源在 https://github.com/lezhang7/RiT 。所有核心组件均为标准件：vanilla DiT 主干规格明确（28 层、隐藏维 1152、16 头、FFN 展开比 4、676M 参数），DINOv2-S 与 RAE 解码器均来自开源 RAE 工作可直接加载，噪声调度（logit-normal + 时间平移）、Heun 求解器、CFG 都有闭式定义，附录 D 给出完整超参。训练规模适中：ImageNet 上 8 块 H200，约 12 分钟/epoch、共 800 epoch，对工业界或大算力学术组可承受，但对独立研究者仍是较高门槛。几何分析（TwoNN 本征维数估计、PCA、峰度、插值 MSE）仅需 10,000 张 ImageNet 图像，复现成本低。潜在风险：FID-50K 评估需 class-balanced 采样与冻结解码器流水线，细节若与 RAE 完全对齐则可信；作者提供消融表与截断误差测量，关键数字可交叉验证。整体属「高可信、中高算力」的可复现工作。

RiT：原生扩散Transformer在表征空间中已足够 RiT: Vanilla Diffusion Transformers Suffice in Representation Space