RiT: Vanilla Diffusion Transformers 在表示空间中即足够

TL;DR RiT 用 vanilla DiT 在 DINOv2 表示空间做 x-prediction 流匹配，利用表示学习赋予的优良几何性质，以更少参数和极少采样步数 (如 5 步 FID 2.0) 超越 DiT-DH-XL，证明表示空间比像素空间更适合作高效生成。

问题

问题背景

扩散模型在像素空间生成图像通常需要大型网络和数百步采样，计算成本高昂。流匹配 (flow matching) 与 x-预测 (回归干净数据点而非环境速度) 能利用低维流形结构提升效率，但在像素空间仍有统计性质不佳的瓶颈。

现有方法局限

• 像素空间扩散：高维像素空间具备低内在维度但有效秩低、协方差条件数差 (conditioning)，导致训练回归目标病态，需要大模型与大量步数。
• VAE 潜在空间：通过压缩获得更紧凑表示，但统计性质仅中等改善，未根本解决分布的高峰度与非高斯性。
• DINOv2 表示空间扩散的前期工作 (如 MDT、TACTO) 需特殊预测头或 Riemannian 传输，因为 DINOv2 特征被认为是非欧几何。这些额外组件增加了复杂度、限制了效率，且无法直接使用 vanilla DiT 训练。

为什么这个问题难/重要

• 几何分析缺失：预训练表示空间 (如 DINOv2) 的内在维度与像素空间相似 (~33)，但其有效秩、协方差条件数、峰度等统计量存在数量级差异 (秩高 7.3×、条件数优 35×、峰度低 11.5×)。这种分布差异使得流匹配回归更良态，但之前未被系统揭示。
• 模型简化与加速潜力：若在表示空间直接用简单 x-预测训练 vanilla DiT，可大幅减少参数 (676M vs. 839M)，并实现极低步数采样 (5 步 Heun 方法 FID 2.0，10 步 1.25)，无需蒸馏或一致性训练。这对实际部署至关重要。

行业类比

类似在 CLIP 嵌入空间中直接训练轻量扩散模型以生成高保真图像，无需繁杂的后处理或重训练编码器，显著降低推理成本。

核心洞察

• 表示空间的协方差条件数、有效秩和峰度等统计特性对流匹配训练的难易度影响远大于内在维度。该工作系统对比了像素、SD-VAE 和 DINOv2 特征，发现尽管像素与 DINOv2 的内在维度相近，后者的协方差条件数优 35 倍、有效秩高 7.3 倍、峰度低 11.5 倍，这些使 x-prediction 回归天生良态，无需 Riemannian 运输或专门预测头。这解释了为何普通扩散 Transformer 在 DINOv2 空间就能高效学习，也提示选择表示空间时应更多关注分布的正则性而非仅看压缩率。
• RiT 证明在预训练表示空间上执行 x-prediction 的流匹配时，极简架构足以超越复杂设计。仅用普通 DiT 底座，叠加维度感知噪声调度和联合 CLS-patch 建模，参数比 DiT-DH-XL 少 19%，FID 却从 1.48 降至 1.14（有引导）。这消除了对专用预测头、流形投影或蒸馏的依赖，大幅降低了工程负债，使表示空间扩散模型更易复现、扩展和部署。
• 良好的表示几何让 ODE 求解器在极少步数下收敛，实现快速采样。RiT 在 5 步 Heun 离散化下即达到 FID 2.0，10 步达 1.25，无需任何蒸馏或一致性训练。这直接得益于 DINOv2 特征的良态协方差结构与低插值误差，使生成过程可由粗粒度的确定性求解器高效逼近，为实时高保真图像生成提供了新的轻量路径。

方法

输入与特征提取

RiT 以 ImageNet 256×256 图像为输入，使用冻结的预训练 DINOv2 编码器提取表征，包括 patch tokens（序列长度为 N）和全局 [CLS] token。所有特征在送入扩散模型前经过 标准化（zero-mean unit-variance），以稳定训练并消除跨维度的尺度差异。

扩散过程与目标

在表征空间而非像素空间执行 flow matching，采用 x-prediction（直接回归干净数据点）而非常见的 velocity prediction。训练时对标准化特征按照 dimension-aware noise schedule 采样噪声：噪声水平随特征维度的增大而自适应缩放，使得高维特征仍能有效混合。加噪后的特征与时间步嵌入一同输入 vanilla Diffusion Transformer (DiT)。

联合 CLS–Patch 建模

DiT 同时处理 patch tokens 和 [CLS] token，所有 token 在自注意力层中交互。这一 joint [CLS]-patch modeling 允许模型利用 [CLS] token 携带的全局语义信息指导 patch 级别的生成，无需额外模块或条件注入设计。

损失与训练

预测目标为去噪后的干净特征，损失函数等价于 velocity loss 的重新加权形式（详见论文附录），但不需要 velocity 的显式计算。训练完全在冻结的 DINOv2 特征空间进行，不更新编码器。

采样与图像重建

采样时从纯噪声出发，应用 ODE solver（Heun steps） 逐步积分去噪，最终得到干净的 DINOv2 特征。之后通过一个 learned decoder（类似轻量卷积网络）将表征映射回像素空间，生成最终图像。得益于 DINOv2 特征空间的高有效秩、低过剩峰度等优良几何性质，ODE 离散化要求极低：仅 5 步 Heun 采样即可达到 FID 2.0。

与同类方法的区别

先前在 DINOv2 特征空间的扩散模型通常需要特殊预测头或黎曼流形上的传输来应对表征空间的各向异性，而 RiT 仅凭 vanilla DiT + x-prediction 结合维度感知噪声调度即达到更优性能，证明表征学习的几何优势本身足以取代复杂的架构适配。

实验

实验设计

在 ImageNet 256×256 上训练 RiT，基于冻结的 DINOv2 特征进行 x-prediction 流匹配，采用 vanilla DiT 架构，仅添加维度感知噪声调度与联合 [CLS]-patch 建模。评估生成质量（FID）、收敛速度、少步采样能力，并与 DiT^DH-XL 等先前工作对比。

关键发现

• RiT 无引导 FID 达 1.45，分类器引导下 FID 1.14，参数量比 DiT^DH-XL 少 19%（676M vs 839M），性能却更优。
• 表示空间的优越几何特性（有效秩、协方差条件数、峰度、流形插值误差）使得流匹配回归条件良好，无需 先前方法所用的特殊预测头或黎曼传输。
• ODE 求解器在高散化步数下效率极高：5 步 Heun 即达 FID 2.0，10 步 至 1.25，无需蒸馏或一致性训练。

对比基线深度解读

相较于 DiT^DH-XL 在像素空间使用标准 DiT，RiT 证明了表示空间的分布更适合流匹配学习，而非仅仅是降维。这与 Li et al.(2025) 在像素空间中利用低维流形的发现一脉相承，但通过 DINOv2 这种对比学习表示，获得了更有利的统计特性。消融实验进一步确认：维度感知噪声调度和联合 CLS-patch 建模对性能至关重要，每个设计选择都是必要的。这一结果暗示，在表示空间进行生成建模，结合简单的扩散 Transformer，即可在图像生成中达到新高度，并为高效率采样开辟道路。

行业影响

落地场景

RiT 的轻量级扩散生成范式适用于需要高效图像合成与编辑的产品线：

• 内容创作工具：广告素材生成、社交媒体配图、虚拟试穿等场景，可借助 RiT 在消费级 GPU 上实现秒级出图，降低云端推理延迟和成本。
• 大规模数据增强：在自动驾驶、医疗影像等需要合成高质量训练样本的领域，用 RiT 替代传统扩散模型，可在 5~10 步内生成逼真图像，提升数据闭环效率。
• 实时交互式应用：如在线设计平台、教育互动白板中的 AI 画笔，受益于 RiT 的 ODE 解算器，少量 Heun 步数即可稳定生成。

商业价值

• 推理成本骤降：传统扩散模型常需 50+ 步生成，RiT 在 5 步内达到 FID 2.0，10 步内 1.25，大幅降低 GPU 时长。对于日调用百万次的云服务，可直接减少 80% 以上推理算力支出。
• 模型部署轻量化：RiT 以 676M 参数超越 839M 的 DiT-DH-XL，同等质量下模型更小，利于移动端或边缘设备部署，拓宽了 AIGC 产品的覆盖终端。
• 训练效率提升：在 frozen DINOv2 特征上做 flow matching，无需复杂的 Riemannian 传递或预测头，训练简单且收敛快，可缩短产品迭代周期。

与现有技术栈的集成

RiT 具有 即插即用 特性，与主流扩散模型流水线兼容：

• 可直接替换 DiT 类架构的生成引擎，配合现成的 VAE 解码器 (如 SD-VAE) 将 DINOv2 特征重建为图像。
• 与 classifier-free guidance 无缝结合，接口标准，无需改动调度器或采样器代码。
• 使用预训练 DINOv2 编码器作为冻结的特征提取器，可与各类图像理解模型共享 backbone，实现“编码一次，多任务使用”。

具体落地用例

1. 电商产品图生成
   平台需要为大量商品生成不同背景、角度的营销图片。RiT 可在 5 步内生成高质量图像，显著降低批量生成耗时，同时保持输出一致性。冻结的特征空间还能与商品分类 embedding 联动，方便按类目控制生成风格。

2. 医疗图像合成与数据增强
   医学影像数据稀缺且标注昂贵，RiT 可用于生成病理切片、X 光图像等，少量采样步数即可产出多样本，帮助训练下游分割/分类模型。由于表示空间维度低且几何性质佳，合成图像分布偏移更小，更贴近真实数据流形。

局限

• - **冻结编码器限制**：RiT 使用**冻结的 DINOv2 编码器**将图像映射到表示空间，训练时编码器参数不更新。这固然隔离了扩散模型与原始像素的复杂性，但也意味着生成模型的性能上限受限于该预训练表示的质量与泛化性。若下游任务的数据分布与 DINOv2 训练分布存在明显偏移，重建保真度可能下降。此外，编码器固定使得整个系统无法端到端针对感知或生成目标进行联合优化，灵活性弱于可微分解码器方案。
• - **分辨率与扩展性未充分验证**：实验仅在 **ImageNet 256×256** 上进行，缺乏对更高分辨率（如 512×512 或 1024×1024）及多尺度训练的系统探索。表示空间扩散模型的维度与计算开销随特征图尺寸增长迅速，文中未讨论在更高分辨率下的内存占用、训练效率以及如何适配多尺度架构。对于追求高分辨率生成的工业级应用，该文的验证范围仍显不足。
• - **评估指标单一与分布外泛化未知**：主要指标使用 **FID** 评估样本真实感，但 FID 对分布偏差和异常值敏感，且不能反映语义一致性、属性编辑能力或多样性-保真度权衡的细粒度表现。此外，论文未评估在分布外（OOD）图像或跨域生成场景下的性能，难以判断 RiT 是否仅仅在训练分布上过拟合，还是会继承 DINOv2 表示的鲁棒性优势。