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

Leveraging representation encoders for generative modelingoffers a path for efficient, high-fidelity synthesis. However, standard diffu-sion transformers fail to converge on these representations directly. Whilerecent work attributes this to a capacity bottleneck—proposing computa-tionally expensive “width scaling” of diffusion transformers—we demon-strate that the failure is fundamentally geometric. We identify GeometricInterference as the root cause: standard Euclidean flow matching forcesprobability paths through the low-density interior of the hypersphericalfeature space of representation encoders, rather than following the mani-fold surface. To resolve this, we propose Riemannian Flow Matchingwith Jacobi Regularization (RJF). By constraining the generativeprocess to the manifold geodesics and correcting for curvature-inducederror propagation, RJF enables standard Diffusion Transformer architec-tures to converge without width scaling. Our method RJF enables thestandard DiT-B architecture (131M parameters) to converge effectively,achieving an FID of 3.37 where prior methods fail to converge.