RiT: Vanilla Diffusion Transformers Suffice in Representation Space

📄 arXiv: 2605.21981v1 📥 PDF

作者: Le Zhang, Ning Mang, Aishwarya Agrawal

分类: cs.CV

发布日期: 2026-05-21

🔗 代码/项目: GITHUB

💡 一句话要点

RiT：仅用Vanilla Diffusion Transformer在表征空间实现高效图像生成

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱四：生成式动作 (Generative Motion)

关键词: 扩散模型 图像生成 Transformer 表征学习 Flow Matching DINOv2 自监督学习

📋 核心要点

  1. 现有像素空间Flow Matching方法在利用低维流形结构方面表现出色，但可能存在训练困难。
  2. RiT通过在预训练DINOv2特征空间进行Flow Matching，利用其优良的统计特性简化训练过程。
  3. 实验表明，RiT在ImageNet图像生成任务上，以更少的参数实现了优于DiT-XL的性能。

📝 摘要（中文）

本文研究了在预训练表征空间中进行Flow Matching学习的有效性。通过比较像素空间、SD-VAE和DINOv2特征，发现DINOv2特征具有更高的有效秩、更好的协方差条件、更低的超额峰度和更低的流形插值误差，尽管其内在维度与像素空间相似。这些统计特性使得Flow Matching回归具有良好的条件，无需像之前的DINOv2扩散方法那样使用专门的预测头或黎曼传输。因此，本文提出了表征图像Transformer（RiT），它是一个Vanilla Diffusion Transformer，通过在冻结的DINOv2特征上进行x-prediction训练，并辅以维度感知的噪声调度和联合[CLS]-patch建模。在ImageNet 256×256上，RiT在无引导的情况下达到FID 1.45，在分类器自由引导的情况下达到1.14，优于参数量更大的DiT^DH-XL。

🔬 方法详解

问题定义：现有的扩散模型，如DiT，通常在像素空间或VAE的隐空间进行训练。然而，直接在像素空间训练可能面临数据分布复杂、训练不稳定的问题。而基于DINOv2等预训练模型的扩散方法，通常需要复杂的预测头或黎曼传输来处理表征空间的特性，增加了模型的复杂度和计算成本。

核心思路：本文的核心思路是利用预训练的DINOv2特征空间的良好统计特性，例如更高的有效秩、更好的协方差条件和更低的超额峰度，来简化扩散模型的训练。通过在DINOv2特征空间进行Flow Matching，可以避免复杂的预测头或黎曼传输，从而降低模型的复杂度和计算成本。

技术框架：RiT的整体框架包括以下几个主要步骤：1) 使用预训练的DINOv2模型提取图像的特征；2) 在DINOv2特征空间中，使用Diffusion Transformer进行Flow Matching训练，目标是预测clean data point (x-prediction)；3) 使用维度感知的噪声调度来适应DINOv2特征空间的维度；4) 采用联合[CLS]-patch建模来提高模型的性能。

关键创新：本文最重要的技术创新点在于，通过在预训练的DINOv2特征空间进行Flow Matching，可以显著简化扩散模型的训练过程，并获得更好的性能。与现有的方法相比，RiT无需复杂的预测头或黎曼传输，从而降低了模型的复杂度和计算成本。

关键设计：RiT的关键设计包括：1) 使用维度感知的噪声调度，根据DINOv2特征空间的维度来调整噪声的添加方式；2) 采用联合[CLS]-patch建模，将[CLS] token与图像patch一起输入到Transformer中，从而提高模型的性能；3) 使用x-prediction作为Flow Matching的目标，即直接预测clean data point，而不是预测速度场。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

RiT在ImageNet 256×256图像生成任务上取得了显著的成果。在无引导的情况下，RiT达到了FID 1.45，在分类器自由引导的情况下达到了1.14，优于参数量更大的DiT^DH-XL（676M vs. 839M）。此外，RiT的ODE可以使用较少的Heun steps进行求解，例如，在分类器自由引导的情况下，仅需5个Heun steps即可达到FID 2.0，10个steps即可达到1.25。

🎯 应用场景

RiT具有广泛的应用前景，包括图像生成、图像编辑、图像修复等。由于其高效性和高性能，RiT可以应用于资源受限的设备上，例如移动设备和嵌入式系统。此外，RiT还可以作为一种通用的图像表征学习方法，用于其他计算机视觉任务。

📄 摘要（原文）

Flow matching with $x$-prediction -- regressing the clean data point rather than the ambient velocity -- is known to exploit low-dimensional manifold structure effectively in pixel space \cite{li2025back}. We ask whether a pretrained representation space, while containing a low-dimensional data manifold of comparable intrinsic dimensionality, offers a distribution more favorable for flow-matching learning. Comparing pixel, SD-VAE, and DINOv2 features along four geometric axes, we find that pixel and DINOv2 share nearly identical intrinsic dimensionalities (both $\hat{d}!\approx!33$) yet DINOv2 exhibits $7.3\times$ higher effective rank, $35\times$ better covariance conditioning, $11.5\times$ lower excess kurtosis, and $1.7\times$ lower on-manifold interpolation error; SD-VAE latents are consistently intermediate, indicating that the advantage stems from representation-learning objectives rather than mere compression. These statistical properties render the flow-matching regression well-conditioned and remove the need for the specialized prediction heads or Riemannian transport used by prior DINOv2 diffusion methods. We propose the \emph{Representation Image Transformer} (RiT): a vanilla Diffusion Transformer trained by $x$-prediction on frozen DINOv2 features, augmented only by a dimension-aware noise schedule and joint \texttt{[CLS]}-patch modeling. On ImageNet $256{\times}256$, RiT attains FID 1.45 without guidance and 1.14 with classifier-free guidance, outperforming DiT$^\text{DH}$-XL with $19\%$ fewer parameters (676M vs.\ 839M). The resulting ODE is efficiently solvable at coarse discretizations: with classifier-free guidance, $5$ Heun steps already reach FID 2.0 and $10$ steps reach 1.25, without distillation or consistency training. Code at https://github.com/lezhang7/RiT.