Technical Summary: Learning on the Manifold with Riemannian Flow Matching

1. Problem Statement

Recent advancements in generative modeling have shifted toward leveraging frozen, high-dimensional semantic representation encoders (e.g., DiNOv2, SigLIP) to improve synthesis fidelity and convergence speed. However, standard Diffusion Transformers (DiT) and Flow Matching (FM) frameworks fail to converge effectively when trained directly on these representations.

Prior work, specifically Representation Autoencoders (RAE), attributed this failure to a capacity bottleneck, proposing that models require "width scaling" (increasing the transformer width to match the latent dimensionality) to resolve the issue. This paper challenges that hypothesis, arguing that the root cause is not insufficient model capacity but a fundamental geometric mismatch, termed Geometric Interference.

The core issue arises because representation encoders project features onto a hyperspherical manifold ($S^{d-1}$) with near-zero radial variance (a "hard shell" geometry). In contrast, standard Euclidean Flow Matching assumes a linear probability path (a chord) between the data and Gaussian noise. This linear interpolation cuts through the low-density interior of the hypersphere, forcing the model to learn velocity fields in undefined regions of the feature space. Consequently, the model wastes its capacity minimizing radial errors (which are structurally zero on the manifold) rather than learning the semantic angular transitions.

2. Methodology: Riemannian Flow Matching with Jacobi Regularization (RJF)

To resolve Geometric Interference without architectural modifications like width scaling, the authors propose RJF, a framework that aligns the generative process with the intrinsic geometry of the latent space.

2.1 Riemannian Flow Matching (Trajectory Correction)

Instead of linear interpolation in Euclidean space, RJF defines the probability path directly on the hyperspherical manifold using Spherical Linear Interpolation (SLERP).

• Geodesic Paths: The trajectory between a data point $x$ and noise $\epsilon$ (both projected onto the unit sphere) follows the geodesic (great circle arc) rather than a chord.
• Tangent Space Velocity: The target velocity field is computed by differentiating the SLERP path. This ensures the predicted velocity vector lies strictly in the tangent space of the manifold ($v_t \cdot x_t = 0$), eliminating the need to learn radial components.
• Sampling: During inference, the model uses a Geodesic (Exponential Map) integrator to update the latent state, ensuring the trajectory remains on the manifold surface throughout the generation process.

2.2 Jacobi Regularization (Curvature Correction)

The authors observe that simply correcting the trajectory is insufficient because the standard Mean Squared Error (MSE) loss treats velocity errors uniformly. On a positively curved hypersphere, geodesics focus (converge), meaning velocity errors propagate non-linearly.

• Jacobi Fields: To account for this, the authors introduce a weighting factor derived from Jacobi fields, which quantify the separation between geodesics caused by velocity perturbations.
• Loss Reweighting: The loss function is multiplied by a geometric weighting factor $\lambda(t, \Omega) = \text{sinc}^2((1-t)\Omega)$, where $\Omega$ is the geodesic distance. This down-weights errors near the data ($t=0$) where focusing mitigates perturbations and prioritizes precision near the noise ($t=1$), effectively annealing the learning signal to match the manifold's curvature.

3. Key Contributions

1. Geometric Analysis of Convergence Failure: The paper identifies Geometric Interference as the primary bottleneck for standard diffusion on representation encoders. It demonstrates that the failure stems from Euclidean objectives forcing models to learn through the off-manifold interior of the hypersphere, rather than a lack of model capacity.
2. RJF Framework: The proposal of a geometrically consistent training framework combining Riemannian Flow Matching (for trajectory alignment) and Jacobi Regularization (for curvature-aware error propagation).
3. Efficiency without Width Scaling: The demonstration that standard DiT architectures (e.g., DiT-B with 131M parameters) can converge effectively on high-dimensional semantic latents without the computationally expensive width scaling proposed by prior works.

4. Experimental Results

Experiments were conducted on ImageNet-1K at 256×256 resolution using various model scales (DiT-B, DiT-L, DiT-XL) and representation encoders (DiNOv2-B, SigLIP, MAE).

• Convergence and FID Performance:
  • DiT-B (131M): Standard Euclidean Flow Matching on DiNOv2-B features failed to converge (FID ~24.21). RJF achieved an FID of 6.77 in 80 epochs and 3.37 (with guidance) in 200 epochs.
  • DiT-XL (677M): RJF achieved an FID of 3.62 in 80 epochs without guidance, outperforming the Euclidean baseline (FID 4.28) and the VAE-based LightingDiT (FID 4.29).
  • Scaling: The method showed consistent improvements across model sizes, achieving state-of-the-art Inception Scores (IS 186.2) and Precision (0.82).
• Ablation Studies:
  • Radial vs. Angular: Masking the radial loss provided only marginal gains (FID 21.99), confirming that the linear path itself (the chord) is the primary issue.
  • Trajectory: Switching to Riemannian Flow Matching (SLERP) alone drastically improved performance (FID 7.06), validating the trajectory mismatch hypothesis.
  • Regularization: Adding Jacobi Regularization further refined the results (FID 6.77).
• Generalization: The method successfully converged on other encoders with hyperspherical constraints, including SigLIP (FID 10.39) and MAE (FID 19.82), where standard methods failed completely.

5. Significance and Claims

The paper claims that the difficulty in training diffusion models on representation encoders is fundamentally geometric, not architectural. By respecting the hyperspherical topology of the latent space, standard Diffusion Transformers can learn high-fidelity semantics efficiently.

The authors assert that RJF unlocks the potential of standard DiT architectures, allowing them to match or exceed the performance of much larger, width-scaled models or complex VAE-based pipelines. The work establishes that efficient generative modeling on representation manifolds requires geometric alignment (correct trajectories and curvature-aware losses) rather than brute-force capacity scaling. This approach enables high-fidelity synthesis with fewer parameters and faster convergence, as evidenced by the DiT-B model achieving an FID of 3.37 where prior methods failed to converge.