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

TL;DR

This work identifies a geometric mismatch—the Geometry Gap—between Euclidean diffusion objectives and hyperspherical latent spaces produced by representation encoders, which causes standard diffusion transformers to fail to converge. It introduces Riemannian Flow Matching with Jacobi Regularization (RJF), enforcing geodesic trajectories on the hypersphere and curvature-aware loss weighting to account for geodesic focusing. RJF enables conventional Diffusion Transformer architectures to learn on representation encoders without width scaling, achieving state-of-the-art FID and IS on ImageNet (e.g., FID $=3.37$ with guidance and $=4.95$ without) and demonstrating robust generalization across architectures and encoders. By aligning the generative process with the latent topology, RJF delivers efficient, high-fidelity generation, validating the practicality of diffusion models directly in representation spaces.

Abstract

Leveraging representation encoders for generative modeling offers a path for efficient, high-fidelity synthesis. However, standard diffusion transformers fail to converge on these representations directly. While recent work attributes this to a capacity bottleneck proposing computationally expensive width scaling of diffusion transformers we demonstrate that the failure is fundamentally geometric. We identify Geometric Interference as the root cause: standard Euclidean flow matching forces probability paths through the low-density interior of the hyperspherical feature space of representation encoders, rather than following the manifold surface. To resolve this, we propose Riemannian Flow Matching with Jacobi Regularization (RJF). By constraining the generative process to the manifold geodesics and correcting for curvature-induced error propagation, RJF enables standard Diffusion Transformer architectures to converge without width scaling. Our method RJF enables the standard DiT-B architecture (131M parameters) to converge effectively, achieving an FID of 3.37 where prior methods fail to converge. Code: https://github.com/amandpkr/RJF

Figures (8)

• Figure 1: Bridging the Geometric Gap. We demonstrate that respecting the intrinsic geometry of pre-trained representations encoders enables the use of standard Diffusion Transformers without any architectural modification such as Width Scaling zheng2025diffusion. Our method, Riemannian Flow Matching with Jacobi Regularization (+DiNO+RJF), achieves an FID of 4.95 using standard LightingDiT-B yao2025reconstruction architecture without guidance, significantly outperforming the VAE-based LightingDiT-B (FID 15.83). In contrast, applying standard Flow Matching to DINOv2-B features (+DiNO) fails to converge (FID 21.64) due to Geometric Interference. Even restricting the noise to the hypersphere to strictly learn the angular component (+DiNO+SN) yields only marginal improvement (FID 19.07), as the Euclidean linear paths still traverse the low-probability interior of the feature manifold.
• Figure 2: Geometric Trajectories on the Hypersphere. Visualization of flow matching paths on the manifold $\mathcal{S}^{d-1}$. Standard Euclidean Flow Matching constructs linear paths that ignore the manifold geometry. Whether targeting standard Gaussian noise $\epsilon$ (orange) or projecting noise onto the sphere $\epsilon_s$ (purple) to strictly learn the angular component, the linear interpolation forms a chord that cuts through the low-density interior. This forces the model to learn a velocity field in undefined regions regardless of the endpoint. In contrast, Riemannian Flow Matching follows the geodesic (blue curve), ensuring the intermediate state $x_t$ remains strictly on the manifold surface. The resulting velocity field $u_t^M(x_t)$ is correctly defined within the tangent space (pink plane), naturally respecting the geometry of the representations.
• Figure 3: The Geometry Gap. A comparison of radial feature norms ($r = \|z\|_2$) between DINOv2-B representations and a standard Gaussian prior in $\mathbb{R}^{768}$. While the Gaussian prior (blue) is distributed across a diffuse shell, DINOv2-B features (orange) are rigidly constrained to a hypersphere with near-zero radial variance. This extreme geometric mismatch prevents standard diffusion models from converging effectively.
• Figure 4: Geometric Interference vs. Capacity. We train DiT-S models of varying widths on DINOv2 tokens ($d=768$). Top Row: When minimizing Euclidean MSE, narrower models ($d < 768$) suffer from collapse; the Angular Loss (semantics) gets stuck. Bottom Row: When the radial loss is ignored, even narrow models ($d=384$) converge perfectly on the angular component. This proves the bottleneck is not the dimensionality of the data, but the geometric conflict in the objective.
• Figure 5: Qualitative results of LightingDiT-XL+RJF trained for 80 epochs on ImageNet 256$\times$256. We show uncurated results on the five classes .