Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Riemannian Denoising Diffusion Probabilistic Models

Zichen Liu, Wei Zhang, Christof Schütte, Tiejun Li

TL;DR

We introduce Riemannian Denoising Diffusion Probabilistic Models (RDDPMs), extending DDPMs to learn distributions on level-set submanifolds $\mathcal{M}=\{x: \xi(x)=0\}$ via a projection-based forward process that yields explicit transition densities. The approach avoids geodesic computations and heat-kernel approximations, and its continuous-time limit connects to Riemannian score-based diffusion, establishing a principled link to manifold-level generative modeling. The method is demonstrated on diverse, high-dimensional manifolds including $\mathrm{SO}(10)$ and alanine dipeptide, achieving accurate density estimation and sample generation while preserving geometric constraints such as SE(3) invariance. These results broaden the applicability of diffusion models to constrained manifolds and complex geometric settings, enabling new applications in physics, chemistry, and geometry-aware learning.

Abstract

We propose Riemannian Denoising Diffusion Probabilistic Models (RDDPMs) for learning distributions on submanifolds of Euclidean space that are level sets of functions, including most of the manifolds relevant to applications. Existing methods for generative modeling on manifolds rely on substantial geometric information such as geodesic curves or eigenfunctions of the Laplace-Beltrami operator and, as a result, they are limited to manifolds where such information is available. In contrast, our method, built on a projection scheme, can be applied to more general manifolds, as it only requires being able to evaluate the value and the first order derivatives of the function that defines the submanifold. We provide a theoretical analysis of our method in the continuous-time limit, which elucidates the connection between our RDDPMs and score-based generative models on manifolds. The capability of our method is demonstrated on datasets from previous studies and on new datasets sampled from two high-dimensional manifolds, i.e. $\mathrm{SO}(10)$ and the configuration space of molecular system alanine dipeptide with fixed dihedral angle.

Riemannian Denoising Diffusion Probabilistic Models

TL;DR

We introduce Riemannian Denoising Diffusion Probabilistic Models (RDDPMs), extending DDPMs to learn distributions on level-set submanifolds via a projection-based forward process that yields explicit transition densities. The approach avoids geodesic computations and heat-kernel approximations, and its continuous-time limit connects to Riemannian score-based diffusion, establishing a principled link to manifold-level generative modeling. The method is demonstrated on diverse, high-dimensional manifolds including and alanine dipeptide, achieving accurate density estimation and sample generation while preserving geometric constraints such as SE(3) invariance. These results broaden the applicability of diffusion models to constrained manifolds and complex geometric settings, enabling new applications in physics, chemistry, and geometry-aware learning.

Abstract

We propose Riemannian Denoising Diffusion Probabilistic Models (RDDPMs) for learning distributions on submanifolds of Euclidean space that are level sets of functions, including most of the manifolds relevant to applications. Existing methods for generative modeling on manifolds rely on substantial geometric information such as geodesic curves or eigenfunctions of the Laplace-Beltrami operator and, as a result, they are limited to manifolds where such information is available. In contrast, our method, built on a projection scheme, can be applied to more general manifolds, as it only requires being able to evaluate the value and the first order derivatives of the function that defines the submanifold. We provide a theoretical analysis of our method in the continuous-time limit, which elucidates the connection between our RDDPMs and score-based generative models on manifolds. The capability of our method is demonstrated on datasets from previous studies and on new datasets sampled from two high-dimensional manifolds, i.e. and the configuration space of molecular system alanine dipeptide with fixed dihedral angle.
Paper Structure (43 sections, 5 theorems, 69 equations, 6 figures, 8 tables, 4 algorithms)

This paper contains 43 sections, 5 theorems, 69 equations, 6 figures, 8 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Riemannian submanifolds
  4. Denoising diffusion probabilistic models
  5. Method
  6. Projection scheme
  7. Forward process
  8. Reverse process
  9. Training objective
  10. Algorithmic details
  11. Generation of trajectory data
  12. Choice of $N$, $\sigma_k$, and $\beta_{k+1}$
  13. Method for solving constraint equations \ref{['projection-kstep-forward']} and \ref{['projection-kstep-backward']}
  14. Choices of $b$ and sampling of the prior $p(x^{(N)})$
  15. Theoretical results
  16. ...and 28 more sections

Key Result

Theorem 4.1

Let $T>0$ and $g: [0,T]\rightarrow \mathbb{R}^+$ be a continuous function. Define $h=\frac{T}{N}$ and $t_k=kh$, for $k=0,1,\dots,N-1$. Assume that $\sigma_k = \beta_{k+1} = \sqrt{h} g(t_k)$. Also assume that, for any parameter $\theta$, there is a $C^1$ function $s_\theta: \mathbb{R}^n\times [0,T]\r where $\mathbb{E}_{\mathbb{Q}}$ on the right hand side denotes the expectation with respect to the

Figures (6)

  • Figure 1: First row: datasets and true distributions. Second row: learned samples and distributions.
  • Figure 2: Results for $\mathrm{SO}(10)$ with $m=5$. Histograms of the statistics $\operatorname{tr}(S), \operatorname{tr}(S^2), \operatorname{tr}(S^4)$, and $\operatorname{tr}(S^5)$ for the forward process (solid line) and the learned reverse process (dashed line) at different steps $k = 0, 50, 200, 500$, colored in black, red, green, and blue, respectively.
  • Figure 3: (a) Illustration of the system. Names and $1$-based indices are shown for atoms that define the dihedral angles. The dihedral angles $\phi$ and $\psi$ are defined by atoms whose $1$-based indices are $5,7,9, 15$ and $7,9, 15, 17$, respectively. (b) Histograms of the angle $\psi$, $\mathrm{RMSD}_1$, and $\mathrm{RMSD}_2$ for the forward process (solid line) and the learned reverse process (dashed line) at steps $k = 0, 10, 40, 200$, colored in black, red, green, blue, respectively. The $\psi$ values of the two reference points that are used to define $\mathrm{RMSD}_1$ and $\mathrm{RMSD}_2$ are $-20^\circ$ and $150^\circ$, respectively (as shown by the two vertical dashed lines in the left panel).
  • Figure 4: The learned densities on earth and climate science datasets, with the standard dataset splitting. Darker green color indicates areas of higher likelihood. Red dots and blue dots show points in test set and generated samples, respectively.
  • Figure 5: Volcano dataset. (a) Red, blue, and green points represent the training set, validation set, and test set, respectively, obtained from the standard dataset splitting with random seed $4$. Triangles indicate the isolated points in the validation set and in the test set. (b) The training NLL, the validation NLL, and the test NLL during the training with random seed $4$. The solid lines correspond to the training where the standard dataset splitting is employed. The dashed lines correspond to the training where the isolated points (triangles in (a)) are included in the training set. (c) The best NLLs for five different runs with random seeds $0$, $1$, $2$, $3$, and $4$.
  • ...and 1 more figures

Theorems & Definitions (10)

  • Theorem 4.1
  • Corollary 4.2
  • proof : Proof of Theorem \ref{['thm-continuous-limit']}
  • proof : Proof of Corollary \ref{['corollary-continuous-limit-kl']}
  • Lemma 1.1
  • proof
  • Proposition 3.1
  • proof
  • Proposition 3.2
  • proof