PaperlayerView exact revision on arXiv

Source-linked AI summary

RiT: Vanilla Diffusion Transformers Suffice in Representation Space

Le Zhang, Ning Mang, Aishwarya Agrawal

arXiv:2605.21981v1cs.CV

TL;DR

Flow matching in pretrained representation spaces may benefit from geometry beyond intrinsic dimensionality, but this question has lacked direct evidence. RiT studies this geometry and trains a vanilla DiT with x-prediction on frozen DINOv2 features, achieving FID 1.45 without guidance and 1.14 with classifier-free guidance on ImageNet 256×256.

  • Problem

    The paper asks whether pretrained representation spaces offer distributions more favorable for flow-matching learning than pixel space despite comparable intrinsic dimensionality.

  • Method

    RiT trains a vanilla Diffusion Transformer with x-prediction flow matching on frozen DINOv2 features, using a dimension-aware noise schedule and joint [CLS]-patch modeling.

  • Results

    FID 1.45 without guidance and 1.14 with classifier-free guidance on ImageNet 256×256, using 19% fewer denoiser parameters than DiT^DH-XL.

  • Takeaways & Limitations

    Representation-space geometry can support architecturally simple flow matching and efficient few-step generation without distillation or consistency training.

  • Takeaways & Limitations

    RiT inherits DINOv2’s weaker feature resolution for fine textures, thin structures, and small objects because its encoder remains frozen.

Abstract

from arXiv · show

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.

1 Introduction

The introduction argues that pretrained representation spaces can make flow-matching x-prediction substantially better conditioned than pixel-space learning, despite comparable intrinsic dimensionality. It motivates RiT’s target-side solution: regress clean DINOv2 features on the data manifold rather than adding specialized architectures or transport paths.

  • Motivation: x-prediction regresses the clean data point on the low-dimensional manifold instead of the ambient-space velocity, extending JiT’s pixel-space insight to pretrained representations.This target-side parameterization keeps the regression target on-manifold even when intermediate flow states are off-manifold.
  • Geometric comparison: Both pixel and DINOv2 features have intrinsic dimensionality d̂ ≈33, but DINOv2 is near-isotropic, near-Gaussian, and supports approximately on-manifold linear interpolation.Pixel features are anisotropic, strongly non-Gaussian per coordinate, and their linear chords can cross low-density regions.
  • Geometric comparison: 7.3× higher effective rank, 35× better covariance conditioning, 11.5× lower excess kurtosis, and 1.7× lower interpolation error distinguish DINOv2 from pixels.SD-VAE latents fall consistently between pixel and DINOv2 features, suggesting representation-learning objectives—not compression alone—drive the advantage.
  • Geometric challenge: DINOv2’s norm-concentrated features make linear paths enter encoder-unseen ambient regions, where velocity targets acquire a large radial component and create radial ambiguity.Prior methods address this with a specialized DDT prediction head, a ViT decoder, or Riemannian Flow Matching, adding architectural or transport complexity.
  • Proposed approach: RiT resolves radial ambiguity at the output by regressing ẑ0 on the data manifold, leveraging DINOv2’s isotropic variance and near-Gaussian marginals.The reparameterization is not new; the introduction attributes its effectiveness to these favorable DINOv2 distributional properties.

2 The Geometry of Representation Spaces for Flow Matching

Pixel and DINOv2 have essentially identical intrinsic dimensionality, but DINOv2’s manifold is substantially better aligned with isotropic Gaussian noise for flow matching. Across effective rank, conditioning, Gaussianity, and interpolation geometry, SD-VAE is consistently intermediate.

  • Intrinsic dimensionality: Both pixel and DINOv2 spaces have intrinsic dimensionality d̂ ≈33, so their flow-matching difficulty differs primarily by manifold positioning rather than manifold size.The reported estimates are ˆd = 32.6 ± 0.8 for DINOv2 and a nearly identical pixel estimate.
  • Effective rank: 7.3× higher effective rank distinguishes DINOv2 from pixels, with effective ranks of 327 and 45 respectively, indicating more uniformly distributed variance.SD-VAE is intermediate at effective rank 98; DINOv2’s LayerNorm also places features near an approximately isotropic shell.
  • Optimization conditioning: 35× better covariance conditioning gives DINOv2 κ ≈56 versus pixel-space κ ≈2,000 at t=0.9, allowing variance directions to be learned at comparable rates.The distribution’s proximity to N(0, I) also reduces the irreducible variance of the per-pair velocity target.
  • Marginal Gaussianity: 11.5× lower excess kurtosis makes DINOv2 marginals markedly more Gaussian than pixels, with 98.7% satisfying |κj| < 0.5 versus 0% for pixels.SD-VAE is intermediate: 74.2% of dimensions satisfy the same threshold, and DINOv2’s median |κj| is 2.7× lower than SD-VAE.
  • On-manifold interpolation: 1.7× lower on-manifold interpolation error further favors DINOv2, whose representation-space interpolations remain close to the manifold while pixel interpolations produce ghosting artifacts.This joint-geometry measure complements the variance-based axes by testing whether linear chords leave the data manifold.

3 RiT: A Vanilla DiT for Representation-Space Diffusion

RiT applies vanilla DiT with x-prediction to standardized frozen DINOv2 representations, jointly modeling [CLS] and patch tokens. Dimension-aware noise scheduling and coarse ODE integration complete the representation-space diffusion design.

  • Architecture: RiT uses a vanilla DiT backbone trained with x-prediction on frozen DINOv2 features, with a frozen encoder and ViT decoder connecting representations and pixels.The backbone includes SwiGLU, RMSNorm, 2D-RoPE, QK-normalization, and in-context class tokens.
  • Preprocessing: Raw DINOv2 training diverges entirely, so RiT standardizes patch and [CLS] tokens element-wise before diffusion and inversely transforms them before decoding.The diagonal preconditioner reduces covariance conditioning and relaxes the near-constant-norm constraint imposed by LayerNorm.
  • Prediction objective: x-prediction makes the network directly predict the clean standardized feature z0, avoiding the off-manifold radial component and ambient-space target required by v-prediction.The corresponding predicted velocity is v̂θ = (ẑ0−zt)/(1−t), and the x- and v-losses coincide up to (1−t)^−2 reweighting.
  • Token modeling: RiT jointly diffuses the intrinsic [CLS] token and local patch tokens through shared self-attention, using separate prediction heads for ẑ0 and ẑcls,0.The [CLS] token provides a global semantic summary complementary to local patch content.
  • Noise schedule: 5× lower median SNR from the dimension-dependent time shift improves FID from 3.17 to 1.44 at 800 epochs versus the pixel-space logit-normal baseline.The schedule uses s = hwd/4096 ≈4.9 and shifts the median t from ≈0.31 to ≈0.17.

4 Experiments

RiT’s ablations identify dimension-aware scheduling, feature standardization, and joint [CLS]-patch modeling as key design choices, while experiments show faster convergence and strong few-step generation. RiT achieves state-of-the-art ImageNet quality with a smaller vanilla denoiser and DINOv2-S encoder.

  • Ablations: FID >300 throughout training: raw DINOv2 features diverge because heterogeneous channel variances span a ≳10^2 range.Element-wise standardization is therefore required for stable training.
  • Ablations: 3.17 → 1.44 at 800 epochs: the time-shift schedule closes a 2× FID gap over JiT’s original logit-normal schedule.The schedule reallocates training density toward higher noise for per-token dimensionality d=384 versus pixel d=3.
  • Ablations: FID 1.63 → 1.44: modeling the [CLS] token with λ=0.2 improves results over omitting CLS modeling.Attention visualization indicates that [CLS] aggregates scene cues, integrates object-context relations, and broadcasts refined guidance.
  • Ablations: FID 1.44 versus 1.56 at 800 epochs: DINOv2-Small consistently outperforms DINOv2-Base despite half the feature dimensionality.Both manifolds have intrinsic dimensionality ˆd ≈33, but the larger target does not provide a corresponding structural gain.
  • Few-step sampling: 12.9× versus 3.6×: RiT’s truncation error decays more steeply than JiT’s from K=2 to K=50.At K=25, RiT is within 22.2 Frobenius units of its 125-step endpoint, while its FID 1.47 matches the K=50 value 1.46.
  • Comparison: FID 1.45 unguided and 1.14 guided: RiT achieves the best reported representation-based ImageNet results with a 676M-parameter denoiser.The denoiser is 19% smaller than DiTDH-XL’s 839M parameters, and RiT uses DINOv2-S rather than larger encoder variants.

5 Related Work

This section situates RiT among diffusion and flow-matching methods, representation-based generation approaches, few-step sampling methods, and unified understanding–generation models. It emphasizes that RiT models DINOv2 features with a vanilla Transformer using x-prediction, without specialized architectural or geometric modifications.

  • Diffusion and flow matching for images: x-prediction outperforms ϵ- and v-prediction in pixel space by placing the target on the low-dimensional data manifold.JiT established this prediction-target advantage, which RiT extends to pretrained representation spaces.
  • Leveraging representations for generation: RiT operates natively in DINOv2 space, unlike REPA and REG, which use DINOv2 alignment or [CLS] entanglement while retaining v-prediction on SD-VAE trajectories.The related methods shape or align diffusion latents rather than directly modeling DINOv2 features.
  • Leveraging representations for generation: Element-wise standardization and x-prediction suffice for modeling DINOv2 features with a vanilla DiT, requiring no architectural modification or Riemannian reformulation.This is presented as RiT’s core distinction from prior DINOv2 generation methods.
  • Few-step generation and distillation: RiT reaches competitive few-step FID with its base model, without distillation or consistency loss, while remaining a natural teacher for such methods.Progressive distillation, consistency models, and rectified flow instead reduce sampling cost through dedicated students or trajectory straightening.
  • Toward unified understanding–generation: RiT suggests using a single semantic representation and vanilla Transformer backbone for both understanding and generation, without DDT heads, Riemannian reformulation, or alignment loss.This contrasts with unified vision models that typically separate perception and synthesis encoders and add task-specific generative components.

6 Conclusion

RiT is a vanilla DiT trained with x-prediction on frozen DINOv2 features, achieving strong ImageNet 2562 generation quality with fewer denoiser parameters than DiTDH-XL. It also supports efficient few-step generation without distillation or consistency training.

  • RiT trains a vanilla DiT with x-prediction on frozen DINOv2 features, using the smallest DINOv2 variant, DINOv2-S (d=384).
  • 1.45 FID without guidance and 1.14 with classifier-free guidance on ImageNet 2562, using 19% fewer denoiser parameters than DiTDH-XL (676M vs. 839M).
  • 2.0 guided FID at 5 Heun steps and 1.25 at 10 steps, without distillation or consistency training.

A Limitations · B Equivalence of Velocity Loss and Reweighted x-Prediction Loss

RiT is limited by frozen DINOv2 biases, restricted evaluation settings, and assumptions underlying its geometric analysis. Its velocity-MSE objective is functionally equivalent to a reweighted x-prediction loss, while parameterization determines the fitted function.

  • A Limitations: Frozen DINOv2 features may weaken RiT’s generation of fine textures, thin structures, and small objects because the encoder emphasizes semantic content over photometric detail.Joint encoder fine-tuning could mitigate these biases but would sacrifice the simpler frozen-encoder setup.
  • A Limitations: RiT has only been evaluated for class-conditional ImageNet generation at 256×256, leaving text conditioning, higher resolutions, and non-image modalities untested.The geometric analysis was also limited to ImageNet images using DINOv2-Base and DINOv2-Small scales.
  • A Limitations: Whether the four geometric axes persist at larger model or data scales, or under text conditioning, remains future work.This limitation follows from the restricted evaluation settings and analysis scales.
  • A Limitations: The covariance-conditioning diagnostic assumes a local Gaussian model p(z0) ≈ N(µ, H), while the other geometric summaries cannot exclude unfavorable manifold regions.The favorable properties may therefore fail in adversarial pockets of the manifold.
  • A Limitations: The flow-matching experiments corroborate the geometric claims empirically but do not establish that any individual geometric axis is necessary for the efficiency gains.This limits causal interpretation of the four diagnostics.
  • B Equivalence of Velocity Loss and Reweighted x-Prediction Loss: Under zt = t z0 + (1 − t) ϵ, velocity-MSE training with x-prediction parameterization is equivalent to a reweighted x-prediction loss.The equivalence follows by substituting the forward-process expression for ϵ and the predicted velocity derived from ˆz0.
  • B Equivalence of Velocity Loss and Reweighted x-Prediction Loss: The reweighting factor is (1 − t)^−2, which upweights errors at high t near clean data.The two objectives are equivalent as functionals, but their parameterizations determine which function the network fits.
  • B Equivalence of Velocity Loss and Reweighted x-Prediction Loss: x-prediction directly outputs ˆz0 on the data manifold, whereas v-prediction must produce the ambient velocity (z0 − zt)/(1 − t).This parameterization difference, rather than the functional loss alone, distinguishes the learned regression target.

C Manifold Analysis Details · D Architecture and Hyperparameters

The appendix formalizes manifold-analysis metrics and validates the pixel–DINOv2 intrinsic-dimensionality comparison with bootstrap uncertainty. It concludes that DINOv2’s advantage lies in global geometry rather than manifold dimensionality, while detailing an unbiased interpolation-error evaluation.

  • C Manifold Analysis Details: 10,000 randomly sampled ImageNet training images underpin all manifold-analysis experiments.The appendix provides formal definitions and implementation details for the metrics used in Section 2.
  • C Manifold Analysis Details: 512 PCA components are fit to flattened feature vectors, with effective rank measuring variance isotropy.Effective rank equals 1 for maximally anisotropic variance and k for perfectly uniform variance.
  • C Manifold Analysis Details: TwoNN estimates intrinsic dimensionality from first- and second-nearest-neighbor distance ratios under local manifold uniformity.The implementation subsamples 5,000 points and computes pairwise Euclidean distances in chunks.
  • C Manifold Analysis Details: 10 independent 5,000-point subsamples provide bootstrap means and standard deviations for TwoNN at high ambient dimensionality.This addresses nearest-neighbor distance concentration and noise in single-run estimates.
  • C Manifold Analysis Details: 1.0 dimension is the pixel–DINOv2 gap, below the combined standard deviation of approximately 1.5 dimensions.The reported z ≈ 0.7 and p ≫ 0.05 make the estimates statistically indistinguishable.
  • C Manifold Analysis Details: DINOv2’s advantage is attributed to global geometry rather than manifold dimensionality.The appendix therefore relies on TwoNN as the primary estimator because larger-k MLE variants can suffer upward bias at high ambient dimension.
  • C Manifold Analysis Details: 100 same-class pairs with 11 interpolation steps each are evaluated through an identical encode→decode→MSE pipeline.Pixel interpolation blends images before encoding, whereas DINOv2 interpolation blends representations before decoding and re-encoding.
  • C Manifold Analysis Details: A feature-space cosine-similarity sanity check tests whether DINOv2 interpolation errors are artificially reduced by encoder re-projection.The check compares each interpolant z_t with its re-encoded decoded representation.

D.1 Model Architecture … E Encoder Size Ablation

RiT uses a modernized DiT backbone over DINOv2 feature grids, with specified normalization, positional encoding, conditioning, and in-context-token components. DINOv2-Small outperforms DINOv2-Base despite half the feature dimensionality, reaching FID 1.44 versus 1.56 at 800 epochs.

  • D.1 Model Architecture: RiT operates on a 16×16 DINOv2 feature grid, using DINOv2-Small (d=384) for main experiments and DINOv2-Base (d=768) for ablations.Input dimensions are 384 × 16 × 16 or 768 × 16 × 16, respectively.
  • D.1 Model Architecture: Each DiT block combines adaLN conditioning, QK-normalized multi-head attention with VisionRoPE, and a SwiGLU feed-forward network.Timestep and class embeddings are summed before projection into per-layer scale and shift parameters.
  • D.1 Model Architecture: 32 learnable in-context tokens are injected at an intermediate layer, participate in later self-attention, and are discarded before final projection.They add capacity for class-conditional generation without changing the core DiT block.
  • D.2 Training Hyperparameters: RiT-XL training hyperparameters are summarized in Table 7 for ImageNet 2562.The supplied passage provides the table’s scope but no individual hyperparameter values.
  • D.3 Sampling Hyperparameters: Sampling hyperparameters are summarized in Table 8.The supplied passage provides the table’s scope but no individual sampling values.
  • D.4 Pseudocode: RiT pseudocode uses x-prediction and joint CLS modeling, with main experiments sampled by a second-order Heun solver.The figure shows Euler for clarity, while Heun averages velocity at the current and predicted next steps.
  • E Encoder Size Ablation: FID 1.44 versus 1.56 at 800 epochs is achieved by DINOv2-Small over DINOv2-Base despite half the feature dimensionality.The lower-dimensional latent space is easier to model, while TwoNN intrinsic dimensionality remains comparable across encoder sizes.

F Uncurated Sample Grid

RiT-XL produces uncurated ImageNet samples spanning diverse semantic categories, including animals, food, landscapes, vehicles, and plants. Using Heun sampling with 100 steps and classifier-free guidance scale 3.7, the samples demonstrate high-fidelity and diverse image generation.

  • F Uncurated Sample Grid: Heun sampling uses 100 steps with classifier-free guidance scale 3.7.The samples are generated by RiT-XL with a DINOv2-Small encoder after 760 training epochs.
  • F Uncurated Sample Grid: The uncurated samples demonstrate RiT’s ability to produce high-fidelity, diverse images across many semantic categories.The displayed categories include animals, food, landscapes, vehicles, and plants.
  • F Uncurated Sample Grid: 28 samples span diverse ImageNet categories, including animals, food, landscapes, vehicles, and plants.Examples include macaw, jellyfish, flamingo, golden retriever, pizza, volcano, and daisy.

G Sampling Schedule Analysis · H Random Samples

The paper evaluates six ODE time-discretization schedules for RiT sampling and examines their effects on FID across Heun step counts. It also reports uncurated random samples showing consistently high quality and diversity across ImageNet categories.

  • G Sampling Schedule Analysis: Six schedules map normalized step indices i/K to timesteps t_i ∈ [0, 1], from pure noise to clean data.K denotes the number of Heun steps.
  • G Sampling Schedule Analysis: Uniform distributes evaluations evenly, whereas EDM, power-2, and time-shift concentrate them near t=0.Cosine and log-SNR place more evaluations near both endpoints.
  • G Sampling Schedule Analysis: Schedules emphasizing the high-noise regime target the region where the velocity field varies most rapidly.Figure 12 compares these schedules through FID as a function of Heun step count.
  • H Random Samples: 192 randomly generated samples—8 per class across 24 ImageNet classes—are shown without curation or cherry-picking.Each row corresponds to one ImageNet class.
  • H Random Samples: RiT-XL produces consistently high-quality and diverse samples across all displayed categories.This qualitative assessment accompanies the uncurated random-sample grid.
  • H Random Samples: Figure 12 compares FID-50K against Heun step count without and with classifier-free guidance for RiT-XL on DINOv2-S.The evaluation uses coupled noise and 800 training epochs.
  • G Sampling Schedule Analysis: Table 10 provides a full sampling-schedule ablation across ODE schedules and Heun step counts, including 2-step sampling.Each cell reports independent and coupled noise results for RiT-XL on DINOv2-S.

I CLS–Patch Attention Analysis

RiT exhibits stage-wise bidirectional communication between [CLS] and patch tokens. [CLS] aggregates foreground and contextual information before refining salient details, while patches increasingly query it to support global guidance and generation quality.

  • CLS→Patch: In early layers, [CLS] attends to salient foreground regions, expands to contextual background regions in middle layers, and refocuses on critical details in late layers.Late-layer attention emphasizes details such as the head and eyes, which influence structural consistency and perceptual realism.
  • Patch→CLS: At deeper layers, semantically important patches increasingly query [CLS], whereas low-information background patches rely less on it.The asymmetric reliance concentrates global information exchange on semantically important regions.
  • CLS–Patch communication: Together, these interactions make [CLS] a global message hub that collects distributed evidence, integrates object–context relations, and broadcasts refined guidance to patches.The communication pattern is associated with improved object–background disentanglement and final generation quality.
Loading 2605.21981v1…