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

2026-02-10

AI Summary

Standard diffusion Transformers fail to converge when directly utilizing encoder-derived feature representations due to geometric interference, hindering high-fidelity generative modeling. This work introduces Riemannian geometry into diffusion models for the first time, proposing Riemannian flow matching and Jacobian regularization to constrain the generative process along geodesics of the representation manifold. This approach avoids probability paths traversing low-density regions in Euclidean space. Notably, it achieves stable training without increasing model width—using only a 131M-parameter DiT-B—and attains a state-of-the-art FID of 3.37 on image generation, significantly outperforming prior baselines that failed to converge.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Knowledge Representation and Reasoning: Geometric, Spatial, and Temporal Reasoning

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation

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

Problem

• Diffusion Transformers
• Representation Encoders
• Manifold Learning
• Geometric Interference
• Flow Matching

Innovation

• Riemannian Flow Matching
• Geometric Interference
• Manifold Learning
• Diffusion Transformers
• Jacobi Regularization