Generative Modeling on Manifolds Through Mixture of Riemannian Diffusion Processes

Learning the distribution of data on Riemannian manifolds is crucial for modeling data from non-Euclidean space, which is required by many applications from diverse scientific fields. Yet, existing generative models on manifolds suffer from expensive divergence computation or rely on approximations of heat kernel. These limitations restrict their applicability to simple geometries and hinder scalability to high dimensions. In this work, we introduce the Riemannian Diffusion Mixture, a principled framework for building a generative process on manifolds as a mixture of endpoint-conditioned diffusion processes instead of relying on the denoising approach of previous diffusion models, for which the generative process is characterized by its drift guiding toward the most probable endpoint with respect to the geometry of the manifold. We further propose a simple yet efficient training objective for learning the mixture process, that is readily applicable to general manifolds. Our method outperforms previous generative models on various manifolds while scaling to high dimensions and requires a dramatically reduced number of in-training simulation steps for general manifolds. Deep generative models have shown great success in learning the distributions of the data represented in Euclidean space, e.g., images and text. While the focus of the previous works has been biased toward data in the Euclidean space, modeling the distribution of the data that naturally resides in the Riemannian manifold with specific geometry has been underexplored, while they are required for wide application: For example, the earth and climate science data (Karpatne et al., 2018; Mathieu & Nickel, 2020) lives in the sphere, whereas the protein structures (Jumper et al., 2021; Watson et al., 2022) and the robotic movements (Simeonov et al., 2022) are best represented by the group SE(3). Moreover, 3D computer graphics shapes (Hoppe et al., 1992) can be identified as a general closed manifold. However, previous generative methods are ill-suited for modeling these data as they do not take into consideration the specific geometry describing the data space and may assign a non-zero probability to regions outside the desired space.