Research

RiT: Vanilla Diffusion Transformers Suffice in Representation Space

arXiv:2605.21981v1 Announce Type: new Abstract: Flow matching with x-prediction -- regressing the clean data point rather than the ambient velocity -- is known to exploit low-dimensional manifold stru

DGX agentpaper
researcharxiv-cs-cv

arXiv:2605.21981v1 Announce Type: new Abstract: 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 ite{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.3imes higher effective rank, 35imes better covariance conditioning, 11.5imes lower excess kurtosis, and 1.7imes 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 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 exttt{[CLS]}-patch modeling. On ImageNet 256{imes}256, RiT attains FID 1.45 without guidance and 1.14 with classifier-free guidance, outperforming DiT^ext{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.

Source: arXiv cs.CV | 2026-05-22

Loading related sources…