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

TL;DR

This paper identifies geometric interference as the cause of convergence failure in diffusion transformers with representation encoders and proposes Riemannian Flow Matching with Jacobi Regularization (RJF) to enable effective training on the data manifold.

Contribution

The paper introduces RJF, a novel geometric regularization method that allows standard diffusion transformers to converge on representation manifolds without width scaling.

Findings

RJF enables the DiT-B architecture to converge with an FID of 3.37.

Standard Euclidean flow matching causes probability paths to pass through low-density regions.

RJF constrains the generative process to manifold geodesics, improving convergence.

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