Generative Modeling on Manifolds Through Mixture of Riemannian Diffusion Processes

TL;DR

This work constructs a diffusion process using a mixture of bridge processes derived on general manifolds without requiring heat kernel estimations, and develops a geometric understanding of the mixture process, deriving the drift as a weighted mean of tangent directions to the data points that guides the process toward the data distribution.

Abstract

Learning the distribution of data on Riemannian manifolds is crucial for modeling data from non-Euclidean space, which is required by many applications in 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 diffusion process on manifolds. Instead of following the denoising approach of previous diffusion models, we construct a diffusion process using a mixture of bridge processes derived on general manifolds without requiring heat kernel estimations. We develop a geometric understanding of the mixture process, deriving the drift as a weighted mean of tangent directions to the data points that guides the process toward the data distribution. We further propose a scalable training objective for learning the mixture process that readily applies to general manifolds. Our method achieves superior performance on diverse manifolds with dramatically reduced number of in-training simulation steps for general manifolds.

Figures (14)

• Figure 1: We construct a generative process on general manifolds as a mixture of bridge processes (Eq. \ref{['eq:mixture_process']}). The drift of the mixture process (purple vector) corresponds to the weighted mean of the tangent vectors pointing to the directions of the endpoints (blue vector), guiding the diffusion process (black dotted) to the data distribution.
• Figure 2: (Left) Test NLL results on protein datasets. Best performance and its comparable results ($p>0.05$) from the t-test are highlighted in bold. (Right) Comparison on high-dimensional tori. We compare the log-likelihood in bits against RSGM and RFM where the results are obtained by running the open-source codes.
• Figure 2: (Left) Test NLL results on mesh datasets. We report the mean of 5 different runs. Best performance and its comparable results ($p>0.05$) from the t-test are highlighted. (Right) Comparison of the training time. We report the relative training time of the baselines with respect to ours on high-dimensional torus and Spot.
• Figure 3: Visualization of the generated samples and the learned density of our method and RFM on the mesh datasets. Blue dots represent the generated samples and darker red colors indicate higher likelihood. The numbers in the parentheses denote the number of in-training simulation steps used to train the model.
• Figure 4: Visualization of the generated samples and learned density of our model on earth and climate science datasets. Red dots denote samples from the test set and green dots denote the generated samples. Darker green colors denote a higher likelihood modeled by our approach.