PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Amandeep Kumar, Vishal M. Patel

arXiv:2602.10099v2cs.LGcs.CV

TL;DR

Standard diffusion transformers fail to converge directly on representation-encoder features because Euclidean flow-matching paths conflict with their hyperspherical geometry. RJF addresses this by combining manifold geodesics with Jacobi Regularization, achieving 3.37 FID with a standard 131M-parameter DiT-B where standard flow matching fails to converge.

  • Problem

    Standard diffusion recipes fail to converge effectively on high-dimensional representation-encoder latents, while the failure is debated as either a capacity bottleneck or a geometric mismatch.

  • Method

    RJF combines Riemannian Flow Matching with Jacobi Regularization to align generative trajectories and optimization with the latent manifold’s topology and curvature.

  • Results

    3.37 FID is achieved with RJF on the standard 131M-parameter DiT-B with guidance, whereas standard flow matching fails to converge.

  • Takeaways & Limitations

    RJF enables standard DiT architectures to learn directly from representation-encoder features without computationally expensive width scaling.

Abstract

from arXiv · show

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

1 Introduction

Standard diffusion transformers struggle to converge directly on representation-encoder features because Euclidean flow matching conflicts with their hyperspherical geometry. RJF aligns trajectories and optimization with the manifold, enabling standard DiT architectures to generate effectively without width scaling.

  • 1 Introduction: Representation encoders provide semantically rich features, but standard diffusion recipes fail to converge effectively on these high-dimensional latents.Prior work attributes this failure to a capacity bottleneck and proposes transformer width scaling.
  • 1 Introduction: 4.95 FID is achieved without guidance using standard LightingDiT-B, compared with 15.83 for the VAE-based LightingDiT-B.The result demonstrates the benefit of respecting representation geometry without architectural width scaling.
  • 1 Introduction: Geometric Interference arises when Euclidean probability paths traverse the low-density interior of the hyperspherical feature manifold instead of following its surface.The objective also spends capacity minimizing radial errors even though representation semantics are encoded on the manifold.
  • 1 Introduction: RJF combines Riemannian Flow Matching, which corrects the trajectory, with Jacobi Regularization, which accounts for geodesic focusing and curvature.The resulting optimization is consistent with the topology and curvature of the latent space.
  • 1 Introduction: 3.37 FID is achieved by RJF with DINOv2-B and a 131M-parameter DiT-B with guidance, while standard flow matching fails to converge.Without guidance, the same setup achieves FID 4.95 in 200 epochs.

2 Geometrical Analysis

Representation-encoder features lie on a hypersphere with semantic information concentrated in angular directions, while standard Euclidean flow matching creates off-manifold trajectories and radial conflicts. Experiments show that masking radial loss removes the apparent capacity bottleneck, motivating manifold-constrained flow matching.

  • 2.1 The Geometry Gap: DINOv2-B features are projected onto a fixed-radius hypersphere, with near-zero radial variance and semantic information encoded in angular components.This differs from the diffuse-shell geometry assumed by a standard Gaussian prior.
  • 2.1 The Geometry Gap: The linear Euclidean path between noise and data forms a chord through the hypersphere’s interior rather than remaining on the manifold surface.The model must therefore learn a velocity field in off-manifold regions where the representation space is undefined.
  • 2.2 Revisiting the Capacity Hypothesis: Geometric Interference: The apparent width bottleneck reflects geometric conflict: Euclidean objectives force models to spend capacity minimizing radial errors while learning angular semantics.Narrow models fail under the full loss because radial noise dominates gradient updates.
  • 2.2 Revisiting the Capacity Hypothesis: Geometric Interference: Riemannian Flow Matching replaces the chord with a geodesic trajectory on the hypersphere, eliminating the radial conflict by aligning transport with the manifold.Projecting endpoints or masking radial loss alone does not fix the off-manifold Euclidean path.

3 Method

The method reformulates flow matching on the hyperspherical feature manifold, replacing Euclidean chords with geodesic paths and correcting curvature-dependent error propagation. This produces tangent-space velocity fields and manifold-preserving integration throughout generation.

  • Euclidean Flow Matching: Flow Matching learns a time-dependent vector field that transforms Gaussian noise into the data distribution through a probability path.Conditional Flow Matching trains the field against the velocity of a specified path.
  • Geometric Motivation: Standard Euclidean flow matching is ill-suited to hyperspherical features because linear interpolation cuts through the sphere’s low-density interior.This geometric mismatch forces the model to learn trajectories off the feature manifold.
  • Riemannian Flow Matching: RJF projects features and Gaussian noise onto the unit hypersphere, then uses SLERP to define geodesic probability paths that preserve unit norm.The geodesic path remains on the manifold for every interpolation time, avoiding norm collapse.
  • Riemannian Flow Matching: The Riemannian objective trains the network to predict tangent-space velocities, eliminating radial error by construction and focusing learning on angular changes.The resulting loss is the squared ambient-space norm against the Riemannian target velocity.
  • Sampling: Geodesic exponential-map integration wraps predicted tangent velocities around the sphere, maintaining the learned manifold geometry during sampling.A final rotate-and-normalize operation corrects minor numerical drift over multiple integration steps.
  • Jacobi Regularization: Jacobi regularization weights velocity errors by curvature-induced geodesic focusing, down-weighting errors near data and prioritizing precision near noise.The weighting factor is λ(t, Ω) = sinc^2((1-t)Ω).

4 Experiments

Experiments show that geometric corrections, rather than width scaling, let standard diffusion transformers model representation latents effectively. RJF improves convergence and FID across training settings.

  • Main Results: FID 3.62: LightingDiT-XL with RJF outperforms Euclidean Flow Matching at FID 4.28 in the 80-epoch setting.The comparison uses ImageNet 256×256 without guidance in the reported main-results setting.
  • Main Results: RJF produces high-quality, semantically diverse ImageNet 256×256 samples after 80 epochs.The qualitative results are uncurated samples from LightingDiT-XL trained with RJF.
  • Ablation Study: FID 24.32: standard Euclidean Flow Matching fails to converge effectively on DINOv2-B features.The baseline wastes capacity on radial errors and velocity fields inside the manifold.
  • Ablation Study: FID 21.99: projecting noise and target latents onto the sphere yields only a marginal 2.33-FID improvement.The Euclidean interpolation remains a chord through the manifold interior.
  • Ablation Study: FID 7.06: Riemannian Flow Matching substantially improves performance by following geodesic trajectories.This result supports correcting the trajectory mismatch rather than only constraining endpoints.
  • Ablation Study: FID 6.77: adding Jacobi Regularization further improves geometric alignment after Riemannian Flow Matching.With 200 epochs, RJF reaches FID 4.95 without guidance and FID 3.37 with guidance, with IS 180.26.

5 Discussion

RJF generalizes across architectures and representation encoders, while inference-time radius selection affects performance. These results support geometric alignment as a broadly applicable approach.

  • Generality Across Architectures: RJF achieves consistent performance gains across the evaluated Diffusion Transformer architectures.The evaluation includes LightingDiT-XL, DDT-XL, and an architecture explicitly designed with width scaling for RAE latents.
  • Generality Across Architectures: FID 3.62: RJF on LightingDiT-XL outperforms the VAE baseline at FID 4.29 and Euclidean Flow Matching at FID 4.28.On DDT-XL, RJF also reaches FID 5.82 versus 6.55 for the Euclidean baseline.
  • Inference Radius: FID 6.77: increasing the inference projection radius to R ≈45 outperforms reprojection to the original DINOv2-B norm, which yields FID 7.79.The result indicates sensitivity of the RAE decoder to the latent norm used at inference.

6 Related Works

Related work improves diffusion generation through representation alignment or direct feature-space modeling, but standard diffusion recipes can still fail on representation latents. RJF frames this failure as geometric rather than purely architectural.

  • Representation Alignment: Representation alignment methods such as REPA and DDT improve diffusion transformer training by aligning diffusion features with pretrained representations.DDT extends this paradigm with alignment applied to a decoupled encoder.
  • Representation Autoencoders: RAE replaces the VAE with representation encoders but uses width scaling to mitigate convergence failure in direct feature-space diffusion.The paper instead attributes the failure to geometric mismatch in Euclidean flow matching.
  • Geometric Perspective: RJF addresses the representation-latent problem by treating Euclidean trajectories through the hypersphere interior as a geometric mismatch.The supplied related-work passages connect this diagnosis to the limitations of direct diffusion on representation encoders.
  • Flow Matching: Standard Flow Matching constructs Euclidean linear probability paths, which are ill-suited to data residing on non-Euclidean spaces.This motivates methods that account for the geometry of the data manifold.

7 Conclusion

The paper concludes that convergence failure on representation-encoder latents is geometric rather than a capacity deficit. RJF enforces geodesic trajectories and achieves strong results without relying on width scaling.

  • Conclusion: RJF identifies Geometric Interference as Euclidean flow matching’s wasteful radial-error and off-manifold trajectory problem.The method combines Riemannian Flow Matching with Jacobi Regularization to align trajectories with latent topology.
  • Conclusion: FID 3.37: RJF enables the standard 131M-parameter DiT-B to converge where baselines fail.The conclusion also reports FID 3.62 for DiT-XL after 80 epochs.
  • Conclusion: Efficient generation requires respecting latent topology rather than simply scaling model width.This is the paper’s stated conclusion about the relationship between geometric alignment and architectural scaling.

8 Theoretical Derivation of Jacobi Field Regularization

The derivation shows that curvature makes tangent-space velocity errors distort endpoint errors, motivating Jacobi-based weighting in the Riemannian flow objective.

  • Geometric Motivation: Positive curvature makes tangent-to-endpoint mapping non-isometric, so uniform velocity-error weighting does not ensure accurate endpoint reconstruction.The objective therefore targets the error at the target noise manifold rather than only instantaneous velocity error.
  • Jacobi Fields and Error Propagation: The Jacobi-field formulation models trajectory perturbations caused by errors in the initial velocity.A geodesic family shares the same starting point, while its initial velocity is perturbed by w.
  • Hyperspherical Specialization: For a hypersphere, curvature K = 1/R^2 and geodesic length L = (1−t)ΩR determine the Jacobi-field magnitude along the trajectory.Substitution reduces the curvature-dependent expression to a normalized sinc function of the remaining angular distance.
  • Jacobian Scaling: The Jacobian scaling factor matches tangent-space error to its effective displacement at the target endpoint.It is defined as the ratio between endpoint displacement and the corresponding linearized displacement.
  • Training and Sampling: RJF training normalizes data and noise onto the sphere, uses SLERP geodesic interpolation, and samples shifted Logit-Normal times.The sampling procedure projects velocity onto the tangent space, updates position by rotation, renormalizes, and rescales the final output.
  • Jacobi-Regularized Objective: The regularized loss weights squared tangent-space error by λ(t, Ω) = sinc^2((1−t)Ω) to minimize squared endpoint error.The weight is the squared Jacobian factor induced by the manifold geometry.
Loading 2602.10099v2…