Learning on the Manifold: Unlocking Standard Diffusion Transformers with Representation Encoders

Conference: ECCV2026
Paper: ECCV Paper
Code: https://github.com/amandpkr/RJF
Area: Image Generation
Keywords: representation encoders, spherical latent space, Riemannian flow matching, Jacobi regularization, diffusion Transformer

TL;DR

RJF replaces straight-line flow matching in representation latent spaces with spherical geodesic flow matching and weights training by curvature-dependent error propagation, enabling a standard-width, 131M-parameter DiT-B to achieve FID 4.95 without guidance and 3.37 with guidance on ImageNet 256 class-conditional generation, both after 200 epochs.

Background & Motivation

Latent diffusion typically compresses images with a VAE before training a generative model to recover the image distribution in that space. Optimized for reconstruction, VAE latents do not necessarily exhibit the strong semantic organization of representation encoders such as DINOv2, motivating methods such as REPA to align generative models with pretrained features. RAE instead replaces the VAE encoder with a frozen representation encoder and pairs it with an image decoder, allowing the generator to model semantic features directly. However, replacing latents with high-dimensional representations does not guarantee that standard DiTs converge. Whereas RAE addresses this difficulty by increasing model width, this paper asks whether the real requirement is greater capacity or a change to the task imposed by the training objective.

The authors observe that 768-dimensional DINOv2-B features have highly concentrated norms, approximately occupying a spherical shell of radius \(\sqrt{768}\) rather than filling the ambient Euclidean space. Standard flow matching nevertheless connects data and noise with straight lines: even if both endpoints are projected onto the sphere, the connecting segment passes through its interior. The network must therefore learn semantic directions while also fitting radial variation introduced by the training path. In single-image overfitting experiments, models with widths of 384 or 512 fail under the full Euclidean loss but learn the angular component after radial error is removed. This supports a mismatch between the objective and data geometry, but does not establish that all capacity limitations disappear in large-scale generation.

The proposal is therefore not simply to normalize features, but to make the path, supervision, and sampling consistent with the same spherical geometry. Normalizing a noise endpoint makes the starting point valid without making intermediate states valid; even after correcting the path, velocity errors at different times may have different consequences because of curvature. Core Idea: eliminate off-manifold training paths with spherical geodesics, then use Jacobi-field-inspired weighting to account for error propagation, so standard DiTs devote capacity to learning semantic directions instead of compensating for artificial radial motion.

Method

Overall Architecture

Riemannian Flow Matching with Jacobi Regularization, abbreviated RJF, retains the standard diffusion Transformer backbone. A frozen representation encoder extracts features from training images. Both features and Gaussian noise are normalized to the unit sphere, and a Spherical Geodesic Path constructs intermediate states and tangent-velocity targets; Jacobi Curvature Weighting affects only the training loss. At inference time, Geodesic Sampling and Decoding starts from spherical noise, progressively generates features, and passes them to an RAE decoder to produce images.

The conditioning signal is an ImageNet class, not a text prompt. The paper defines the data endpoint as \(t=0\) and the noise endpoint as \(t=1\), so generation integrates backward from noise to data; following the training interpolation direction would instead move an image toward noise. Real images provide training supervision, but inference does not require a real image or its encoded features for the sample being generated.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Training images and Gaussian noise"] --> B["Frozen encoder<br/>and unit-sphere normalization"]
    B --> C["Spherical Geodesic Path"]
    C -->|Intermediate state, time, and class| D["Standard DiT<br/>predicts velocity"]
    C -.->|Tangent target, training only| E["Jacobi Curvature Weighting"]
    D -.->|Velocity error, training only| E
    N["Inference: spherical noise and class"] --> F["Geodesic Sampling and Decoding"]
    D -->|Trained velocity field| F
    F --> G["Generated image"]

Key Designs

1. Spherical Geodesic Path: change the supervision trajectory, not just endpoint distributions

The authors first normalize each feature vector and sampled Gaussian noise vector to place both on the same unit sphere. Spherical linear interpolation, or SLERP, then connects the endpoints along the shorter great-circle arc to construct a training state at any time. This is not synonymous with normalizing a straight-line interpolation afterward: SLERP parameterizes angular distance, defining both the path and its target velocity through spherical geometry. Consequently, model inputs retain unit norm rather than taking the low-norm intermediate values introduced by a chord passing through the sphere's interior.

The target velocity is obtained by differentiating this geodesic path with respect to time and lies in the tangent space at the current position. On the unit sphere, tangency means that position and velocity are orthogonal:

\[ v_t \cdot x_t = 0. \]

This condition means that local motion changes direction rather than moving inward or outward. The paper trains DiT with the squared error between predicted velocity and geodesic target velocity; the metric induced on the sphere allows this error to be evaluated in the ambient space. A distinction matters here: a target with no radial velocity does not imply that arbitrary network outputs are automatically tangent. The former follows from the path, while the latter also requires output constraints or implementation details that the main text does not fully specify. Target tangency alone therefore does not establish that the radial component of prediction error is identically zero.

2. Jacobi Curvature Weighting: make velocity-error importance reflect geometric propagation

A geodesic path determines where to learn, but does not fully establish whether errors everywhere are equally important. On a positively curved space such as a sphere, neighboring geodesics can converge, so the effect of a velocity perturbation on a later position differs from linear accumulation in a flat space. Jacobi fields describe how nearby geodesics separate under perturbations. The authors use this idea to derive a weight depending on time and the angle between endpoints, reweighting flow-matching error rather than adding another semantic teacher or feature-alignment network.

Under the paper's time convention, this weight reduces the relative contribution of errors near the data endpoint and emphasizes velocity accuracy near the noise endpoint. Intuitively, sampling starts at noise, where the model determines its initial semantic direction; curvature-induced focusing also changes the eventual effect of subsequent perturbations. Equation (13) expresses the weight using a squared sinc form, and Equation (14) applies it to velocity error. The local PDF text contains missing parentheses and merged symbols, so this note retains the mechanism and weighting direction without presenting a repaired expression as the authors' exact formula. The ablation shows that this step refines the path correction rather than supplying most of the improvement.

3. Geodesic Sampling and Decoding: keep generation on the sphere used for training

If training uses spherical paths but inference simply adds a tangent velocity to the current position with a Euclidean Euler step, a finite step can still leave the sphere. The authors therefore use exponential-map integration: rotate the current position along the predicted tangent direction by an angle determined by velocity magnitude and step size. For a given tangent vector, this update stays on a great circle. It preserves the geometry of an individual step, but does not imply an exact finite-step solution of the entire time-varying neural ODE. The authors additionally use rotation and normalization to address numerical drift, with 50 sampling steps in the experiments.

The resulting unit directions cannot simply be fed to the decoder as though they were original encoder features. The authors examine the final feature projection radius: restoring the original DINOv2-B norm of approximately 27.7 gives FID 7.79, whereas a radius of approximately 45 reaches 6.77. Thus, a decoder can be sensitive to magnitude even when semantic information is primarily angular. This norm adjustment belongs to decoding and is distinct from learning directions on the unit sphere during training. It is an important reproduction setting, and its optimal value does not follow directly from the spherical hypothesis.

A Worked Example

Consider one ImageNet class and a \(256\times256\) training image. DINOv2-B produces a set of 768-dimensional feature vectors for the image. Focus on one vector for illustration; the actual generator jointly processes the spatial token sequence rather than generating each token independently. The feature vector is normalized, a Gaussian vector is sampled and normalized, and the class label conditions DiT.

To visualize the problem with a straight-line path, consider two approximately orthogonal unit endpoints. The following is a geometric illustration, not an additional experimental result:

\[ \left\|\frac{x+\epsilon}{2}\right\| \approx \frac{1}{\sqrt{2}} \approx 0.707. \]

The straight-line midpoint is therefore well inside the sphere, whereas the SLERP midpoint remains on its surface. Training samples a time, asks DiT to predict the tangent velocity from that spherical state, and measures the error using Jacobi weighting. Inference no longer uses the real image: it starts with new spherical noise tokens, conditions on a class, performs 50 geodesic updates to generate directions, adjusts feature magnitude, and reconstructs an image with the RAE decoder.

Loss & Training

The final objective is geometrically weighted velocity-matching error, not an image-pixel reconstruction loss or an additional DINO feature-distillation loss. The representation encoder is frozen, and the generative pipeline uses an RAE decoder. The reported 131M, 459M, and 677M parameter counts describe the corresponding generative backbones; they should not be interpreted as total system costs including the encoder, decoder, and guidance model.

Default training follows LightningDiT: ImageNet-1K at \(256\times256\) resolution, 80 epochs, and a global batch size of 1024. Adam uses a learning rate of \(2\times10^{-4}\) and momentum coefficients of 0.9 and 0.95; gradient norms are clipped to 1.0, and EMA decay is 0.9995. The authors also retain RAE's dimension-dependent noise-schedule shift with \(n=4096\), which is not the same quantity as the 768-dimensional DINO features.

The small model also receives an extended 200-epoch run. Guided results use Autoguidance with scale 1.6, so they should not all be interpreted as classifier-free guidance results. The paper does not detail the extra resource cost of the auxiliary guidance model here. Consequently, avoiding width scaling is a backbone-architecture claim, not a directly measurable percentage saving in end-to-end cost.

Key Experimental Results

Main Results

The task is ImageNet-1K \(256\times256\) class-conditional generation, evaluated using 50,000 generated images. FID measures the difference between real and generated feature distributions, with lower values preferred; the main text does not clearly identify the exact split used for the reference FID statistics. The following results come from the paper's Table 1. Each same-scale comparison uses DINOv2-B features, 80 epochs, no guidance, and matched backbone parameter counts within the row.

Generative backbone	Parameters	EFM FID	RJF FID	Absolute FID reduction
DiT-B	131M	24.21	6.77	17.44
DiT-L	459M	6.31	4.21	2.10
DiT-XL	677M	4.28	3.62	0.66

In Table 2, the unguided XL results also include IS 186.2, Precision 0.82, and Recall 0.52. Higher IS generally indicates stronger class recognizability and class diversity, while Precision and Recall characterize generated-distribution quality and coverage, respectively. Within the same 80-epoch table block, REPA-E has unguided FID 3.46 and REG has 3.40, both lower than RJF's 3.62. The advantage should therefore not be described as the best FID across all methods; Recall of 0.52 also cautions against inferring the best coverage from high Precision alone.

Ablation Study

The following results follow the paper's Table 3 and the accompanying text in Section 4.3. All use standard DiT-B/1 with DINOv2-B under the evaluation setting above; only the last row uses guidance. SN projects noise and latents onto the sphere, RFM changes the path to a geodesic, and RJF additionally introduces Jacobi weighting.

Config	Epochs	FID	Note
Euclidean flow matching EFM	80	24.32	Straight-line baseline
EFM + SN	80	21.99	Valid endpoints, but an interior-crossing path
RFM	80	7.06	Geodesic paths provide the main improvement
RJF	80	6.77	Jacobi weighting reduces FID by another 0.29
RJF	200	4.95	Unguided result after longer training
RJF + Autoguidance	200	3.37	Guidance scale 1.6

The source contains numerical inconsistencies: the B-model EFM baseline is 24.21 in Table 1 but 24.32 in Table 3. Section 4.2 states an XL improvement of 1.19 over the Euclidean baseline, but subtracting 3.62 from 4.28 in Table 1 gives 0.66. This note preserves the values from each table separately, computes differences within the relevant table, and neither merges the baselines nor repeats the incorrect difference.

Key Findings

• Changing the path provides the largest gain: FID falls from 21.99 with SN to 7.06 with RFM in the same ablation, much more than the 0.29 improvement from RFM to RJF.
• The geometric benefit persists at XL scale but is smaller in absolute terms than at B scale. Complete convergence failure should not be generalized to every model size, since XL EFM already reaches 4.28.
• Section 5.2 reports baseline/RJF FIDs of 130.21/10.39 for SigLIP and 50.48/19.82 for MAE. These support transfer across encoders without implying equivalent final quality across representations.

Highlights & Insights

• The work turns the difficulty of generating high-dimensional features into a specific conflict in training-path geometry, offering an actionable diagnostic. Separating radial and angular error is more explanatory than immediately attributing failure to an undersized model.
• The ablation separates endpoint normalization, geodesic paths, and curvature weighting. Its most transferable lesson is to inspect intermediate states, not just the marginal distributions of real data and noise.

Limitations & Future Work

• The authors explicitly cite computational constraints and primarily compare methods within an 80-epoch budget, extending the small model to 200 epochs. This does not establish comprehensive superiority over strong, long-trained methods under fully matched budgets.
• This note's assessment: single-image overfitting supports geometric interference but does not rule out capacity bottlenecks in large-scale, multimodal distributions. An approximate shell also does not establish that the true semantic distribution uniformly covers the sphere.
• This note's assessment: tangent-output handling, decoder-radius selection, and extra guidance cost need more complete reproduction details. A useful next step is to measure actual geometry after normalization rather than assuming that every LayerNorm output strictly follows the same spherical model.
• This note's assessment: validation mainly covers ImageNet 256 image generation and does not directly establish transfer to text-to-image, video, or other resolutions. Multiple random seeds and training-time measurements would also clarify the stability of the gains.

Related Work & Insights

• vs RAE: Both use representation encoders and decoders. RAE addresses high-dimensional generation through width design, whereas RJF changes probability paths and training objectives. They are not mutually exclusive; the paper also observes smaller additional gains on a width-scaled architecture.
• vs REPA / REPA-E: REPA aligns intermediate diffusion features, while REPA-E jointly optimizes the VAE and generator. RJF models representation space directly, so its gains should not be attributed to an additional semantic-alignment loss.
• vs general Riemannian flow matching: Geodesic flow matching already exists. The contribution here centers on diagnosing representation-encoder generation failures and empirically combining Jacobi weighting with standard DiTs.

Rating

• Novelty: 4/5. Connecting geometric diagnosis to representation-space generation is valuable, but foundational Riemannian flow matching is not new.
• Experimental Thoroughness: 3/5. The study spans model scales, encoders, and component ablations, but training budgets, reproduction details, and numerical inconsistencies constrain the conclusions.
• Writing Quality: 3/5. The central intuition is accessible, but SOTA wording, arithmetic, and statements of geometric guarantees need greater care.
• Value: 4/5. The work offers a clear, testable direction for representation-space generation without widening the backbone.