Riemannian Denoising Diffusion Probabilistic Models
Extracted source text · arXiv 2505.04338
\documentclass{article}
% Recommended, but optional, packages for figures and better typesetting:
\usepackage{microtype}
\usepackage{graphicx}
% \usepackage{subfigure}
\usepackage{booktabs} % for professional tables
% hyperref makes hyperlinks in the resulting PDF.
% If your build breaks (sometimes temporarily if a hyperlink spans a page)
% please comment out the following usepackage line and replace
% \usepackage{icml2025} with \usepackage[nohyperref]{icml2025} above.
\usepackage{hyperref}
% Attempt to make hyperref and algorithmic work together better:
\newcommand{\theHalgorithm}{\arabic{algorithm}}
% Use the following line for the initial blind version submitted for review:
% \usepackage{icml2025}
% If accepted, instead use the following line for the camera-ready submission:
\usepackage[accepted]{icml2025}
% For theorems and such
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{mathtools}
\usepackage{amsthm}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%% my settings %%%%%%%%%%%%%%%%
% \usepackage{caption}
\usepackage{subcaption}
% \newcommand{\MYCOMMENT}[1]{\footnotesize\textcolor{gray}{// #1}}
\newcommand{\MYCOMMENT}[1]{\textcolor{gray}{// #1}}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% if you use cleveref..
\usepackage[capitalize,noabbrev]{cleveref}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% THEOREMS
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\theoremstyle{plain}
\newtheorem{theorem}{Theorem}[section]
\newtheorem{proposition}[theorem]{Proposition}
\newtheorem{lemma}[theorem]{Lemma}
\newtheorem{corollary}[theorem]{Corollary}
\theoremstyle{definition}
\newtheorem{definition}[theorem]{Definition}
\newtheorem{assumption}[theorem]{Assumption}
\theoremstyle{remark}
\newtheorem{remark}[theorem]{Remark}
% Todonotes is useful during development; simply uncomment the next line
% and comment out the line below the next line to turn off comments
%\usepackage[disable,textsize=tiny]{todonotes}
\usepackage[textsize=tiny]{todonotes}
% The \icmltitle you define below is probably too long as a header.
% Therefore, a short form for the running title is supplied here:
\icmltitlerunning{Riemannian Denoising Diffusion Probabilistic Models}
\begin{document}
\twocolumn[
\icmltitle{Riemannian Denoising Diffusion Probabilistic Models}
% It is OKAY to include author information, even for blind
% submissions: the style file will automatically remove it for you
% unless you've provided the [accepted] option to the icml2025
% package.
% List of affiliations: The first argument should be a (short)
% identifier you will use later to specify author affiliations
% Academic affiliations should list Department, University, City, Region, Country
% Industry affiliations should list Company, City, Region, Country
% You can specify symbols, otherwise they are numbered in order.
% Ideally, you should not use this facility. Affiliations will be numbered
% in order of appearance and this is the preferred way.
\icmlsetsymbol{equal}{*}
\begin{icmlauthorlist}
\icmlauthor{Zichen Liu}{equal,cds}
\icmlauthor{Wei Zhang}{equal,zib}
\icmlauthor{Christof Schütte}{zib,mathFUB}
\icmlauthor{Tiejun Li}{cds,math,CMLR}
\end{icmlauthorlist}
\icmlaffiliation{cds}{Center for Data Science, Peking University}
\icmlaffiliation{zib}{Zuse Institute Berlin}
\icmlaffiliation{mathFUB}{Institute of Mathematics, Freie Universität Berlin}
\icmlaffiliation{math}{School of Mathematical Sciences, Peking University}
\icmlaffiliation{CMLR}{Center for Machine Learning Research, Peking University}
\icmlcorrespondingauthor{Tiejun Li}{
tieli@pku.edu.cn}
% \icmlcorrespondingauthor{Firstname2 Lastname2}{first2.last2@www.uk}
% You may provide any keywords that you
% find helpful for describing your paper; these are used to populate
% the "keywords" metadata in the PDF but will not be shown in the document
% \icmlkeywords{Machine Learning, ICML}
\vskip 0.3in
]
% this must go after the closing bracket ] following \twocolumn[ ...
% This command actually creates the footnote in the first column
% listing the affiliations and the copyright notice.
% The command takes one argument, which is text to display at the start of the footnote.
% The \icmlEqualContribution command is standard text for equal contribution.
% Remove it (just {}) if you do not need this facility.
% \printAffiliationsAndNotice{} % leave blank if no need to mention equal contribution
\printAffiliationsAndNotice{\icmlEqualContribution} % otherwise use the standard text.
\begin{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.
\end{abstract}
\section{Introduction}
\label{sec-intro}
Diffusion models are generative models that learn data distributions by gradually perturbing data into noise and then reconstructing data from noise using stochastic processes. Two primary frameworks of diffusion models are Denoising Diffusion Probabilistic Models (DDPMs; \citet{NEURIPS2020_ddpm}), where the models are trained to minimize the variational bound in variational inference, and Score-based Generative Models (SGMs; \citet{song2019generative,song2021scorebased}), where the models are trained to learn the score function~\citep{hyvarinen2005estimation}. Both frameworks have achieved remarkable success in various applied fields.
In recent years, there has been a growing interest in developing generative models for data on manifolds~\citep{de2022riemannian,NEURIPS2023_ScalingRDM,chen2024flow,jo2024generativemodelingmanifoldsmixture}.
However, existing methods on manifolds rely on substantial geometric information, e.g.\ geodesics~\citep{de2022riemannian}, heat kernel or its approximations~\citep{de2022riemannian,NEURIPS2023_ScalingRDM}, or eigenfunctions and (pre)metrics~\citep{chen2024flow}. As a result, their applications are restricted to manifolds where such information can be obtained.
In this work, we introduce Riemannian Denoising Diffusion Probabilistic Models (RDDPMs), an extension of DDPMs to Riemannian submanifolds.
A key ingredient is the projection scheme used in Monte Carlo methods for sampling under constraints~\citep{goodman-submanifold,LelievreStoltzZhang2022}, which allows us to develop Markov chains on manifolds with explicit transition densities. The main advantages of our method over existing methods are summarized below.
\begin{itemize}
\item Our method is developed for submanifolds that are level sets of smooth functions in Euclidean space. This general setting includes most of the often studied manifolds such as spheres and matrix groups. More importantly, it fits well with applications where constraints are involved, e.g.\ applications in molecular dynamics and statistical mechanics.
\item Our method only requires the computation of the value and the first order derivatives of the function that defines the submanifold. Therefore, it is applicable to more general manifolds.
\item We present a theoretical analysis for the loss function of our method in the continuous-time limit, elucidating its connection to the existing methods~\citep{de2022riemannian}. This analysis also shows the equivalence between loss functions derived from variational bound in variational inference and from learning score function.
\end{itemize}
We successfully apply our method to datasets from prior works, and to new datasets from the special orthogonal group $\mathrm{SO}(10)$ and from alanine dipeptide with fixed dihedral angle, both of which are challenging for existing methods due to their geometric complexity.
\section{Background}
\paragraph{Riemannian submanifolds}
We consider the zero level set $\mathcal{M}=\{x \in \mathbb{R}^n |\xi(x)=0\}$ of a smooth function $\xi:\mathbb{R}^n\rightarrow \mathbb{R}^{n-d}$.
We assume that $\mathcal{M}$ is non-empty and the matrix $\nabla\xi(x)\in \mathbb{R}^{n\times(n-d)}$, i.e.\ the Jacobian of $\xi$, has full rank at each $x\in \mathcal{M}$.
Under this assumption, $\mathcal{M}$ is a $d$-dimensional submanifold of $\mathbb{R}^n$. We further assume that $\mathcal{M}$ is a smooth compact connected manifold without boundary.
The Riemannian metric on $\mathcal{M}$ is endowed from the standard Euclidean distance on $\mathbb{R}^n$.
For $x\in \mathcal{M}$, we denote by $T_x\mathcal{M}$ the tangent space of
$\mathcal{M}$ at~$x$. The orthogonal
projection matrix $P(x) \in \mathbb{R}^{n\times n}$ mapping
$T_x\mathbb{R}^n=\mathbb{R}^n$ to $T_x\mathcal{M}$ is given by $P(x) = I_n -
\nabla\xi(x) \big(\nabla\xi(x)^\top\nabla\xi(x)\big)^{-1} \nabla\xi(x)^\top$. Let $U_x\in\mathbb{R}^{n\times{d}}$ be a matrix whose column vectors form an
orthonormal basis of $T_{x}\mathcal{M}$ such that $U_x^\top U_x=I_d$. It is straightforward to verify that $P(x)=U_xU_x^\top$.
The volume element over $\mathcal{M}$ is denoted by $\sigma_{\mathcal{M}}$. All probability densities that appear in this paper refer to relative probability densities with respect to either $\sigma_{\mathcal{M}}$ or the product of $\sigma_{\mathcal{M}}$ over product spaces.
For notational simplicity, we also use the shorthand $\int p(x^{(1:N)}) dx^{(1:N)} := \int_{\mathcal{M}}\dots \int_{\mathcal{M}} p(x^{(1)},\dots, x^{(N)}) d\sigma_{\mathcal{M}}(x^{(1)})\cdots d\sigma_{\mathcal{M}}(x^{(N)})$.
% $\int p(x^{(1:N)})\, dx^{(1:N)} := \int_{\mathcal{M}}\dots \int_{\mathcal{M}} p(x^{(1)}, x^{(2)},\dots, x^{(N)})\, d\sigma_{\mathcal{M}}(x^{(1)})\, d\sigma_{\mathcal{M}}(x^{(2)})\cdots d\sigma_{\mathcal{M}}(x^{(N)})$.
\paragraph{Denoising diffusion probabilistic models}
We formulate the DDPMs \citep{pmlr-v37-sohl-dickstein15,NEURIPS2020_ddpm} to the general Riemannian manifold setting.
Assume that the data distribution is $q_0(x) d\sigma_{\mathcal{M}}(x)$.
DDPMs are a class of generative models built on Markov chains.
Specifically, states $x^{(1)}, \dots, x^{(N)}\in \mathcal{M}$ are generated by
gradually corrupting the data $x^{(0)}$ according to a Markov chain on
$\mathcal{M}$, i.e. the forward process. The joint probability
density of $x^{(1)}, \dots, x^{(N)}$ given $x^{(0)}$ is
\begin{equation}
q(x^{(1:N)}\,|\,x^{(0)}) = \prod_{k=0}^{N-1} q(x^{(k+1)}\,|\,x^{(k)})\,.
\label{joint-forward}
\end{equation}
The generative process, also called the reverse process, is a Markov chain on $\mathcal{M}$ that is learnt to reproduce the data by reversing the forward process.
Its joint probability density is
\begin{equation}
p_\theta(x^{(0:N)}) = p(x^{(N)}) \prod_{k=0}^{N-1} p_\theta(x^{(k)}\,|\,x^{(k+1)})\,,
\label{joint-backward}
\end{equation}
where $p(x^{(N)})$ is a (fixed) prior density. The probability density of
$x^{(0)}$ generated by the reverse process is therefore $p_\theta(x^{(0)})=
\int p_\theta(x^{(0:N)})\, dx^{(1:N)}$.
The learning objective is based on the standard variational bound on the negative log-likelihood. Specifically,
using \eqref{joint-forward} and \eqref{joint-backward}, and applying Jensen's inequality, we can derive
\begin{equation}
\begin{aligned}
&\mathbb{E}_{q_0}\big(-\log p_\theta(x^{(0)})\big) \\
= & \mathbb{E}_{q_0}\Big(-\log \int p_\theta(x^{(0:N)})\, dx^{(1:N)}\Big) \\
=& \mathbb{E}_{q_0}\Big(-\log \int
\frac{p_\theta(x^{(0:N)})}{q(x^{(1:N)}\,|\,x^{(0)})}
q(x^{(1:N)}\,|\,x^{(0)})\, dx^{(1:N)}\Big) \\
\le & \mathbb{E}_{\mathbb{Q}^{(N)}}\Big( -\log \frac{p_\theta(x^{(0:N)})}{q(x^{(1:N)}\,|\,x^{(0)})}\Big)\\
=& \mathbb{E}_{\mathbb{Q}^{(N)}}\Big( -\log p(x^{(N)}) - \sum_{k=0}^{N-1}
\log\frac{p_\theta(x^{(k)}\,|\,x^{(k+1)})}{q(x^{(k+1)}\,|\,x^{(k)})}\Big)\,,
\end{aligned}
\label{variational-bound}
\end{equation}
where $\mathbb{E}_{q_0}, \mathbb{E}_{\mathbb{Q}^{(N)}}$ denote the expectation with respect to the data distribution on $\mathcal{M}$, and the expectation with respect to
the joint density $q(x^{(0:N)})$ under the forward process, respectively.
In order to derive an explicit training objective, we have to construct Markov chains on $\mathcal{M}$ with explicit transition densities.
We discuss how this can be achieved in the next section.
We conclude this section by reformulating the variational bound (\ref{variational-bound})
using relative entropy (see~\citet{NEURIPS2021_likelihood_training_sbdm} for a similar formulation of score-based diffusion models). Recall that the relative entropy, or Kullback-Leibler (KL) divergence, from a probability density $Q_2$ to another probability density $Q_1$ in the same measure space, where $Q_1$ is absolutely continuous with respect to $Q_2$, is defined as
$H(Q_1\,|\,Q_2) := \mathbb{E}_{Q_1}\Big(\log \frac{Q_1}{Q_2}\Big)$.
For simplicity, we also use the same notation for two probability measures.
Adding the term $\mathbb{E}_{q_0}(\log q_0)$ to both sides of the inequality (\ref{variational-bound}), we see that it is equivalent to (the data processing inequality)
\begin{equation}
H(q_0\,|\,p_\theta) \le
H(\overleftarrow{\mathbb{Q}}^{(N)}\,|\,\mathbb{P}^{(N)}_\theta) \,,
\label{relative-entropy-formulation-of-variational-bound}
\end{equation}
where the upper bound is the relative entropy from the path measure
$\mathbb{P}^{(N)}_\theta$ of the reverse process to the path measure $\overleftarrow{\mathbb{Q}}^{(N)}$
of the forward process (we include the arrow in the notation to indicate that paths of the forward process are viewed backwardly).
Therefore, learning DDPMs using the variational bound (\ref{variational-bound}) can be viewed as approximating probability measures in path space by the cross-entropy method~\citep{zwhws2014}.
\section{Method}
\label{sec-method}
\subsection{Projection scheme}
\label{subsec-projection-scheme}
We recall a projection scheme from
Monte Carlo sampling methods on manifolds~\citep{projection_diffusion,goodman-submanifold,LelievreStoltzZhang2022}, and we show that it allows us to construct Markov chains on $\mathcal{M}$ with tractable transition densities.
Given $x\in\mathcal{M}$ and a tangent vector $v\in T_x\mathcal{M}$ that is drawn from the standard Gaussian distribution on $T_x\mathcal{M}$, we compute
the intermediate state $x'=x+ \sigma^2 b(x)+\sigma v\in \mathbb{R}^n$, where $\sigma>0$ is a
positive constant and $b: \mathbb{R}^n \rightarrow \mathbb{R}^n$ is a smooth function.
In general, $x'$ does not belong to $\mathcal{M}$. We
consider the projection $y\in \mathcal{M}$ of $x'$ onto $\mathcal{M}$ along
an orthogonal direction in the column space of $\nabla\xi(x)$. Precisely, the projected state
$y$ is found by (numerically) solving the constraint equation for $c\in \mathbb{R}^{n-d}$
\begin{equation}
y = x + \sigma^2 b(x) + \sigma v + \nabla\xi(x)c, ~~ \mbox{s.t.}~~ \xi(y) = 0\in \mathbb{R}^{n-d} \,.
\label{projection}
\end{equation}
The choice of $b$ will affect the final invariant distribution and the convergence rate to equilibrium of the resulting Markov chain (see Section~\ref{subsec-algo-details} for further discussion). There are $n-d$ constraints in Eq.~\eqref{projection} with the same number of unknown variables. In particular, when $\xi$ is scalar-valued, i.e.\ $n-d=1$, solving Eq.~\eqref{projection} amounts to finding a root of a (nonlinear) scalar function.
While for some vectors $v$ there may be multiple solutions to Eq.~\eqref{projection}, the projected state, and hence the resulting Markov chain built on a numerical solver, is uniquely defined as long as the numerical solver finds one solution in a deterministic way. For example, this is the case when Newton's method is adopted with fixed initial condition $c=0$. To handle the case where no solution can be found for some $v$, let $\mathcal{F}_{x}^{(\sigma)}$ be the set of $v$ for which a solution can be found and denote by $\epsilon_{x}^{(\sigma)}=\mathbb{P}(v\notin \mathcal{F}_{x}^{(\sigma)})$, i.e.\ the probability that no solution can be found. Since $\epsilon_{x}^{(\sigma)}=0$ when $\sigma=0$ ($c=0$ is a solution for any $v$), it is natural to expect that $\epsilon_{x}^{(\sigma)}\rightarrow 0$ as $\sigma$ decreases to zero. That is, when $\sigma$ is small, a solution to \eqref{projection} can be found with probability close to one. However, we emphasize that we do not require this as an assumption in our analysis. We denote $\mathcal{M}_{x}^{(\sigma)}$ the set of all states in $\mathcal{M}$ that can be reached from $x$ by solving Eq.~\eqref{projection} with certain~$v\in \mathcal{F}_{x}^{(\sigma)}$.
To derive the transition density of $y$ from $x$, we notice that, by applying the orthogonal projection matrix $P(x)$ to both sides of Eq.~\eqref{projection} and using the fact that $P(x)\nabla\xi(x)=0$,
we have the relation $\sigma v=P(x)(y-x-\sigma^2 b(x))$. This indicates that, given a state $x\in\mathcal{M}$ and $y\in \mathcal{M}_{x}^{(\sigma)}$,
there is a unique tangent vector $v\in \mathcal{F}_{x}^{(\sigma)}\subseteq T_x\mathcal{M}$ that leads to $y$
by solving Eq.~\eqref{projection}. In other words, the mapping
from $v\in \mathcal{F}_{x}^{(\sigma)}$ to $y\in \mathcal{M}_{x}^{(\sigma)}$ is a bijection. Moreover, its inverse is explicitly given by $G_{x}^{(\sigma)}: \mathcal{M}_{x}^{(\sigma)}\rightarrow \mathcal{F}_{x}^{(\sigma)} \subseteq T_{x}\mathcal{M}$, where
\begin{equation}\label{general_g-x}
G_{x}^{(\sigma)}(y;b) = \frac{1}{\sigma} P(x)(y-x - \sigma^2 b(x))\,.
\end{equation}
To simplify the notation, we also write
$G_{x}^{(\sigma)}(y)$ when omitting the dependence on $b$ does not cause ambiguity.
Recall that $U_x, U_y\in \mathbb{R}^{n\times d}$ denote the matrices whose columns form the
orthonormal basis of $T_x\mathcal{M}$ and $T_y\mathcal{M}$, respectively.
Using Eq.~\eqref{general_g-x}, we can derive (see~\citet{LelievreStoltzZhang2022} for more detailed discussions)
\begin{equation}
\det (D G_{x}^{(\sigma)}(y)) = \sigma^{-d} \det (U_x^\top U_y)\,,
\end{equation}
where the left hand side denotes the determinant of the Jacobian $D G_{x}^{(\sigma)}(y): T_{y}\mathcal{M}\rightarrow T_vT_{x}\mathcal{M}\cong \mathbb{R}^d$ of the map $G_{x}^{(\sigma)}$ at $y$. Since $v$ is a Gaussian variable in $\mathcal{F}_{x}^{(\sigma)}$ (with a normalizing constant rescaled by $(1-\epsilon_{x}^{(\sigma)})^{-1}$), applying the change of variables formula for probability densities, we obtain the probability density of landing at $y \in \mathcal{M}_{x}^{(\sigma)}$ from~$x$:
% \small{
% \begin{equation}\label{transition-density-single-step}
% \begin{aligned}
% q(y\,|\,x) =& (2\pi)^{-\frac{d}{2}} (1-\epsilon^{(\sigma)}_{x})^{-1}\mathrm{e}^{-\frac{1}{2} |G_{x}^{(\sigma)}(y)|^2} |\det D G_{x}^{(\sigma)}(y)|\\
% =& (2\pi\sigma^2)^{-\frac{d}{2}} (1-\epsilon^{(\sigma)}_{x})^{-1} \mathrm{e}^{-\frac{1}{2} |G_{x}^{(\sigma)}(y)|^2}
% |\det (U_x^\top U_y)|\,.
% \end{aligned}
% \end{equation}
% }
% \normalsize
\begin{equation}\label{transition-density-single-step}
\begin{aligned}
&q(y\,|\,x) \\
=& (2\pi)^{-\frac{d}{2}} (1-\epsilon^{(\sigma)}_{x})^{-1}\mathrm{e}^{-\frac{1}{2} |G_{x}^{(\sigma)}(y)|^2} |\det D G_{x}^{(\sigma)}(y)|\\
=& (2\pi\sigma^2)^{-\frac{d}{2}} (1-\epsilon^{(\sigma)}_{x})^{-1} \mathrm{e}^{-\frac{1}{2} |G_{x}^{(\sigma)}(y)|^2}
|\det (U_x^\top U_y)|\,.
\end{aligned}
\end{equation}
\subsection{Forward process}
\label{subsec-forward-process}
We construct the forward process in our model as a Markov chain on $\mathcal{M}$ whose transitions are defined by the projection scheme in Eq.~\eqref{projection}.
Specifically, given the current state $x^{(k)}\in \mathcal{M}$ at step $k$, where $k=0,1,\dots, N-1$, the next state $x^{(k+1)}\in \mathcal{M}$ is determined by solving the constraint equation (for $c\in \mathbb{R}^{n-d}$):
\begin{equation} \label{projection-kstep-forward}
\begin{gathered}
x^{(k+1)} = x^{(k)} + \sigma_k^2 b(x^{(k)}) + \sigma_k v^{(k)} + \nabla\xi(x^{(k)})c, \\
\mbox{s.t.}~~ \xi(x^{(k+1)}) = 0\in\mathbb{R}^{n-d}\,,
\end{gathered}
\end{equation}
where $\sigma_k>0$ and $v^{(k)}\in T_{x^{(k)}}\mathcal{M}$ is a standard Gaussian variable in $T_{x^{(k)}}\mathcal{M}$. According to Eq.~\eqref{joint-forward} and Eq.~\eqref{transition-density-single-step}, we obtain the joint probability density of the forward process as
\begin{equation}\label{joint-forward-projection}
\begin{gathered}
q(x^{(1:N)}\,|\,x^{(0)}) = \prod_{k=0}^{N-1} q(x^{(k+1)}\,|\,x^{(k)})\,,\\
\mbox{where~} q(x^{(k+1)}\,|\,x^{(k)}) = (2\pi\sigma_k^2)^{-\frac{d}{2}}|\det (U_{x^{(k)}}^\top U_{x^{(k+1)}})| \\
\big(1-\epsilon_{x^{(k)}}^{(\sigma_k)}\big)^{-1} \exp\left(-\frac{1}{2}\left|G_{x^{(k)}}^{(\sigma_k)}(x^{(k+1)};b)\right|^2\right)\,,
\end{gathered}
\end{equation}
where
the definition of $G_{x^{(k)}}^{(\sigma_k)}(x^{(k+1)};b)$ can be found in Eq.~\eqref{general_g-x}.
\subsection{Reverse process}
The reverse process in our model is a Markov chain on $\mathcal{M}$
whose transitions (from $x^{(k+1)}$ to $x^{(k)}$) are defined by the constraint equation
\begin{equation}\label{projection-kstep-backward}
\begin{gathered}
x^{(k)} = x^{(k+1)} - \beta_{k+1}^2 b(x^{(k+1)}) + \beta_{k+1}^2 s^{(k+1),\theta}(x^{(k+1)}) \\
+ \beta_{k+1} \bar{v}^{(k+1)} + \nabla\xi(x^{(k+1)})c, ~~\mbox{s.t.}~~\xi(x^{(k)}) = 0 \,,
\end{gathered}
\end{equation}
for $k=N-1, N-2, \dots, 0$,
where $\beta_{k+1}>0$, $\bar{v}^{(k+1)}$ is a standard Gaussian variable in
$T_{x^{(k+1)}}\mathcal{M}$, and $s^{(k+1),\theta}(x^{(k+1)})\in\mathbb{R}^n$ depends
on the learning parameter $\theta$. Combining \eqref{joint-backward} and \eqref{transition-density-single-step}, we obtain the joint probability density of the reverse process as
% \begin{equation}\label{joint-backward-projection}
% \begin{gathered}
% p_\theta(x^{(0:N)}) = p(x^{(N)}) \prod_{k=0}^{N-1}
% p_\theta(x^{(k)}\,|\,x^{(k+1)})\,, \\
% \mbox{where~} p_\theta(x^{(k)}\,|\,x^{(k+1)}) = (2\pi\beta_{k+1}^2)^{-\frac{d}{2}} |\det (U_{x^{(k+1)}}^\top U_{x^{(k)}})| \\
% \left(1-\epsilon_{x^{(k+1)},\theta}^{(\beta_{k+1})}\right)^{-1} e^{-\frac{1}{2}\left|G_{x^{(k+1)}}^{(\beta_{k+1})}(x^{(k)}; s^{(k+1),\theta}-b)\right|^2}\,.
% \end{gathered}
% \end{equation}
\begin{equation}\label{joint-backward-projection}
\begin{gathered}
p_\theta(x^{(0:N)}) = p(x^{(N)}) \prod_{k=0}^{N-1}
p_\theta(x^{(k)}\,|\,x^{(k+1)})\,, \\
\mbox{where~} p_\theta(x^{(k)}\,|\,x^{(k+1)}) = (2\pi\beta_{k+1}^2)^{-\frac{d}{2}} \Big(1-\epsilon_{x^{(k+1)},\theta}^{(\beta_{k+1})}\Big)^{-1} \\
\cdot|\det (U_{x^{(k+1)}}^\top U_{x^{(k)}})| e^{-\frac{1}{2}\left|G_{x^{(k+1)}}^{(\beta_{k+1})}(x^{(k)}; s^{(k+1),\theta}-b)\right|^2}\,.
\end{gathered}
\end{equation}
In the above, $\epsilon_{x^{(k+1)},\theta}^{(\beta_{k+1})}$ denotes the probability of having $\bar{v}^{(k+1)}$ with which no solution to (\ref{projection-kstep-backward}) can be found and $G_{x^{(k+1)}}^{(\beta_{k+1})}(x^{(k)}; s^{(k+1),\theta}-b)$ is defined in Eq.~\eqref{general_g-x}.
\subsection{Training objective}
\label{subsec-training-objective}
The training objective follows directly from the variational bound~(\ref{variational-bound}) on
the negative log-likelihood, as well as the explicit expressions of transition densities in Eq.~\eqref{joint-forward-projection} and Eq.~\eqref{joint-backward-projection}.
Concretely, substituting Eq.~\eqref{joint-forward-projection} and Eq.~\eqref{joint-backward-projection} into the last line of (\ref{variational-bound}), we get
\begin{equation}
\mathbb{E}_{q_0}(-\log p_\theta(x^{(0)})) \le \mbox{Loss}^{(N)}(\theta) + C^{(N)} \,,
\label{variational-bound-involving-loss}
\end{equation}
where
\begin{equation}
\mbox{Loss}^{(N)}(\theta)=\frac{1}{2}\mathbb{E}_{\mathbb{Q}^{(N)}}\sum_{k=0}^{N-1} \left|G_{x^{(k+1)}}^{(\beta_{k+1})}(x^{(k)}; s^{(k+1),\theta}-b)\right|^2
\label{training-loss}
\end{equation}
is our objective for training the parameter $\theta$ in the reverse process
(recall that $\mathbb{E}_{\mathbb{Q}^{(N)}}$ denotes the expectation with respect to the forward process), the constant
\begin{equation}\label{constant-c}
\begin{aligned}
&C^{(N)} \\
=&-\mathbb{E}_{\mathbb{Q}^{(N)}}\bigg[\frac{1}{2}\sum_{k=0}^{N-1}\left|G_{x^{(k)}}^{(\sigma_k)}(x^{(k+1)};b)\right|^2+\log p(x^{(N)})\bigg] \\
&+d \sum_{k=0}^{N-1} \log \frac{\beta_{k+1}}{\sigma_k} - \mathbb{E}_{\mathbb{Q}^{(N)}} \sum_{k=0}^{N-1}\log \big(1-\epsilon_{x^{(k)}}^{(\sigma_k)}\big)
\end{aligned}
\end{equation}
is independent of the training parameter $\theta$, and we have used $\log\left(1-\epsilon_{x^{(k+1)},\theta}^{(\beta_{k+1})}\right)\le 0$ in deriving Eq.~\eqref{variational-bound-involving-loss}.
\subsection{Algorithmic details}
\label{subsec-algo-details}
The algorithms for training and data generation are summarized in Algorithms~\ref{algo_training} and~\ref{algo_generation}, respectively. Algorithm for sampling the forward process, which is similar to Algorithm~\ref{algo_generation}, and algorithm for solving constraint equations are presented in Algorithms~\ref{algo_forward} and \ref{Newton_solver} in Appendix~\ref{Exp_details}.
In the following, we discuss several algorithmic details of our method. Further details are left in Appendix~\ref{Exp_details}.
\paragraph{Generation of trajectory data} The optimal parameter $\theta$ is sought by minimizing the objective in Eq.~\eqref{training-loss}, for which trajectory data of the forward process is required as training data.
We sample trajectories of the
forward process in a preparatory step before training, and train the model with min-batches sampled from this trajectory dataset. The trajectory dataset is updated by re-sampling trajectories every several training epochs (see line 2 and lines 10--13 in Algorithm~\ref{algo_training}).
\paragraph{Choice of $N$, $\sigma_k$, and $\beta_{k+1}$}
The total number of steps $N$ should be large enough such that the forward Markov chain is able to reach equilibrium (approximately) starting from the data distribution.
While larger $\sigma_k$, $\beta_{k+1}$ allow the Markov chains to make larger jumps, their size should be chosen properly (depending on the manifold) such that the solution to the constraint equations \eqref{projection-kstep-forward} and \eqref{projection-kstep-backward} can be found with high probability.
\paragraph{Method for solving constraint equations~\eqref{projection-kstep-forward} and \eqref{projection-kstep-backward}}
As in Monte Carlo sampling methods on
submanifolds~\citep{projection_diffusion,goodman-submanifold,LelievreStoltzZhang2022},
we employ Newton's method to solve the constraint equations~\eqref{projection-kstep-forward} and \eqref{projection-kstep-backward}. This method has quadratic convergence (locally) and its implementation is simple.
In most cases, a solution with high precision can be found within a few iteration steps (e.g.\ less than 5 steps). When no solution is found, one can re-generate the state by sampling a new tangent vector or re-generate the entire trajectory. Our implementation of Newton's method is summarized in Algorithm~\ref{Newton_solver}.
\paragraph{Choices of $b$ and sampling of the prior $p(x^{(N)})$}
For relatively simple manifolds, we can for simplicity set $b=0$ and choose the prior (see line 2 in Algorithm~\ref{algo_generation}) as the uniform distribution on $\mathcal{M}$. When $\mathcal{M}$ is non-compact or when the convergence of Markov chain to equilibrium is slow with $b=0$, we can choose non-zero $b$ such as $b=-\nabla V$, i.e. the (full space) gradient of a function $V:\mathbb{R}^n \rightarrow \mathbb{R}$ in the ambient space. In this case, sampling the prior can be done by simulating a single long trajectory of the forward process. See Section~\ref{subsec-ad} and Appendix~\ref{subsec-dipeptide_details} for a concrete choice of nonzero $b$ in practice.
\begin{algorithm}[t!]
\caption{Training procedure}\label{algo_training}
\begin{algorithmic}[1]
\STATE \textbf{Input: } training data $(y^{i})_{1\le i\le M}$, functions $(s^{(k+1),\theta}(x))_{0\le k\le N-1}$, constants $\sigma_{k}, \beta_{k}>0$, function $b: \mathbb{R}^n\rightarrow\mathbb{R}^n$, integer $N$, batch size $B>0$, number of total training epochs $N_{\mathrm{epoch}}$, integer $l_{\mathrm{f}}>0$, learning rate $r>0$.
\STATE generate a path $(x^{(0),i},x^{(1),i},\dots,x^{(N),i})$ using Algorithm \ref{algo_forward}, for each $x^{(0),i}=y^{i}$
\FOR{$l = 1$ to $N_{\mathrm{epoch}}$}
\FOR{$j = 1$ to $\lfloor M/B \rfloor $}
\STATE sample a min-batch $\mathcal{I}=(i_1,i_2,\dots,i_{B})$ from the set of indices $\{1,2,\dots, M\}$
\STATE calculate loss:
\STATE \hspace{-0.7cm} $\ell(\theta)=\frac{1}{2|\mathcal{I}|}\sum\limits_{i\in\mathcal{I}}\sum\limits_{k=0}^{N-1} \left|G_{x^{(k+1),i}}^{(\beta_{k+1})}(x^{(k),i}; s^{(k+1),\theta}-b)\right|^2$
\STATE $\theta=\operatorname{optimizer\_update}(\theta, \ell(\theta), r)$
\ENDFOR
\STATE \MYCOMMENT{update trajectories every $l_{\mathrm{f}}$ epochs}
\IF{$l\, \%\, l_{\mathrm{f}} == 0$}
\STATE re-generate paths $(x^{(0),i},x^{(1),i},\dots,x^{(N),i})$ using Algorithm \ref{algo_forward}, for each $x^{(0),i}=y^{i}$
\ENDIF
\ENDFOR
\STATE \textbf{Return} $\theta$
\end{algorithmic}
\end{algorithm}
\begin{figure*}[!t]
\centering
\includegraphics[width=0.9\linewidth]{figs/mesh_summary_all.pdf}
\caption{First row: datasets and true distributions. Second row: learned samples and distributions.}\label{fig:mesh_result}
\end{figure*}
\section{Theoretical results}
\label{sec-continuous-limit}
In this section, we study the continuous-time limit of our proposed method.
Let $T>0$ and $g: [0,T]\rightarrow \mathbb{R}^+$ be a continuous function. Define $h=\frac{T}{N}$ and consider the case where $\sigma_k = \sqrt{h} g(kh)$, for $k=0,1,\dots,N-1$.
It is shown in~\citet{projection_diffusion} that the forward process (\ref{projection-kstep-forward}) converges strongly to the SDE on $\mathcal{M}$
\begin{equation}\label{sde-manifold}
dX_t = g^2(t) P(X_t) b(X_t) dt + g(t) dW^\mathcal{M}_t\,,~t\in[0,T]\,,
\end{equation}
where $W^\mathcal{M}_t$ is a Brownian motion over $\mathcal{M}$. Denote by
$p(\cdot,t)$ the probability density of $X_t$
with respect to $\sigma_{\mathcal{M}}$ at time $t\in [0,T]$.
We have the following result, which characterizes the loss function in Eq.~\eqref{training-loss} as $N\rightarrow +\infty$.
\begin{theorem}
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]\rightarrow
\mathbb{R}^n$ such that $s^{(k+1),\theta}(x) = s_\theta(x,t_{k+1})$ for all $k=0,1,\dots,N-1$ and $x\in\mathcal{M}$.
For the loss function defined in Eq.~\eqref{training-loss}, we have
\begin{equation*}
\begin{aligned}
& \lim_{N\rightarrow +\infty}\left(\mathrm{Loss}^{(N)}(\theta) -
\frac{Nd}{2}\right) = \\
& \mathbb{E}_{\mathbb{Q}}\Big[\frac{1}{2}\int_{0}^T \big| P(X_t)s_{\theta}(X_t, t)-\nabla_{\mathcal{M}} \log p(X_t,t)\big|^2 g^2(t)\,dt \\
& + \int_{0}^T \Big(P(X_t)b(X_t)-\frac{1}{2}\nabla_{\mathcal{M}} \log p(X_t,t)\Big) \\
& \hspace{3.0cm} \cdot \nabla_{\mathcal{M}} \log p(X_t,t)\, g^2(t)\, dt \Big]\, ,
\end{aligned}
\end{equation*}
where $\mathbb{E}_{\mathbb{Q}}$ on the right hand side denotes the expectation with respect to the paths of SDE~(\ref{sde-manifold}) and $\nabla_{\mathcal{M}}$ denotes the gradient operator on $\mathcal{M}$.
\label{thm-continuous-limit}
\end{theorem}
Based on Theorem~\ref{thm-continuous-limit}, the variational bound~(\ref{variational-bound}), and its relative entropy formulation in (\ref{relative-entropy-formulation-of-variational-bound}), we obtain the following corollary, which elucidates the connection between our RDDPMs and Riemannian score-based generative models \citep{de2022riemannian} as $N\rightarrow +\infty$.
\begin{corollary}
Under the same assumptions of Theorem~\ref{thm-continuous-limit}, we have, for any parameter $\theta$,
\begin{equation}
\begin{aligned}
& \lim_{N\rightarrow +\infty} H\big(\overleftarrow{\mathbb{Q}}^{(N)}\,|\,\mathbb{P}^{(N)}_\theta\big) \\
=& \frac{1}{2}\mathbb{E}_{\mathbb{Q}}\int_{0}^T \big| P(X_t)s_{\theta}(X_t, t)-\nabla_{\mathcal{M}} \log p(X_t,t)\big|^2 g^2(t) dt \\
=& H(\overleftarrow{\mathbb{Q}}\,|\,\mathbb{P}_\theta)\,,
\end{aligned}
\end{equation}
where $\overleftarrow{\mathbb{Q}}$ denotes the path measure of the time-reversal of SDE~(\ref{sde-manifold}), and $\mathbb{P}_\theta$ denotes the path measure of the SDE
\begin{equation}\label{sde-manifold-reversed-theta}
\begin{gathered}
dY_t = g^2(T-t) P(Y_t) \big(-b(Y_t) + s_\theta(Y_t, T-t)\big) dt \\
+ g(T-t) dW^\mathcal{M}_t\,, ~~ t\in[0,T]\,,
\end{gathered}
\end{equation}
starting from $Y_0=X_T$.
\label{corollary-continuous-limit-kl}
\end{corollary}
The proofs of Theorem~\ref{thm-continuous-limit} and
Corollary~\ref{corollary-continuous-limit-kl} are presented in Appendix~\ref{app-sec-proofs}.
\begin{algorithm}[htbp]
\caption{Sampling trajectory of reverse process}\label{algo_generation}
\begin{algorithmic}[1]
\STATE \textbf{Input: } trained functions $(s^{(k+1),\theta}(x))_{0\le k\le N-1}$, constants $\beta_{k}$, function $b: \mathbb{R}^n\rightarrow\mathbb{R}^n$, and integer $N$
\STATE draw sample $x^{(N)}$ from the prior distribution $p(x^{(N)})$
\FOR {$k = N-1$ to $0$}
\STATE sample $\bar{z}^{(k+1)}\sim \mathcal{N}(0,I_n)$
\STATE $\bar{v}^{(k+1)}= P(x^{(k+1)})\bar{z}^{(k+1)}$
\STATE $\bar{u}^{(k+1)}=P(x^{(k+1)}) \left(s^{(k+1),\theta}-b\right)(x^{(k+1)}) $
\STATE $x^{(k+\frac{1}{2})} = x^{(k+1)} +\beta_{k+1}^2 \bar{u}^{(k+1)} +\beta_{k+1} \bar{v}^{(k+1)}$
\STATE \MYCOMMENT{solve (\ref{projection-kstep-backward}) by Algorithm~\ref{Newton_solver}}
\STATE $c, \mathrm{flag}=\text{newton\_solver}(x^{(k+1)}, x^{(k+\frac{1}{2})};\xi)$
\IF {flag == true}
\STATE $x^{(k)} := x^{(k+\frac{1}{2})} + \nabla\xi(x^{(k+1)})c$
\ELSE
\STATE discard the trajectory and re-generate
\ENDIF
\ENDFOR
\STATE \textbf{Return} $(x^{(N)},x^{(N-1)},\dots,x^{(0)})$
\end{algorithmic}
\end{algorithm}
\section{Related work}
\begin{table*}[!t]
\centering
\caption{Test negative log-likelihood (NLL) results on earth and climate science datasets. A smaller value indicates better performance.}
\label{earth_comparison}
% \begin{tabular}{@{}lcccc@{}}
\begin{tabular}{lcccc}
\toprule
% & \multicolumn{1}{c}{Volcano (827)} & \multicolumn{1}{c}{Earthquake (6120)} & \multicolumn{1}{c}{Flood (4875)} & \multicolumn{1}{c}{Fire (12809)} \\
& Volcano & Earthquake & Flood & Fire\\
{\scriptsize Dataset size} & {\scriptsize 827} & {\scriptsize 6120} & {\scriptsize 4875} & {\scriptsize 12809} \\
\midrule
RCNF {\scriptsize \citep{mathieu2020riemannian}} & -6.05{\scriptsize $\pm$0.61} & ~0.14{\scriptsize $\pm$0.23} & 1.11{\scriptsize $\pm$0.19} & -0.80{\scriptsize $\pm$0.54} \\
Moser Flow {\scriptsize \citep{rozen2021moser}} & -4.21{\scriptsize $\pm$0.17} & -0.16{\scriptsize $\pm$0.06} & 0.57{\scriptsize $\pm$0.10} & -1.28{\scriptsize $\pm$0.05} \\
CNFM {\scriptsize \citep{BenHamu2022MatchingNF}} & -2.38{\scriptsize $\pm$0.17} & -0.38{\scriptsize $\pm$0.01} & 0.25{\scriptsize $\pm$0.02} & -1.40{\scriptsize $\pm$0.02} \\
RSGM {\scriptsize \citep{de2022riemannian}} & -4.92{\scriptsize $\pm$0.25} & -0.19{\scriptsize $\pm$0.07} & 0.48{\scriptsize $\pm$0.17} & -1.33{\scriptsize $\pm$0.06}\\
RDM {\scriptsize \citep{huang2022riemannian}} & -6.61{\scriptsize $\pm$0.96} & -0.40{\scriptsize $\pm$0.05} & 0.43{\scriptsize $\pm$0.07} & -1.38{\scriptsize $\pm$0.05} \\
% RSGM-improved {\scriptsize \citep{NEURIPS2023_ScalingRDM}} & -4.69{\scriptsize $\pm$0.29} & -0.27{\scriptsize $\pm$0.05} & 0.44{\scriptsize $\pm$0.03} & -1.51{\scriptsize $\pm$0.13} \\
RFM {\scriptsize \citep{chen2024flow}} & -7.93{\scriptsize $\pm$1.67} & -0.28{\scriptsize $\pm$0.08} & 0.42{\scriptsize $\pm$0.05} & -1.86{\scriptsize $\pm$0.11} \\
LogBM {\scriptsize \citep{jo2024generativemodelingmanifoldsmixture} }&-9.52{\scriptsize $\pm$0.87} &-0.30{\scriptsize $\pm$0.06} & 0.42{\scriptsize $\pm$0.08} & -2.47{\scriptsize $\pm$0.11} \\
\midrule
\hspace{-0.2cm} \textit{\scriptsize Ours}\\
RDDPM & -2.16{\scriptsize $\pm$1.92} & -0.17{\scriptsize $\pm$0.10} & 0.49{\scriptsize $\pm$0.09} & -1.48{\scriptsize $\pm$0.10} \\
RDDPM w/ isolated points & -3.57{\scriptsize $\pm$1.05} & -0.29{\scriptsize $\pm$0.04} & 0.43{\scriptsize $\pm$0.05} & -1.56{\scriptsize $\pm$0.08} \\
\bottomrule
\end{tabular}
\end{table*}
\begin{table*}[!t]
\centering
\caption{Test negative log-likelihood (NLL) results on mesh datasets. A smaller value indicates better performance.}\label{sdf_compare}
\begin{tabular}{lcccc}
\toprule
& \multicolumn{2}{c}{\textbf{Stanford Bunny}} & \multicolumn{2}{c}{\textbf{Spot the Cow}} \\
\cmidrule(lr){2-3} \cmidrule(lr){4-5}
& $ k = 50 $ & $ k = 100 $ & $ k = 50 $ & $ k = 100 $ \\
\midrule
RFM w/ Diff. {\scriptsize \citep{chen2024flow}} & 1.48\scriptsize{$\pm$ 0.01} & 1.53\scriptsize{$\pm$ 0.01} & 0.95\scriptsize{$\pm$ 0.05} & 1.08\scriptsize{$\pm$ 0.05} \\
RFM w/ Bihar. {\scriptsize \citep{chen2024flow}}&1.55\scriptsize{$\pm$ 0.01} & 1.49\scriptsize{$\pm$ 0.01} & 1.08\scriptsize{$\pm$ 0.05} & 1.29\scriptsize{$\pm$ 0.05} \\
LogBM w/ Diff. {\scriptsize \citep{jo2024generativemodelingmanifoldsmixture} }&1.42\scriptsize{$\pm$ 0.01} & 1.41\scriptsize{$\pm$ 0.00} & 0.99\scriptsize{$\pm$ 0.03} & 0.97\scriptsize{$\pm$ 0.03}\\
LogBM w/ Bihar. {\scriptsize \citep{jo2024generativemodelingmanifoldsmixture} }& 1.55\scriptsize{$\pm$ 0.02} & 1.45\scriptsize{$\pm$ 0.01} & 1.09\scriptsize{$\pm$ 0.06} & 0.97\scriptsize{$\pm$ 0.02} \\
\midrule
\hspace{-0.2cm} \textit{\scriptsize Ours}\\
RDDPM & 1.36\scriptsize{$\pm$ 0.00} & 1.31\scriptsize{$\pm$ 0.01} & 0.84\scriptsize{$\pm$ 0.00} & 0.77\scriptsize{$\pm$ 0.00} \\
\bottomrule
\end{tabular}
\end{table*}
\paragraph{Denoising diffusion probabilistic models}
DDPMs \citep{NEURIPS2020_ddpm} employ a forward Markov chain to perturb data into noise and a reverse Markov chain to incrementally recover data from noise.
The models are trained to minimize a variational bound on the negative log-likelihood. DDPMs have made remarkable achievement in generative modeling on Euclidean space. However, to our knowledge, prior to this study there was no successful extension of DDPMs to manifolds.
\paragraph{Diffusion models on manifolds}
Riemannian Score-based Generative Models (RSGMs; \citet{de2022riemannian} extend SGMs to Riemannian manifolds.
A major difficulty of RSGMs is due to the fact that the denoising score-matching objective involves heat kernel, which is not known analytically except for very special manifolds. In addition, RSGMs also require sampling of geodesic curves on manifolds.
\citet{de2022riemannian} and~\citet{NEURIPS2023_ScalingRDM} proposed to approximate heat kernel by eigenfunction expansion or Varadhan’s approximation, but these approximations bring in extra errors. In contrast, our method requires neither geodesic curves nor heat kernel.
The Riemannian Diffusion Model (RDM; \citet{huang2022riemannian}) adopts a variational diffusion model framework on Riemannian manifolds. Similar to our method, RDM framework considers
submanifolds embedded in an Euclidean space and it also utilizes
a variational upper bound on the negative log-likelihood as loss function. However, RDM requires to sample SDEs on manifolds using closest-point projection, which in general may be a difficult task. In contrast, the projection adopted in our method is applicable to general submanifolds and the computational cost is lower.
\paragraph{Flow-based generative models on manifolds}
Riemannian continuous normalizing flow models \citep{lou2020neural,mathieu2020riemannian,rozen2021moser,BenHamu2022MatchingNF,chen2024flow} extend the continuous-time flow framework \citep{chen2018neural} to manifolds.
In particular, the methods proposed in \cite{rozen2021moser,BenHamu2022MatchingNF,chen2024flow} are simulation-free for manifolds with simple geometry. The method in
\cite{chen2024flow} can deal with manifolds with general geometry, but it is not feasible for
high-dimensional manifolds, where eigenfunctions of the Laplace-Beltrami operator are typically difficult to compute. Our method requires sampling trajectories in order to evaluate loss function. However, the cost for trajectory simulation can be alleviated by working with a pre-prepared trajectory set that is updated during training with a tunable frequency. Our method does not use eigenfunctions nor metrics and it can be easily applied to high-dimensional manifolds.
\paragraph{Markov Chain Monte Carlo on submanifolds}
Markov Chain Monte Carlo (MCMC) under equality constraints has been studied in several works~\citep{goodman-submanifold,hmc-submanifold-tony,LelievreStoltzZhang2022}. We adopt the same projection scheme from these works to develop our method.
\section{Experiments}
\begin{figure*}[!t]
\centering
\includegraphics[width=0.95\linewidth]{figs/SOn_Statistics_M=5_smooth.pdf}
\caption{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.}\label{fig:SOn_result_5}
\end{figure*}
\begin{figure*}[!t]
\centering
\begin{subfigure}[b]{0.28\textwidth}
\includegraphics[width=0.98\linewidth]{figs/ad.png}
\caption{Alanine dipeptide\label{dipeptide_system}}
\end{subfigure}
\begin{subfigure}[b]{0.70\textwidth}
\includegraphics[width=0.98\linewidth]{figs/dipeptide_fwd_bwd_smooth.pdf}
\caption{Results for alanine dipeptide\label{dipeptide_result}}
\end{subfigure}
\caption{
(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).
\label{dipeptide}}
\end{figure*}
We evaluate our method on datasets from earth and climate science, mesh data on learned manifolds, synthetic dataset of high-dimensional special orthogonal matrices, and dataset of molecular conformations under constraints. The last two novel datasets have not been studied by existing methods.
% We have released the source code to promote reproducibility.
\subsection{Earth and climate science datasets}
We utilize public datasets \citep{data_earthquake,data_volcano,data_flood,data_fire} on 2-D sphere, which are compiled by \citet{mathieu2020riemannian}.
Table~\ref{earth_comparison} summarizes the negative log-likelihood (NLL) of our method alongside results of prior methods. The learned densities are displayed in Figure~\ref{earth_result} of Appendix~\ref{app-subsec-earth}.
We point out a general issue regarding model evaluation on these datasets. We notice that in each dataset there are a few data points (i.e.\ isolated points) whose vicinity contains no other points, and the size of dataset is relatively small compared to its complex distribution. For such datasets, standard dataset splitting for cross-validation results in
non-overlapping training/validation/test sets, whose empirical distributions are considerably different.
As a consequence, computing the NLL on the test set either results in an overconfident assessment of the model (when the test set contains no isolated points) or
requires evaluating the model on isolated points that are completely unseen during training (when the test set contains isolated points).
We argue that the phenomenon discussed above is general and should appear regardless of the methods employed. As a simple solution, we propose to identify the isolated points in validation and test sets and include them into the training set. In particular, evident improvement is achieved when isolated points are included in training set, as shown in the last line of Table~\ref{earth_comparison}. Further discussions are provided in Appendix~\ref{app-subsec-earth}.
% \begin{table*}[!t]
% \centering
% \caption{Test negative log-likelihood (NLL) results on earth and climate science datasets. A smaller value indicates better performance.}
% \label{earth_comparison}
% % \begin{tabular}{@{}lcccc@{}}
% \begin{tabular}{lcccc}
% \toprule
% % & \multicolumn{1}{c}{Volcano (827)} & \multicolumn{1}{c}{Earthquake (6120)} & \multicolumn{1}{c}{Flood (4875)} & \multicolumn{1}{c}{Fire (12809)} \\
% & Volcano & Earthquake & Flood & Fire\\
% {\scriptsize Dataset size} & {\scriptsize 827} & {\scriptsize 6120} & {\scriptsize 4875} & {\scriptsize 12809} \\
% \midrule
% RCNF {\scriptsize \citep{mathieu2020riemannian}} & -6.05{\scriptsize $\pm$0.61} & ~0.14{\scriptsize $\pm$0.23} & 1.11{\scriptsize $\pm$0.19} & -0.80{\scriptsize $\pm$0.54} \\
% Moser Flow {\scriptsize \citep{rozen2021moser}} & -4.21{\scriptsize $\pm$0.17} & -0.16{\scriptsize $\pm$0.06} & 0.57{\scriptsize $\pm$0.10} & -1.28{\scriptsize $\pm$0.05} \\
% CNFM {\scriptsize \citep{BenHamu2022MatchingNF}} & -2.38{\scriptsize $\pm$0.17} & -0.38{\scriptsize $\pm$0.01} & 0.25{\scriptsize $\pm$0.02} & -1.40{\scriptsize $\pm$0.02} \\
% RSGM {\scriptsize \citep{de2022riemannian}} & -4.92{\scriptsize $\pm$0.25} & -0.19{\scriptsize $\pm$0.07} & 0.48{\scriptsize $\pm$0.17} & -1.33{\scriptsize $\pm$0.06}\\
% RDM {\scriptsize \citep{huang2022riemannian}} & -6.61{\scriptsize $\pm$0.96} & -0.40{\scriptsize $\pm$0.05} & 0.43{\scriptsize $\pm$0.07} & -1.38{\scriptsize $\pm$0.05} \\
% % RSGM-improved {\scriptsize \citep{NEURIPS2023_ScalingRDM}} & -4.69{\scriptsize $\pm$0.29} & -0.27{\scriptsize $\pm$0.05} & 0.44{\scriptsize $\pm$0.03} & -1.51{\scriptsize $\pm$0.13} \\
% RFM {\scriptsize \citep{chen2024flow}} & -7.93{\scriptsize $\pm$1.67} & -0.28{\scriptsize $\pm$0.08} & 0.42{\scriptsize $\pm$0.05} & -1.86{\scriptsize $\pm$0.11} \\
% LogBM {\scriptsize \citep{jo2024generativemodelingmanifoldsmixture} }&-9.52{\scriptsize $\pm$0.87} &-0.30{\scriptsize $\pm$0.06} & 0.42{\scriptsize $\pm$0.08} & -2.47{\scriptsize $\pm$0.11} \\
% \midrule
% \hspace{-0.2cm} \textit{\scriptsize Ours}\\
% RDDPM & -2.16{\scriptsize $\pm$1.92} & -0.17{\scriptsize $\pm$0.10} & 0.49{\scriptsize $\pm$0.09} & -1.48{\scriptsize $\pm$0.10} \\
% RDDPM w/ isolated points & -3.57{\scriptsize $\pm$1.05} & -0.29{\scriptsize $\pm$0.04} & 0.43{\scriptsize $\pm$0.05} & -1.56{\scriptsize $\pm$0.08} \\
% \bottomrule
% \end{tabular}
% \end{table*}
\subsection{Mesh data on learned manifolds}
Our method can effectively handle manifolds with general geometries. For demonstration, we examine the Standard Bunny \citep{turk1994zippered} and Spot the Cow \citep{crane2013robust}, two manifolds defined by triangle meshes. The datasets are created according to the $k$-th eigenfunction of the grid Laplacian operator using the same approach described in \citet{jo2024generativemodelingmanifoldsmixture,chen2024flow}.
Similar to \citet{rozen2021moser}, we first learn the function $\xi: \mathbb{R}^3\rightarrow \mathbb{R}$, whose zero level set matches the manifold.
We adopt the approach in \citet{gropp2020implicit}, where
$\xi$ is represented by a neural network and is trained such that on mesh data $\xi$ is close to zero and $|\nabla\xi|$ is close to one.
Using this approach, we obtain a function $\xi$ whose value is at the order $10^{-2}$ on mesh data. Then, we perform a further refinement to the dataset such that all points belong to the learned manifold $\mathcal{M}=\{x\in \mathbb{R}^3|\xi(x)=0\}$ up to a small error $10^{-5}$. The maximal distance between the original data and the refined data is smaller than
$0.017$. Details for training $\xi$ and refining data are left in Appendix \ref{Bunny_Spot_details}.
We perform the training with the learned function $\xi$. Table \ref{sdf_compare} shows that our learned models achieve
evident improvement over prior works in terms of NLL results. Figure \ref{fig:mesh_result} visualizes the agreement between generated samples (learned distributions) and datasets (target distributions).
% \begin{table*}[!t]
% \centering
% \caption{Test negative log-likelihood (NLL) results on mesh datasets. A smaller value indicates better performance.}\label{sdf_compare}
% \begin{tabular}{lcccc}
% \toprule
% & \multicolumn{2}{c}{\textbf{Stanford Bunny}} & \multicolumn{2}{c}{\textbf{Spot the Cow}} \\
% \cmidrule(lr){2-3} \cmidrule(lr){4-5}
% & $ k = 50 $ & $ k = 100 $ & $ k = 50 $ & $ k = 100 $ \\
% \midrule
% RFM w/ Diff. {\scriptsize \citep{chen2024flow}} & 1.48\scriptsize{$\pm$ 0.01} & 1.53\scriptsize{$\pm$ 0.01} & 0.95\scriptsize{$\pm$ 0.05} & 1.08\scriptsize{$\pm$ 0.05} \\
% RFM w/ Bihar. {\scriptsize \citep{chen2024flow}}&1.55\scriptsize{$\pm$ 0.01} & 1.49\scriptsize{$\pm$ 0.01} & 1.08\scriptsize{$\pm$ 0.05} & 1.29\scriptsize{$\pm$ 0.05} \\
% LogBM w/ Diff. {\scriptsize \citep{jo2024generativemodelingmanifoldsmixture} }&1.42\scriptsize{$\pm$ 0.01} & 1.41\scriptsize{$\pm$ 0.00} & 0.99\scriptsize{$\pm$ 0.03} & 0.97\scriptsize{$\pm$ 0.03}\\
% LogBM w/ Bihar. {\scriptsize \citep{jo2024generativemodelingmanifoldsmixture} }& 1.55\scriptsize{$\pm$ 0.02} & 1.45\scriptsize{$\pm$ 0.01} & 1.09\scriptsize{$\pm$ 0.06} & 0.97\scriptsize{$\pm$ 0.02} \\
% \midrule
% \hspace{-0.2cm} \textit{\scriptsize Ours}\\
% RDDPM & 1.36\scriptsize{$\pm$ 0.00} & 1.31\scriptsize{$\pm$ 0.01} & 0.84\scriptsize{$\pm$ 0.00} & 0.77\scriptsize{$\pm$ 0.00} \\
% \bottomrule
% \end{tabular}
% \end{table*}
% \begin{figure}[htbp]
% \centering
% \includegraphics[width=0.98\linewidth]{figs/mesh_summary_all.pdf}
% \caption{First row: datasets and true distributions. Second row: learned samples and distributions.}\label{fig:mesh_result}
% \end{figure}
\subsection{High-dimensional special orthogonal group}
We apply our method to special orthogonal group $\mathrm{SO}(10)$, viewed as a $45$-dimensional submanifold embedded in $\mathbb{R}^{100}$. The synthetic dataset is sampled from a multimodal distribution on $\mathrm{SO}(10)$ with $m=5$ modes.
To assess the quality of generated data, we consider the statistics $\operatorname{tr}(S)$, $\operatorname{tr}(S^2)$, $\operatorname{tr}(S^4)$, and $\operatorname{tr}(S^5)$, where $\operatorname{tr}$ denotes the trace operator of matrices. Further details for the construction of the dataset and the choice of statistics are provided in Appendix~\ref{SOn_details}.
Figure~\ref{fig:SOn_result_5} indicates that our learned model can generate the data distribution accurately. What is more, the distributions of the forward process at intermediate steps are also faithfully reproduced.
% \begin{figure*}[!t]
% \centering
% \includegraphics[width=0.95\linewidth]{figs/SOn_Statistics_M=5_smooth.pdf}
% \caption{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.}\label{fig:SOn_result_5}
% \end{figure*}
\subsection{Alanine dipeptide}
\label{subsec-ad}
We apply our method to alanine dipeptide, a commonly studied model system in bio-physics. The configuration of the system can be characterized by its two dihedral angles $\phi$ and $\psi$ (see Figure~\ref{dipeptide_system}). In this study, we are interested in the configurations of the $10$ non-hydrogen atoms of the system (in $\mathbb{R}^{30}$) with the fixed angle $\phi=-70^\circ$.
Since the manifold is unbounded, we choose a nonzero function $b=-\nabla V$ in the forward process, where $V$ is proportional to the root mean squared deviation (RMSD) from a pre-selected reference configuration $x^{\mathrm{ref}}$. Accordingly, the prior distribution $p(x^{(N)})$ is a single-well distribution centered at $x^{\mathrm{ref}}$.
Furthermore, we model $s^{(k+1),\theta}(x)$ in the reverse process using a network that preserves rotational equivariance and translational invariance.
This, as well as our choice of $b$, guarantee that the distribution $p_{\theta}(x^{(0)})$ generated by our model is invariant under $\mathrm{SE}(3)~$\citep{xu2022geodiff}. We refer to Appendix~\ref{subsec-dipeptide_details} for implementation details and to Appendix~\ref{Equivariant_NN} for theoretical support.
We employ three metrics to assess the quality of the generated configurations: the angle $\psi$, and two RMSDs (denoted by $\mathrm{RMSD}_1$ and $\mathrm{RMSD}_2$) with respect to two pre-defined reference configurations that are selected from two different wells. Figure~\ref{dipeptide_result} illustrates the histograms of these three metrics for the configurations generated by our model and the configurations in the dataset. The solid and dashed lines show the agreement between the distributions of the learned reverse process and the distributions of the forward process at different time steps $k$. In particular, the overlap between the lines in black, which correspond to step $k=0$, demonstrates that the distribution of the generated samples (dashed) closely matches the data distribution (solid).
\section{Conclusion}
We have proposed Riemannian Denoising Diffusion Probabilistic Models for generative modeling on submanifolds. Our method does not rely on sophisticated geometric objects on manifold and it is applicable to high-dimensional manifolds with nontrivial geometry. We have provided a theoretical analysis of our method in the continuous-time limit, which elucidates its connection to Riemannian score-based generative models. We have demonstrated the strong capability of our method on datasets from previous studies and from high-dimensional manifolds that can not be easily studied by existing methods.
% \section*{Impact Statement}
% This paper aims to advance the study of generative models for data on more general manifolds. While this work may have potential societal consequences, none require specific emphasis at this stage.
% , with a particular focus on applications to molecular systems.
% This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
% In the unusual situation where you want a paper to appear in the
% references without citing it in the main text, use \nocite
% \nocite{langley00}
\bibliography{example_paper}
\bibliographystyle{icml2025}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% APPENDIX
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\newpage
\appendix
\onecolumn
\section{Proofs of the continuous-time limit}
\label{app-sec-proofs}
In this section, we prove Theorem~\ref{thm-continuous-limit} and Corollary~\ref{corollary-continuous-limit-kl} in Section~\ref{sec-continuous-limit}.
For notation simplicity, we denote by $\partial_i$ the derivative with respect to $x_i$ in the ambient space, and by $I$ the identity matrix of order $n$. We use subscripts to denote components of a vector and entries of a matrix. Also recall that the orthogonal projection matrix $P(x)\in \mathbb{R}^{n\times n}$ is well defined for $x\in \mathbb{R}^n$ and has the expression
\begin{equation}
P_{ij}(x) = \delta_{ij}-\sum_ {\alpha,\alpha'=1}^{n-d} \partial_i \xi_\alpha(x)
(\nabla\xi^\top\nabla\xi)^{-1}_{\alpha\alpha'}(x)
\partial_j \xi_{\alpha'}(x), \quad 1 \le i,j \le n\,,
\label{exp-pij}
\end{equation}
where $\delta_{ij}$ is the Dirac delta function.
First, we present the proof of Theorem~\ref{thm-continuous-limit}.
\begin{proof}[Proof of Theorem~\ref{thm-continuous-limit}]
Let us write the forward process in Eq.~\eqref{projection-kstep-forward} as
\begin{equation*}
x^{(k+1)} = x^{(k+\frac{1}{2})} + \nabla\xi(x^{(k)}) c(x^{(k+\frac{1}{2})})\,,
\end{equation*}
where $x^{(k+\frac{1}{2})} = x^{(k)} + \sigma_k^2 b(x^{(k)}) + \sigma_k v^{(k)}$ and the dependence of $c$ on $x^{(k+\frac{1}{2})}$ is made explicit.
Applying Lemma~\ref{lemma-expansion-of-projection} at the end of this section, we obtain the
expansion, for $1\le i\le n$,
\begin{equation}\label{expansion-xk}
\begin{aligned}
x^{(k+1)}_i = & x^{(k)}_i +\sum_{j=1}^{n} P_{ij}(x^{(k)}) \Big(\sigma_k^2 b_j(x^{(k)}) +
\sigma_k v^{(k)}_j\Big)\\
& + \frac{1}{2} \sum_{j,l,r,r'=1}^{n}
\Big((I-P)_{ir} P_{r'l} \,\partial_{r'} P_{rj}
\Big)(x^{(k)})
\Big(\sigma_k^2 b_j(x^{(k)}) + \sigma_k v^{(k)}_j\Big)\Big(\sigma_k^2
b_{l}(x^{(k)}) + \sigma_k v^{(k)}_{l}\Big) \\
& + \frac{1}{6} \sum_{j,l,r=1}^{n} \sum_{\eta=1}^{n-d}
\big(\partial_i \xi_\eta\, \partial_{jlr}^3
c_\eta \big)(x^{(k)})
\Big(
x^{(k+\frac{1}{2})}_j-x^{(k)}_j\Big)
\Big(
x^{(k+\frac{1}{2})}_l-x^{(k)}_l\Big)
\Big(
x^{(k+\frac{1}{2})}_r-x^{(k)}_r\Big) \\
& + o\big(|x^{(k+\frac{1}{2})}-x^{(k)}|^3\big) \\
=& x^{(k)}_i + \sigma_k v^{(k)}_i + \sigma_k^2\sum_{j=1}^{n} P_{ij}(x^{(k)}) b_j(x^{(k)}) + \frac{\sigma_k^2}{2}
\sum_{j,r,r'=1}^{n}
\big((I-P)_{ir}\, \partial_{r'} P_{rj}\big)(x^{(k)}) v^{(k)}_j v^{(k)}_{r'} \\
&+ \sigma_k^3 R^{(k)}_i + o(\sigma_k^3)\,,
\end{aligned}
\end{equation}
where we have used the identity $\sum_{j=1}^n P_{ij}(x^{(k)})
v_j^{(k)}=v^{(k)}_i$ (since $v^{(k)}$ is a tangent vector), and $R_i^{(k)}$ is a term that satisfies $\sum_{i'=1}^n
P_{ii'}(x^{(k)})R_{i'}^{(k)}=0$, for $1\le i \le n$.
With the expansion above, we compute the loss function in
Eq.~\eqref{training-loss}. Using Eq.~\eqref{expansion-xk}, the relation
$\beta_{k+1}=\sigma_k=\sqrt{h}g(kh)$, and the assumption that
$s^{(k+1),\theta}(x^{(k+1)}) = s_\theta(x^{(k+1)}, (k+1)h)\in \mathbb{R}^n$, we can derive
\begin{equation}
\begin{aligned}
& \hspace{0.4cm} \beta_{k+1}^2 \Big|P(x^{(k+1)})\Big(s^{(k+1),\theta}(x^{(k+1)}) - b(x^{(k+1)}) +
\frac{x^{(k+1)}-x^{(k)}}{\beta_{k+1}^2}\Big)\Big|^2\\
&= \sigma_{k}^2 \sum_{i=1}^n\bigg|\sum_{i'=1}^n
P_{ii'}(x^{(k+1)})\Big[s_{\theta, i'}(x^{(k+1)}, (k+1)h) - b_{i'}(x^{(k+1)})
+ \sum_{j=1}^n P_{i'j}(x^{(k)}) b_j(x^{(k)}) \\
& \hspace{3.2cm} + \frac{1}{2} \sum_{j,r,r'=1}^n
\big((I-P)_{i'r}\,\partial_{r'} P_{rj}\big)(x^{(k)}) v^{(k)}_j v^{(k)}_{r'}
+ \frac{ v^{(k)}_{i'}}{\sigma_k} + \sigma_k R^{(k)}_{i'} + o(\sigma_k)\Big]\bigg|^2\\
& = \hspace{0.1cm} \mathcal{I}_1 + \mathcal{I}_2 + \mathcal{I}_3 + o(\sigma_k^2)\,,
\end{aligned}
\label{loss-term-as-sum-of-i123}
\end{equation}
where the three terms on the last line are defined as
\begin{equation*}
\begin{aligned}
\mathcal{I}_1 := & \sigma_{k}^2 \sum_{i=1}^n\bigg|\sum_{i'=1}^n
P_{ii'}(x^{(k+1)})\Big[s_{\theta,i'}(x^{(k+1)}, (k+1)h) - b_{i'}(x^{(k+1)}) + \sum_{j=1}^n P_{i'j}(x^{(k)}) b_j(x^{(k)}) \\
& \hspace{4.2cm} + \frac{1}{2} \sum_{j,r,r'=1}^n \big((I-P)_{i'r}\,\partial_{r'} P_{rj}\big)(x^{(k)}) v^{(k)}_j v^{(k)}_{r'} + \sigma_k R^{(k)}_{i'}
\Big]\bigg|^2\,,\\
\mathcal{I}_2 := & \sum_{i=1}^n\Big(\sum_{i'=1}^n P_{ii'}(x^{(k+1)}) v^{(k)}_{i'}\Big)^2\,,\\
\mathcal{I}_3: = &2\sigma_{k} \sum_{i,i'=1}^n
P_{ii'}(x^{(k+1)})\bigg[s_{\theta,i'}(x^{(k+1)},(k+1)h) - b_{i'}(x^{(k+1)}) + \sum_{j=1}^n P_{i'j}(x^{(k)}) b_j(x^{(k)}) \\
& \hspace{4.2cm} + \frac{1}{2} \sum_{j,r,r'=1}^n \big((I-P)_{i'r}\,\partial_{r'}
P_{rj}\big)(x^{(k)}) v^{(k)}_j v^{(k)}_{r'} + \sigma_k R^{(k)}_{i'}
\bigg] v^{(k)}_{i}\,,
\end{aligned}
\end{equation*}
respectively. In the following, we derive the expansions of the three terms above. For $\mathcal{I}_1$, expanding the functions $P, s_\theta, b$ using Eq.~\eqref{expansion-xk}, we can derive
\begin{equation}
\begin{aligned}
\mathcal{I}_1 =& \sigma_{k}^2 \sum_{i=1}^n\bigg|\sum_{i'=1}^n P_{ii'}(x^{(k)})\Big[s_{\theta,i'}(x^{(k)},kh) - b_{i'}(x^{(k)}) +
\sum_{j=1}^n P_{i'j}(x^{(k)}) b_j(x^{(k)}) \\
& \hspace{4.2cm} + \frac{1}{2}
\sum_{j,r,r'=1}^n \big((I-P)_{i'r}\,\partial_{r'} P_{rj}\big)(x^{(k)}) v^{(k)}_j v^{(k)}_{r'}\Big] + o(1) \bigg|^2 \\
=& \sigma_{k}^2 \sum_{i=1}^n\Big|\sum_{i'=1}^n
P_{ii'}(x^{(k)})s_{\theta,i'}(x^{(k)},kh)\Big|^2 + o(\sigma_k^2)\,,
\end{aligned}
\label{expansion-I1}
\end{equation}
where we have used the relations $P^2=P$ and $P(I-P)=0$ satisfied by the orthogonal projection matrix $P$ to derive the second equality.
For $\mathcal{I}_2$, using the relation $P^2=P$ and Eq.~\eqref{expansion-xk}, we can compute
\begin{align}
\mathcal{I}_2 =& \sum_{i,i'=1}^n P_{ii'}(x^{(k+1)})
v^{(k)}_{i}v^{(k)}_{i'} \notag\\
=&\sum_{i,i'=1}^n P_{ii'}(x^{(k)}) v^{(k)}_{i}v^{(k)}_{i'}
+ \sum_{i,i',r=1}^n \partial_r P_{ii'} (x^{(k)})
\big(x^{(k+1)}_r - x^{(k)}_r\big) v^{(k)}_{i}v^{(k)}_{i'} \notag\\
& + \frac{1}{2} \sum_{i,i'=1}^n\sum_{r,r'=1}^n \partial^2_{rr'}
P_{ii'}(x^{(k)}) v^{(k)}_{i}v^{(k)}_{i'}
(x^{(k+1)}_r - x^{(k)}_r) (x^{(k+1)}_{r'} - x^{(k)}_{r'}) + o(|x^{(k+1)} - x^{(k)}|^2)\notag \\
=& |v^{(k)}|^2 + \sum_{i,i',r=1}^n \partial_r P_{ii'}(x^{(k)})
\big(x^{(k+1)}_r - x^{(k)}_r\big) v^{(k)}_{i}v^{(k)}_{i'} + \frac{\sigma_k^2}{2} \sum_{i,i'=1}^n\sum_{r,r'=1}^n \partial^2_{rr'}
P_{ii'}(x^{(k)}) v^{(k)}_{i}v^{(k)}_{i'} v^{(k)}_r v^{(k)}_{r'} + o(\sigma_k^2) \,.
\label{expansion-I2-pre-1}
\end{align}
Let's compute the three terms in Eq.~\eqref{expansion-I2-pre-1}. Using the expression of $P_{ii'}$ in (\ref{exp-pij}), the fact that $\sum_{i=1}^n \partial_i \xi_\alpha(x^{(k)}) v^{(k)}_i = 0$, and the product rule, it is straightforward to verify that, for $1 \le r\le n$,
\begin{equation}
\begin{aligned}
\sum_{i,i'=1}^n \partial_{r} P_{ii'} (x^{(k)}) v^{(k)}_{i}v^{(k)}_{i'}
=
-\sum_{i,i'=1}^n \partial_{r}
\Big(\sum_{\alpha,\alpha'=1}^{n-d} \partial_i \xi_\alpha
(\nabla\xi^\top\nabla\xi)^{-1}_{\alpha\alpha'}
\partial_{i'} \xi_{\alpha'}\Big)
(x^{(k)}) v^{(k)}_{i}v^{(k)}_{i'}
= 0\,.
\end{aligned}
\label{expansion-I2-pre-2}
\end{equation}
Similarly, we can verify that
\begin{equation}
\sum_{i,i',r,r'=1}^n \partial^2_{rr'}
P_{ii'}(x^{(k)})\, v^{(k)}_{i}v^{(k)}_{i'} v^{(k)}_r v^{(k)}_{r'}
= -2\sum_{i,i',r,r'=1}^n\sum_{\alpha,\alpha'=1}^{n-d}\Big(\partial^2_{ir} \xi_\alpha (\nabla\xi^\top\nabla\xi)^{-1}_{\alpha\alpha'} \partial^2_{i'r'} \xi_{\alpha'}\Big)(x^{(k)}) \,v^{(k)}_{i}v^{(k)}_{i'} v^{(k)}_r v^{(k)}_{r'}\,.
\label{expansion-I2-pre-3}
\end{equation}
Hence, substituting Eq.~\eqref{expansion-I2-pre-2} and Eq.~\eqref{expansion-I2-pre-3} into Eq.~\eqref{expansion-I2-pre-1}, we obtain
\begin{align}
\mathcal{I}_2 =&
|v^{(k)}|^2 - \sigma_k^2
\sum_{i,i',r,r'=1}^n\sum_{\alpha,\alpha'=1}^{n-d}\Big(\partial^2_{ir} \xi_\alpha (\nabla\xi^\top\nabla\xi)^{-1}_{\alpha\alpha'}\partial^2_{i'r'} \xi_{\alpha'}\Big)(x^{(k)}) \,v^{(k)}_{i}v^{(k)}_{i'} v^{(k)}_r v^{(k)}_{r'}
+ o(\sigma_k^2)\,. \label{expansion-I2}
\end{align}
For $\mathcal{I}_3$, we have
\begin{align}
\mathcal{I}_3 = & 2\sigma_{k} \sum_{i,i'=1}^n P_{ii'}(x^{(k)})\bigg[s_{\theta,i'}(x^{(k)},kh) - b_{i'}(x^{(k)}) + \sum_{j=1}^n P_{i'j}(x^{(k)}) b_j(x^{(k)}) \notag \\
& \hspace{4.2cm} + \frac{1}{2} \sum_{j,r,r'=1}^n
\big((I-P)_{i'r}\, \partial_{r'} P_{rj}\big)(x^{(k)}) v^{(k)}_j v^{(k)}_{r'} + \sigma_k R^{(k)}_{i'} \bigg] v^{(k)}_{i}\notag\\
& + 2\sigma^2_k\sum_{i,i',r=1}^n \partial_r (P_{ii'}(s_{\theta,i'} -
b_{i'}))(x^{(k)},kh)\, v^{(k)}_{r} v^{(k)}_{i} \notag\\
& + 2\sigma^2_k\sum_{i,i',r,j=1}^n \partial_r P_{ii'}(x^{(k)})
P_{i'j}(x^{(k)}) b_j(x^{(k)}) \, v^{(k)}_{r} v^{(k)}_{i}\notag\\
& + \sigma^2_k\sum_{i,i',r,j,r',j'=1}^n \partial_r P_{ii'}(x^{(k)})
\big((I-P)_{i'j'}\, \partial_{r'} P_{j'j}\big)(x^{(k)}) v^{(k)}_j v^{(k)}_{r'}
\, v^{(k)}_{r} v^{(k)}_{i}+ o(\sigma_k^2)\notag\\
=&
2 \sigma_{k}
\sum_{i,i'=1}^n P_{ii'}(x^{(k)})s_{\theta,i'}(x^{(k)},kh) v^{(k)}_{i} + 2\sigma^2_k \sum_{i,i',r=1}^n
\partial_{r} (P_{ii'}(s_{\theta,i'} -
b_{i'}))(x^{(k)},kh) v^{(k)}_{r} v^{(k)}_{i} \notag\\
& + 2\sigma^2_k\sum_{i,i',r,j=1}^n \partial_r P_{ii'}(x^{(k)})
P_{i'j}(x^{(k)}) b_j(x^{(k)}) \, v^{(k)}_{r} v^{(k)}_{i} \notag \\
& + \sigma^2_k\sum_{i,i',r,j,r',j'=1}^n \partial_r P_{ii'}(x^{(k)})
\big((I-P)_{i'j'}\, \partial_{r'} P_{j'j}\big)(x^{(k)}) v^{(k)}_j v^{(k)}_{r'}
\, v^{(k)}_{r} v^{(k)}_{i}+ o(\sigma_k^2)\,, \label{expansion-I3-0}
\end{align}
where we have used Taylor expansion with Eq.~\eqref{expansion-xk} and the fact that $|s_{\theta}(x,(k+1)h) - s_{\theta}(x,kh)| =
O(h) = O(\sigma_k^2)$ to derive the first equality, and we have used the relations $P^2=P$, $P(I-P)=0$, and
$\sum_{i'=1}^n P_{ii'}(x^{(k)})R^{(k)}_{i'}=0$ to derive the second equality. We further simplify the last two terms in the expression above. Notice that, similar to Eq.~\eqref{expansion-I2-pre-2}, we can verify that
\begin{equation}
\sum_{i,i',r=1}^n \partial_r P_{ii'}(x^{(k)})
P_{i'j}(x^{(k)}) \, v^{(k)}_{r} v^{(k)}_{i} = 0\,, \quad 1 \le j \le n\,.\label{I3-pre-1}
\end{equation}
For the last term in Eq.~\eqref{expansion-I3-0}, we can derive
\begin{equation}
\begin{aligned}
&\sum_{i,i',r,j,r',j'=1}^n \partial_r P_{ii'}(x^{(k)})
\big((I-P)_{i'j'}\, \partial_{r'} P_{j'j}\big)(x^{(k)}) v^{(k)}_j v^{(k)}_{r'}
\, v^{(k)}_{r} v^{(k)}_{i}\\
= &\sum_{i',r,j,r',j'=1}^n
\big(\partial_r P_{i'j'}\, \partial_{r'} P_{j'j}\big)(x^{(k)}) v^{(k)}_j v^{(k)}_{r'}
\, v^{(k)}_{r} v^{(k)}_{i'}\\
= &\sum_{i',r,j,r',j'=1}^n \bigg[\sum_{\alpha,\alpha',\eta,\eta'=1}^{n-d} \Big(\partial^2_{i'r} \xi_\alpha
(\nabla\xi^\top\nabla\xi)^{-1}_{\alpha\alpha'}
\partial_{j'} \xi_{\alpha'}\Big)
\Big( \partial^2_{jr'} \xi_\eta
(\nabla\xi^\top\nabla\xi)^{-1}_{\eta\eta'}
\partial_{j'} \xi_{\eta'}\Big)\bigg](x^{(k)})v^{(k)}_j v^{(k)}_{r'}
\, v^{(k)}_{r} v^{(k)}_{i'}\\
= &\sum_{i',r,j,r'=1}^n \bigg[\sum_{\alpha,\eta=1}^{n-d} \Big(\partial^2_{i'r} \xi_\alpha
(\nabla\xi^\top\nabla\xi)^{-1}_{\alpha\eta}
\partial^2_{jr'} \xi_\eta
\Big)\bigg](x^{(k)})v^{(k)}_j v^{(k)}_{r'}
\, v^{(k)}_{r} v^{(k)}_{i'}\,,
\end{aligned}
\label{I3-pre-2}
\end{equation}
where the first equation follows by applying the product rule to the identity $P(I-P)=0$ and using the relation $\sum_{i=1}^n P_{ii'}v_i^{(k)}=v_{i'}^{(k)}$, the second equation follows from the expression (\ref{exp-pij}) and the fact that several terms vanish due to the orthogonality relation $\sum_{i=1}^n \partial_i \xi_\alpha(x^{(k)}) v_i^{(k)}=0$, and the last equation follows from the fact that $\sum_{j'=1}^n \partial_{j'}\xi_{\alpha'}\partial_{j'}\xi_{\eta'}=(\nabla\xi^\top \nabla\xi)_{\alpha'\eta'}$.
Combining Eq.~\eqref{expansion-I3-0}, Eq.~\eqref{I3-pre-1},
and Eq.~\eqref{I3-pre-2}, we obtain
\begin{align}
\mathcal{I}_3
= &
2 \sigma_{k}
\sum_{i,i'=1}^n P_{ii'}(x^{(k)})s_{\theta,i'}(x^{(k)},kh) v^{(k)}_{i} + 2\sigma^2_k \sum_{i,i',r=1}^n
\partial_{r} (P_{ii'}(s_{\theta,i'} -
b_{i'}))(x^{(k)},kh) v^{(k)}_{r} v^{(k)}_{i} \notag\\
&\hspace{-0.1cm}+\sigma^2_k\sum_{i',r,j,r'=1}^n \bigg[\sum_{\alpha,\eta=1}^{n-d} \Big(\partial^2_{i'r} \xi_\alpha
(\nabla\xi^\top\nabla\xi)^{-1}_{\alpha\eta}
\partial^2_{jr'} \xi_\eta
\Big)\bigg](x^{(k)})v^{(k)}_j v^{(k)}_{r'}
\, v^{(k)}_{r} v^{(k)}_{i'}+ o(\sigma_k^2)\,.\label{expansion-I3}
\end{align}
Substituting Eq.~\eqref{expansion-I1}, Eq.~\eqref{expansion-I2}, and Eq.~\eqref{expansion-I3} into Eq.~\eqref{loss-term-as-sum-of-i123}, we obtain (after cancellation of terms in $\mathcal{I}_2$ and $\mathcal{I}_3$)
\begin{equation}
\begin{aligned}
& \beta_{k+1}^2 \bigg|P(x^{(k+1)})\Big(s^{(k+1),\theta}(x^{(k+1)}) -
b(x^{(k+1)}) + \frac{x^{(k+1)}-x^{(k)}}{\beta_{k+1}^2}\Big)\bigg|^2\\
=& \mathcal{I}_1 + \mathcal{I}_2 + \mathcal{I}_3 + o(\sigma_k^2) \\
= &
\sigma_{k}^2 \big|
P(x^{(k)})s_{\theta}(x^{(k)},kh)\big|^2
+ \big|v^{(k)}\big|^2+ 2\sigma^2_k \sum_{i,i',r=1}^n
\partial_r (P_{ii'}(s_{\theta,i'} -
b_{i'}))(x^{(k)},kh) v^{(k)}_{r} v^{(k)}_{i} \\
& +2 \sigma_{k} \sum_{i,i'=1}^n P_{ii'}(x^{(k)})s_{\theta,i'}(x^{(k)},kh) v^{(k)}_{i} + o(\sigma_k^2)\,.
\end{aligned}
\label{loss-term-expansion}
\end{equation}
Since $v^{(k)}$ is a standard Gaussian random variable in $T_{x^{(k)}}\mathcal{M}$, taking expectation in Eq.~\eqref{loss-term-expansion} and substituting it into
Eq.~\eqref{training-loss}, we obtain
\begin{equation}
\mbox{Loss}^{(N)}(\theta)
=
\frac{Nd}{2} +
\mathbb{E}_{\mathbb{Q}^{(N)}} \bigg[\sum_{k=0}^{N-1}\sigma_{k}^2
\Big(\frac{1}{2} \big|
P(x^{(k)})s_{\theta}(x^{(k)},kh)\big|^2
+ \sum_{i,i',r=1}^n \partial_r
(P_{ii'}(s_{\theta,i'} - b_{i'}))(x^{(k)},kh) \, P_{ri}(x^{(k)})\Big)\bigg]
+ o(1)\,.
\end{equation}
In the above, we have used the identities
$\mathbb{E}(v^{(k)}_i)=0$, $\mathbb{E}(|v^{(k)}|^2)=d$, and $\mathbb{E}(v^{(k)}_r v^{(k)}_i)=P_{ri}(x^{(k)})$, where the last one can be verified using the fact that $v^{(k)}=P(x^{(k)}) z^{(k)}$, with $z^{(k)}$ being a standard Gaussian random variable in $\mathbb{R}^n$.
Taking the limit $N\rightarrow +\infty$, and using the fact that the forward process $x^{(k)}$ converges to the SDE~(\ref{sde-manifold})~\citep{projection_diffusion}, we can derive
\begin{align*}
&\lim_{N\rightarrow+\infty}\Big(\mbox{Loss}^{(N)}(\theta)-\frac{Nd}{2}\Big) \\
=& \lim_{N\rightarrow+\infty}
\mathbb{E}_{\mathbb{Q}^{(N)}} \bigg[\sum_{k=0}^{N-1}\sigma_{k}^2
\Big(\frac{1}{2} \big|
P(x^{(k)})s_{\theta}(x^{(k)},kh)\big|^2 + \sum_{i,i',r=1}^n \partial_r
(P_{ii'}(s_{\theta,i'} - b_{i'}))(x^{(k)},kh) P_{ri}(x^{(k)})\Big)\bigg] \\
= & \mathbb{E}_{\mathbb{Q}}\int_{0}^T \bigg[\frac{1}{2}\big|P(X_t)s_{\theta}(X_t, t)\big|^2 +
\mbox{div}_{\mathcal{M}} \big(P(s_{\theta} - b)\big)(X_t,t)\bigg] g^2(t) dt\\
= & \int_{0}^T \bigg[\int_{\mathcal{M}} \Big(\frac{1}{2}\big| P(x)s_{\theta}(x, t)\big|^2 +
\mbox{div}_{\mathcal{M}} \big(P(s_{\theta} - b)\big)(x,t)\Big) p(x,t)\, d\sigma_{\mathcal{M}}(x)\bigg] g^2(t)\,dt\\
=& \int_{0}^T \bigg[\int_{\mathcal{M}} \Big(\frac{1}{2}\big| P(x)s_{\theta}(x, t)\big|^2
- P(x)\big(s_{\theta}(x,t) - b(x)\big) \cdot \nabla_{\mathcal{M}} \log
p(x,t)\Big) p(x,t)\, d\sigma_{\mathcal{M}}(x)\bigg] g^2(t)\,dt \\
=& \frac{1}{2}\int_{0}^T \bigg[\int_{\mathcal{M}} \big| P(x)s_{\theta}(x,
t)-\nabla_{\mathcal{M}} \log p(x,t)\big|^2 p(x,t)\,
d\sigma_{\mathcal{M}}(x)\bigg] g^2(t)\,dt
\\
&+ \int_{0}^T \bigg[\int_{\mathcal{M}} \Big(\big(P(x)b(x)-\frac{1}{2}\nabla_{\mathcal{M}}
\log p(x,t)\big)\cdot \nabla_{\mathcal{M}} \log p(x,t)\Big) p(x,t)\,d\sigma_{\mathcal{M}}(x)\bigg] g^2(t)\,dt \\
=& \mathbb{E}_{\mathbb{Q}}\bigg[\frac{1}{2}\int_{0}^T \big| P(X_t)s_{\theta}(X_t,
t)-\nabla_{\mathcal{M}} \log p(X_t,t)\big|^2 g^2(t)\,dt \\
&+ \int_{0}^T \Big(P(X_t)b(X_t)-\frac{1}{2}\nabla_{\mathcal{M}} \log
p(X_t,t)\Big)\cdot \nabla_{\mathcal{M}} \log p(X_t,t) \, g^2(t) \,dt \bigg]\,,
\end{align*}
where we have used integration by parts on $\mathcal{M}$, and the expression $\mbox{div}_{\mathcal{M}} f=\sum_{i,r=1}^n P_{ir}\partial_r f_i$
for $f:\mathcal{M}\rightarrow \mathbb{R}^n$ (which can be verified using Lemma~A.1 in~\citet{zhang2020}).
\end{proof}
Next, we present the proof of Corollary~\ref{corollary-continuous-limit-kl}.
\begin{proof}[Proof of Corollary~\ref{corollary-continuous-limit-kl}]
Using the assumption $\beta_{k+1}=\sigma_k$, the projection
scheme in Eq.~\eqref{projection-kstep-forward} and the relation $P(x^{(k)})\nabla\xi(x^{(k)})=0$, we can simplify the constant $C^{(N)}$ in Eq.~\eqref{constant-c} as
\begin{equation}
C^{(N)}=-\mathbb{E}_{\mathbb{Q}^{(N)}} \Big(\log p(x^{(N)})+\frac{1}{2}\sum_{k=0}^{N-1} |v^{(k)}|^2
+ \sum_{k=0}^{N-1}\log \big(1-\epsilon_{x^{(k)}}^{(\sigma_k)}\big)
\Big).
\end{equation}
Therefore, using
the definition of relative entropy (see Eq.~\eqref{relative-entropy-formulation-of-variational-bound}), the loss function in Eq.~\eqref{training-loss}, the constant $C^{(N)}$ in Eq.~\eqref{constant-c}, and
applying Theorem~\ref{thm-continuous-limit}, we have
\begin{equation}
\begin{aligned}
& \lim_{N\rightarrow +\infty}
H\big(\overleftarrow{\mathbb{Q}}^{(N)}\,|\,\mathbb{P}^{(N)}_\theta\big)\\
=& \lim_{N\rightarrow +\infty} \left[\mathrm{Loss}^{(N)}(\theta) + C^{(N)} + \mathbb{E}_{\mathbb{Q}^{(N)}}\Big(\sum_{k=0}^{N-1}\log \big(1-\epsilon_{x^{(k+1)},_\theta}^{(\sigma_k)}\big)\Big)\right] + \mathbb{E}_{q_0}(\log q_0) \\
=& \lim_{N\rightarrow +\infty} \left[\mathrm{Loss}^{(N)}(\theta) -\mathbb{E}_{\mathbb{Q}^{(N)}} \left(\log p(x^{(N)})
+\frac{1}{2}\sum_{k=0}^{N-1} |v^{(k)}|^2
+ \sum_{k=0}^{N-1}\log \frac{1-\epsilon_{x^{(k)}}^{(\sigma_k)}}{1-\epsilon_{x^{(k+1)},\theta}^{(\sigma_k)}}\right)\right] + \mathbb{E}_{q_0}(\log q_0) \\
=& \lim_{N\rightarrow +\infty} \Big(\mathrm{Loss}^{(N)}(\theta) -\mathbb{E}_{\mathbb{Q}^{(N)}} \log p(x^{(N)})
-\frac{Nd}{2}\Big) + \mathbb{E}_{q_0}(\log q_0) \\
=& \mathbb{E}_{\mathbb{Q}} \bigg[\log p(X_0,0) -\log p(X_T,T) + \frac{1}{2}\int_{0}^T \big|
P(X_t)s_{\theta}(X_t, t)-\nabla_{\mathcal{M}} \log p(X_t,t)\big|^2 g^2(t)\, dt \\
& + \int_{0}^T \Big(P(X_t)b(X_t)-\frac{1}{2}\nabla_{\mathcal{M}} \log
p(X_t,t)\Big)\cdot \nabla_{\mathcal{M}} \log p(X_t,t)\, g^2(t)\, dt
\bigg] \,,
\end{aligned}
\label{proof-corolloary-1}
\end{equation}
where the third equality follows because the terms containing $\epsilon_{x^{(k)}}^{(\sigma_k)}$ and $\epsilon_{x^{(k+1)},\theta}^{(\sigma_k)}$ vanish and $\lim_{N\rightarrow +\infty}\mathbb{E}_{\mathbb{Q}^{(N)}}(\sum_{k=0}^{N-1}|v^{(k)}|^2-Nd)=0$, both of which can be verified using the asymptotic expression of
complementary error function.
Note that the density $p(x,t)$ of SDE~(\ref{sde-manifold}) solves the Fokker-Planck equation
\begin{equation}
\frac{\partial p}{\partial t}(x,t) = -g^2(t) \mbox{div}_{\mathcal{M}} \big(P(x) b(x) p(x,t)\big) + \frac{g^2(t)}{2} \Delta_{\mathcal{M}} p(x,t)\,, \quad x\in \mathcal{M},\quad t\in[0,T]\,.
\label{fp}
\end{equation}
Therefore, we have
\begin{equation}
\begin{aligned}
& \mathbb{E}_{\mathbb{Q}} \Big[\log p(X_0,0) -\log p(X_T,T)\Big] \\
= &\int_{\mathcal{M}} \log p(x,0) p(x,0)\, d\sigma_{\mathcal{M}}(x) -
\int_{\mathcal{M}} \log p(x,T) p(x,T)\, d\sigma_{\mathcal{M}}(x) \\
= & -\int_0^T \frac{d}{dt} \Big(\int_{\mathcal{M}} \log p(x,t)\, p(x,t)\, d\sigma_{\mathcal{M}}(x)\Big) \,dt \\
= & -\int_0^T \bigg[\int_{\mathcal{M}} \big(\log p(x,t) + 1\big) \frac{\partial p}{\partial t} (x,t)\, d\sigma_{\mathcal{M}}(x)\bigg] dt \\
= & -\int_0^T \bigg[\int_{\mathcal{M}} \big(\log p(x,t) + 1\big)
\Big(-g^2(t) \mbox{div}_{\mathcal{M}} \big(P(x) b(x) p(x,t)\big) + \frac{g^2(t)}{2} \Delta_{\mathcal{M}} p(x,t)\Big) d\sigma_{\mathcal{M}}(x)\bigg] dt \\
= & -\int_0^T \bigg[\int_{\mathcal{M}} \Big(\big(P(x) b(x) -
\frac{1}{2}\nabla_{\mathcal{M}} \log p(x,t)\big) \cdot \nabla_{\mathcal{M}} \log p(x,t)\Big)\, p(x,t)\, d\sigma_{\mathcal{M}}(x)\bigg] g^2(t) \,dt \\
=& -\mathbb{E}_{\mathbb{Q}} \bigg[\int_{0}^T \Big(P(X_t)b(X_t)-\frac{1}{2}\nabla_{\mathcal{M}} \log
p(X_t,t)\Big)\cdot \nabla_{\mathcal{M}} \log p(X_t,t) \, g^2(t) \,dt\bigg] \,,
\end{aligned}
\label{proof-corolloary-2}
\end{equation}
where we have used Eq.~\eqref{fp} to derive the fourth equality and integration by parts on $\mathcal{M}$ to derive the fifth equality.
Combining Eq.~\eqref{proof-corolloary-1} and Eq.~\eqref{proof-corolloary-2}, we obtain
\begin{equation}
\lim_{N\rightarrow +\infty}
H\big(\overleftarrow{\mathbb{Q}}^{(N)}\,|\,\mathbb{P}^{(N)}_\theta\big)
= \mathbb{E}_{\mathbb{Q}} \bigg[\frac{1}{2}\int_{0}^T \big| P(X_t)s_{\theta}(X_t, t)-\nabla_{\mathcal{M}} \log p(X_t,t)\big|^2 g^2(t)\, dt\,\bigg]\,.
\label{proof-corrolary-part-1}
\end{equation}
Finally, note that $\overleftarrow{\mathbb{Q}}$ is the path measure of the time-reversal $Y_t=X_{T-t}$ of SDE~(\ref{sde-manifold}), which satisfies~\citep[Theorem~3.1]{de2022riemannian}
\begin{equation}
dY_t = g^2(T-t) \Big(-P(Y_t) b(X_t)+ \nabla_{\mathcal{M}} \log p(Y_t, T-t)\Big)\,dt + g(T-t) dW^\mathcal{M}_t\,, \quad t\in[0,T]\,,
\label{sde-manifold-reversal}
\end{equation}
and $\mathbb{P}_\theta$ is the path measure of SDE~(\ref{sde-manifold-reversed-theta}).
Applying Girsanov's theorem~\citep[Theorem~8.1.2]{analysis_manifold}, we obtain
\begin{equation}
\begin{aligned}
\frac{d \mathbb{P}_\theta}{d \overleftarrow{\mathbb{Q}}} =& \exp\Big(\int_0^T g^2(T-t) \big(P(Y_t) s_\theta(Y_t,T-t)-\nabla_{\mathcal{M}} \log p(Y_t, T-t)\big)\cdot dW_t^\mathcal{M} \\
& \hspace{0.5cm} - \frac{1}{2}\int_0^T \big|P(Y_t)s_\theta(Y_t,T-t)-\nabla_{\mathcal{M}} \log p(Y_t, T-t)\big|^2 g^2(T-t)\,dt\Big)\,,
\end{aligned}
\label{girsanov}
\end{equation}
where $W_t^\mathcal{M}$ is a Brownian motion on $\mathcal{M}$ under $\overleftarrow{\mathbb{Q}}$. Therefore, we have
\begin{equation}
\begin{aligned}
H(\overleftarrow{\mathbb{Q}}\,|\,\mathbb{P}_\theta) = &
\mathbb{E}_{\overleftarrow{\mathbb{Q}}}\Big(\log \frac{d\overleftarrow{\mathbb{Q}}}{d\mathbb{P}_\theta}\Big) \\
= & \mathbb{E}_{\overleftarrow{\mathbb{Q}}}\Big(
\frac{1}{2}\int_0^T \big|P(Y_t)s_\theta(Y_t,T-t)-\nabla_{\mathcal{M}} \log p(Y_t, T-t)\big|^2 g^2(T-t)\,dt
\Big) \\
= & \mathbb{E}_{\mathbb{Q}}\Big(
\frac{1}{2}\int_0^T \big|P(X_t) s_\theta(X_t,t)-\nabla_{\mathcal{M}} \log p(X_t, t)\big|^2 g^2(t)\,dt
\Big)\,,
\end{aligned}
\label{proof-corrolary-part-2}
\end{equation}
where the second equality follows from the fact that the stochastic integration in Eq.~\eqref{girsanov} vanishes after taking logarithm and expectation, and the third equality follows by a change of variable $t\leftarrow T-t$ and the fact that $Y_t=X_{T-t}$. The conclusion is obtained after combining Eq.~\eqref{proof-corrolary-part-1} and Eq.~\eqref{proof-corrolary-part-2}.
\end{proof}
Finally, we present the technical lemma
on the projection scheme in Eq.~\eqref{projection}, which was used in the proof of Theorem~\ref{thm-continuous-limit}.
\begin{lemma}
Given $x\in\mathcal{M}$ and $x'\in\mathbb{R}^n$, the solution to the problem
\begin{equation}
y = x' + \nabla\xi(x) c(x'), \quad c(x')\in \mathbb{R}^{n-d}, ~~\mbox{s.t.}~~\xi(y) = 0
\label{projection-repeat}
\end{equation}
has the following expansion as $x'$ approaches to $x$
\begin{equation}
\begin{aligned}
\partial_j c_\eta (x) =& - \sum_{\alpha=1}^{n-d}(\nabla\xi^\top \nabla\xi)^{-1}_{\eta\alpha}(x) \partial_j
\xi_\alpha(x)\,, \quad 1 \le j \le n\,,\\
\partial^2_{jl} c_\eta(x) =&
\sum_{\alpha=1}^{n-d}(\nabla\xi^\top\nabla\xi)^{-1}_{\eta\alpha}(x)
\sum_{r,r'=1}^{n}\Big(\partial_r\xi_\alpha \partial_{r'}
P_{rj}\, P_{r'l} \Big)(x)\,, \quad 1 \le j,l \le n\,,
\end{aligned}
\label{lagrange-multiplier-diff}
\end{equation}
for $1\le \eta \le n-d$.
Moreover, as $x'$ approaches to $x$, the following expansion of $y$ in Eq.~\eqref{projection-repeat} holds
\begin{equation}
\begin{aligned}
y_i =& x_i + \sum_{j=1}^n P_{ij}(x) (x'_j - x_j) + \frac{1}{2}
\sum_{j,l=1}^n\bigg[\sum_{r,r'=1}^n \big((I-P)_{ir} P_{r'l}\partial_{r'}
P_{rj} \big)(x)\bigg] (x'_j -
x_j)(x'_{l} - x_{l}) \\
& + \frac{1}{6}
\sum_{j,l,r=1}^n
\bigg(
\sum_{\eta=1}^{n-d}
\partial_i \xi_\eta(x)\,
\partial^3_{jlr} c_\eta(x)
\bigg)
(x'_j - x_j)(x'_{l} - x_{l})(x'_{r} - x_{r}) + o(|x'-x|^3)\,,
\end{aligned}
\label{constraint-expansion}
\end{equation}
where $1\le i \le n$.
\label{lemma-expansion-of-projection}
\end{lemma}
\begin{proof}
Differentiating (with respect to $x'$) the constraint equation
\begin{equation*}
\xi_\alpha(x' + \nabla\xi(x)c(x')) = 0\,, \quad \alpha = 1,\dots, n-d\,,
\end{equation*}
we get
\begin{equation}
\sum_{r=1}^n
\partial_r \xi_\alpha \big(x'+\nabla\xi(x)c(x')\big)
\Big(\delta_{rj} + \sum_{\eta=1}^{n-d}
\partial_r \xi_\eta(x)\partial_{j} c_\eta(x')\Big) = 0\,, \quad 1 \le j \le n\,.
\label{1st-derivative-constraint-eqn}
\end{equation}
Setting $x'=x$ in Eq.~\eqref{1st-derivative-constraint-eqn} (notice that $c(x')=0$ when $x'=x$) and multiplying both
sides by $(\nabla\xi^\top \nabla\xi)^{-1}(x)$, we obtain the first identity in Eq.~\eqref{lagrange-multiplier-diff}.
In particular, using Eq.~\eqref{exp-pij}, we have
\begin{equation}
\delta_{rj} + \sum_{\eta=1}^{n-d}\partial_r \xi_\eta(x)\partial_j c_\eta(x) =
\delta_{rj} - \sum_{\eta,\alpha=1}^{n-d}\Big(\partial_r
\xi_\eta (\nabla \xi^\top\nabla\xi)^{-1}_{\eta\alpha}
\partial_j \xi_\alpha\Big)(x) = P_{rj}(x)\,, \quad 1\le r,j\le n\,.
\label{1st-derivative-constraint-eqn-p}
\end{equation}
Next, we show the second identity in Eq.~\eqref{lagrange-multiplier-diff}. Differentiating Eq.~\eqref{1st-derivative-constraint-eqn} again, setting $x'=x$
and using Eq.~\eqref{1st-derivative-constraint-eqn-p}, we get, for $1\le \alpha \le n-d$ and $1\le j,l\le n$,
\begin{equation*}
\begin{aligned}
0 =& \sum_{r,r'=1}^n \partial^2_{rr'} \xi_\alpha(x) \Big(\delta_{rj} +
\sum_{\eta=1}^{n-d} \partial_r \xi_\eta(x)\partial_j c_\eta(x)\Big) \Big(\delta_{r'l} +
\sum_{\eta=1}^{n-d} \partial_{r'} \xi_\eta(x)\partial_l
c_\eta(x)\Big) \\
& + \sum_{r=1}^n \partial_r \xi_\alpha(x) \Big(\sum_{\eta=1}^{n-d} \partial_r \xi_\eta(x)\partial^2_{jl} c_\eta(x)\Big) \\
= &\sum_{r,r'=1}^n \Big(\partial^2_{rr'} \xi_\alpha P_{rj} P_{r'l}\Big)(x)
+ \sum_{\eta=1}^{n-d}\Big((\nabla \xi^\top \nabla\xi)_{\alpha\eta}
\partial^2_{jl} c_\eta\Big)(x)\,,
\end{aligned}
\end{equation*}
from which we can solve, for $1\le \eta \le n-d$ and $1\le j,l\le n$,
\begin{equation*}
\begin{aligned}
\partial^2_{jl} c_\eta(x) =&
- \sum_{\alpha=1}^{n-d} \sum_{r,r'=1}^{n} \Big((\nabla\xi^\top\nabla\xi)^{-1}_{\eta\alpha}
\partial^2_{rr'} \xi_\alpha\,
P_{rj}P_{r'l}\Big)(x) \\
=& \sum_{\alpha=1}^{n-d} \sum_{r,r'=1}^{n}
\Big((\nabla\xi^\top\nabla\xi)^{-1}_{\eta\alpha}\, \partial_r
\xi_\alpha \partial_{r'} P_{rj}\,
P_{r'l}\Big)(x)
\,,
\end{aligned}
\end{equation*}
where the second equality follows from the product rule and the identity
$\sum_{r=1}^{n} P_{rj} \partial_r \xi_{\alpha}=0$.
This shows the second identity in Eq.~\eqref{lagrange-multiplier-diff}.
Lastly, we prove the expansion in Eq.~\eqref{constraint-expansion}.
Note that the second identity in Eq.~\eqref{lagrange-multiplier-diff} and Eq.~\eqref{exp-pij} implies
\begin{equation}
\sum_{\eta=1}^{n-d} (\partial_i \xi_\eta\partial^2_{jl} c_{\eta})(x) =
\sum_{r,r'=1}^{n} \big((I-P)_{ir} P_{r'l}\,\partial_{r'} P_{rj}\big)(x) \,, \quad 1 \le i,j,l \le n\,.
\label{2nd-derivative-constraint-eqn-p}
\end{equation}
By expanding $c(x')$ at $x'=x$ to the third order, noticing that $c(x)=0$,
and using Eq.~\eqref{1st-derivative-constraint-eqn-p} and Eq.~\eqref{2nd-derivative-constraint-eqn-p}
for the first and second order derivatives respectively,
we can derive
\begin{equation*}
\begin{aligned}
y_i =& x'_i + \sum_{\eta=1}^{n-d} \partial_i \xi_\eta(x) c_\eta(x') \\
=& x_i + (x'_i-x_i) + \sum_{\eta=1}^{n-d} \partial_i \xi_\eta(x)
\bigg[\sum_{j=1}^{n} \partial_j c_{\eta}(x) (x'_j - x_j) +
\frac{1}{2} \sum_{j,l=1}^{n} \partial^2_{jl} c_{\eta}(x) (x'_j -
x_j)(x'_{l} - x_{l}) \\
& + \frac{1}{6}\sum_{j,l,r=1}^{n} \partial^3_{jlr} c_\eta(x)\, (x'_j -
x_j)(x'_{l} - x_{l})(x'_{r} - x_{r})\bigg] + o(|x'-x|^3)\\
=& x_i + \sum_{j=1}^{n} P_{ij}(x) (x'_j - x_j) + \frac{1}{2}
\sum_{j,l=1}^{n} \bigg[\sum_{r,r'=1}^{n} \Big((I-P)_{ir} \partial_{r'}
P_{rj} P_{r'l}\Big)(x)\bigg] (x'_j - x_j)(x'_{l} - x_{l}) \\
& + \frac{1}{6} \sum_{j,l,r=1}^{n}
\bigg(\sum_{\eta=1}^{n-d}
\partial_i \xi_\eta(x)
\partial^3_{jlr} c_\eta(x)\bigg) (x'_j -
x_j)(x'_{l} - x_{l})(x'_{r} - x_{r}) + o(|x'-x|^3)\,,
\end{aligned}
\end{equation*}
which proves Eq.~\eqref{constraint-expansion}.
\end{proof}
\section{Details of algorithms and experiments}\label{Exp_details}
We present the algorithms for sampling the forward process and for solving constraint equations in Algorithms~\ref{algo_forward} and \ref{Newton_solver}, respectively. In the following, we discuss several further algorithmic details (see Section~\ref{subsec-algo-details}) that are common in our experiments. Specific details of each experiment are discussed in the subsections below.
\paragraph{Neural networks and training setup} As described in Theorem \ref{thm-continuous-limit}, the functions $(s^{(k+1),\theta}(x))_{0\le k\le N-1}$ are represented by a single function $s_{\theta}(x, t)$ with parameter $\theta$, which is in turn modeled by a multilayer perceptron (MLP). We employ SiLU as the activation function. We do not require that the output of the neural network belongs to the tangent space, thanks to the presence of the projection in both the forward and the reverse processes. Alternative strategies for designing neural networks with outputs in tangent space are proposed in \citet{de2022riemannian}.
We train our models using PyTorch, where we employ the Adam optimizer with fixed learning rate $r=5\times 10^{-4}$ and we clip the gradients of the parameters when the $2$-norm exceeds $10.0$. We also implement an exponential moving average for the model weights \citep{polyak1992acceleration} with a decay rate of $0.999$.
All experiments are run on a single NVIDIA A40 GPU with 48G memory.
\begin{algorithm}[htbp]
\caption{Sampling trajectory of forward process}\label{algo_forward}
\begin{algorithmic}[1]
\STATE \textbf{Input: } $x^{(0)}\in\mathcal{M}$, constants $\sigma_{k}$, function $b: \mathbb{R}^n\rightarrow\mathbb{R}^n$, and integer $N$
\FOR {$k = 0$ to $N-1$}
\STATE generate $z^{(k)}\sim \mathcal{N}(0,I_n)$ and set $v^{k}= P(x^{(k)})z^{(k)}$
\STATE set $x^{(k+\frac{1}{2})}:= x^{(k)} + \sigma_k^2 b(x^{(k)})+ \sigma_k v^{(k)}$
\STATE \MYCOMMENT{solve~(\ref{projection-kstep-forward}) by Algorithm~\ref{Newton_solver}}
\STATE $c, \mathrm{flag} = \text{newton\_solver}(x^{(k)},x^{(k+\frac{1}{2})}; \xi)$.
\IF {flag == true}
\STATE set $x^{(k+1)} := x^{(k+\frac{1}{2})} + \nabla\xi(x^{(k)})c$
\ELSE
\STATE discard the trajectory and re-generate %set $x^{(k+1)} := x^{(k)}$
\ENDIF
\ENDFOR
\STATE \textbf{Return} $(x^{(0)},x^{(1)},\dots,x^{(N)})$
\end{algorithmic}
\label{algo-reverse}
\end{algorithm}
\begin{algorithm}[htbp]
\caption{newton\_solver($x,x';\xi$) \hfill{\footnotesize\textcolor{gray}{// solve $\xi(x'+\nabla\xi(x)c)=0$ by Newton's method}}}\label{Newton_solver}
\begin{algorithmic}[1]
\STATE \textbf{Input: } $x\in \mathcal{M}$, $x'\in\mathbb{R}^n$, $\xi: \mathbb{R}^n\rightarrow \mathbb{R}^{n-d}$, maximal iteration steps $n_{\mathrm{step}}$, tolerance $\mathrm{tol}>0$
\STATE \textbf{Initialization}: set $c=0\in\mathbb{R}^{n-d}$
\FOR {$k = 0$ to $n_{\mathrm{step}}-1$}
\STATE Solve linear equation $\big[\nabla \xi\big(x'+\nabla\xi(x)c\big)^\top\nabla\xi(x)\big] u= -\xi\big(x'+\nabla \xi (x)c\big)$ for $u\in\mathbb{R}^{n-d}$
\STATE $c \leftarrow c+u$
\IF {$|\xi(x'+\nabla\xi(x) c)|<\mathrm{tol}$}
\STATE \textbf{Return} c, true
\ENDIF
\ENDFOR
\STATE \textbf{Return} c, false
\end{algorithmic}
\end{algorithm}
\paragraph{NLL calculation}
Following \citet{de2022riemannian} and \citet{chen2024flow},
for each experiment we train our model in five different runs with random seed values ranging from $0$ to $4$. In each run, the dataset is divided into training, validation, and test sets with ratio $80$:$10$:$10$. We compute the test NLL (i.e.\ NLL on test set) using the model that yields the best validation NLL during the training. On both validation set and test set, the NLL is calculated by (see the second line in (\ref{variational-bound}))
\begin{equation}\label{nll}
-\log p_\theta(x^{(0)})
= -\log \mathbb{E}_{\mathbb{Q}^{(N)}} \bigg[ p(x^{(N)})\prod_{k=0}^{N-1} \frac{p_\theta(x^{(k)}\,|\,x^{(k+1)})}
{q(x^{(k+1)}\,|\,x^{(k)})} \bigg| x^{(0)} \bigg],
\end{equation}
where the conditional expectation is estimated by sampling multiple ($50$ or $100$) paths of the forward process starting from each $x^{(0)}$, and we replace both $\epsilon_{x^{(k)}}^{(\sigma_k)}$ in Eq.~\eqref{joint-forward-projection} and $\epsilon_{x^{(k+1)},_\theta}^{(\beta_{k+1})}$ in Eq.~\eqref{joint-backward-projection} by zero (which is valid since $\sigma_k$ and $\beta_{k+1}$ are small in our experiment).
We use Eq.~\eqref{nll} to compute NLLs in the experiments where the prior distribution $p(x^{(N)})$ is uniform distribution.
\paragraph{Model parameters} As in Theorem \ref{thm-continuous-limit}, we choose
$T>0$, integer $N>0$, function $g(t) = \gamma_{\mathrm{min}} + \frac{t}{T}(\gamma_{\mathrm{max}} - \gamma_{\mathrm{min}})$ for some $\gamma_{\mathrm{max}} \ge \gamma_{\mathrm{min}}>0$, and parameters $\sigma_k = \beta_{k+1} = \sqrt{h} g(kh)$, where $h=\frac{T}{N}$ and $k=0,1,\dots, N-1$.
We choose $b=0$ in all experiments except the alanine dipeptide.
\paragraph{Generation of standard Gaussian variables in $T_x\mathcal{M}$.}
Let $z$ be a standard Gaussian random variable in $\mathbb{R}^n$. It is
straightforward to verify that $v=P(x)z$ is a standard Gaussian random
variable in $T_x\mathcal{M}$, where $P(x)$ is the orthogonal projection matrix to the tangent space $T_x\mathcal{M}$.
We use this fact to generate tangent vectors in the forward and reverse
processes (see line $3$ of Algorithm~\ref{algo_forward} and line 4--5 of Algorithm~\ref{algo_generation}, respectively).
Values of all the parameters in our experiments are summarized in Table~\ref{Hyperparameters}.
\begin{table}[htbp]
\centering
\caption{Parameters in our experiments. $\gamma_{\mathrm{min}},\gamma_{\mathrm{max}},N,T$ are the parameters in our model; $l_{\mathrm{f}},N_{\mathrm{epoch}},B$ are the parameters in Algorithm \ref{algo_training}; $N_{\mathrm{node}},N_{\mathrm{layer}}$ are the numbers of the hidden nodes per layer and the hidden layers of the neural networks, respectively.}
\label{Hyperparameters}
\begin{tabular}{lccccccccc}
\toprule
Datasets & $\gamma_{\mathrm{min}}$ & $\gamma_{\mathrm{max}} $& $N$ & $T$ &$l_{\mathrm{f}}$ & $N_{\mathrm{epoch}}$& $B$ & $N_{\mathrm{node}}$& $N_{\mathrm{layer}}$ \\ \midrule
Volcano &$ 0.01$& $1.0$& $400$&$4.0$ & $1$ & $20000$ & $128$ & $512$ & $5$\\
Earthquake &$ 0.01$& $1.0$& $400$&$4.0$ &$1$ &$20000$ & $512$& $512$ & $5$\\
Flood &$ 0.01$& $1.0$& $400$&$4.0$ & $1$& $20000$ & $512$ & $512$ & $5$\\
Fire &$ 0.01$& $1.0$& $400$&$4.0$ & $1$ & $20000$ & $512$ & $512$ & $5$ \\
\midrule
Bunny, $k=50$& $0.07$& $0.07$& $800$& $8.0$&$100$ & $2000$ &$2048$ & $256$ & $5$ \\
Bunny, $k=100$& $0.07$ & $0.07$& $500$&$5.0$ & $100$& $2000$ &$2048$ & $256$ & $5$ \\
Spot, $k=50$ & $0.1$&$0.1$ & $500$ & $5.0$ &$100$ & $2000$ &$2048$ & $256$ & $5$ \\
Spot, $k=100$ & $0.1$& $0.1$& $300$& $3.0$ &$100$ & $2000$ &$2048$ & $256$ & $5$ \\
\midrule
$\mathrm{SO}(10)$,\,$m=3$ & $0.2$ & $2.0$& $500$& $1.0$ & $100$& $2000$& $512$ & $512$ & $3$\\
$\mathrm{SO}(10)$,\,$m=5$ & $0.2$ & $2.0$& $500$& $1.0$ & $100$ &$2000$ & $512$ & $512$ & $3$ \\
\midrule
Alanine dipeptide & $1.0$ & $1.0$ & $200$&$0.1$ & $100$& $5000$ & $512$ & $512$ & $5$ \\
\bottomrule
\end{tabular}
\end{table}
\subsection{Earth and climate science datasets}
\label{app-subsec-earth}
The unit sphere is viewed as a submanifold of $\mathbb{R}^3$, which is defined by the zero level set of the function $\xi(x)=|x|-1$, for $x\in \mathbb{R}^3$. Newton's method is not necessary for the projection steps (\ref{projection-kstep-forward}) and (\ref{projection-kstep-backward}), since the general scheme (\ref{projection}) has the closed-form solution:
\begin{equation}\label{proj_scheme_S2}
y= \left(1-|\sigma^2 P(x)b(x)+\sigma v|^2 \right)^{\frac{1}{2}} x + \sigma^2 P(x)b(x)+\sigma v,
\end{equation}
as long as $|\sigma^2 P(x)b(x)+\sigma v| < 1$.
In the following, we present a careful study on the dataset splitting and the NLL calculation for these datasets. In fact, the datasets in this example contain isolated data points and the distributions of the data points are complex.
Consequently, the training, validation, and test sets obtained from
the standard data splitting method exhibit significant differences
in their distributions, leading to inaccurate evaluation results and an increased risk of overfitting. Specifically, the presence of isolated data points in the validation set or in the test set (therefore not in the training set) can significantly affect the final test results. Figure \ref{earth:data_all}
illustrates this phenomenon
in the case of the volcano dataset, where the model is not trained on points in the vicinity of the three (one) isolated points in the test (validation) set, since there is no training data in that area. Accordingly, the loss and NLL values computed on different sets are significantly different, as shown by the solid lines in Figures~\ref{earth:nll_epoch}--\ref{best_nll_seed}.
To resolve this issue, we propose to include the
isolated points in the validation set and the test set to the training set. These points can be identified by
binning the data points according to their latitude and longitude values. Figure \ref{earth:nll_epoch} shows that the NLL values computed on different sets become closer within a single run when the isolated points are included in the training set, and Figure \ref{best_nll_seed} shows that the NLL values become smaller (better) among five different runs.
\begin{figure}[htbp]
\centering
\includegraphics[width=1.0\linewidth]{figs/earth_likelihood_all.pdf}
\caption{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.}\label{earth_result}
\end{figure}
% height=3.3cm
% \begin{figure}[htbp]
\begin{figure}[ht]
\centering
\begin{subfigure}[b]{0.98\textwidth}
\centering
\includegraphics[width=0.95\linewidth]{figs/earth_dataset_volcano_4_new.pdf}
\caption{Volcano dataset \label{earth:data_all}}
\end{subfigure}\\
\begin{subfigure}[b]{0.48\textwidth}
\centering
\includegraphics[width=0.95\linewidth]{figs/earth_nll_epoch.pdf}
\caption{NLL vs epoch \label{earth:nll_epoch}}
\end{subfigure}
\begin{subfigure}[b]{0.48\textwidth}
\centering
\includegraphics[width=0.95\linewidth]{figs/earth_nll_seed.pdf}
\caption{Best NLL \label{best_nll_seed}}
\end{subfigure}
\caption{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$.
\label{earth_unseen}}
\end{figure}
\subsection{Mesh data on learned manifolds}\label{Bunny_Spot_details}
To create the datasets, we adopt the approach described by \citet{jo2024generativemodelingmanifoldsmixture} and \citet{chen2024flow}. The data is generated according to the density defined by the $k$-th clamped eigenfunction of the Laplacian operator on a mesh that has been upsampled threefold.
The function $\xi:\mathbb{R}^3\rightarrow\mathbb{R}$ is modeled by a MLP with $3$ hidden layers, each of which has $128$ nodes. Different from the activation function in our model, here we use Softplus activation function, where the parameter $\beta$ is set to $10$. The loss function for learning $\xi$ is
\begin{equation}\label{sdf_loss}
\ell(\xi)=\frac{1}{|\mathcal{D}|} \sum_{x\in \mathcal{D}}|\xi(x)|+ \frac{\lambda}{|\mathcal{D}'|}\sum_{y \in \mathcal{D}'}(|\nabla \xi(y)|-1)^2,
\end{equation}
where $\lambda = 0.1$, $\mathcal{D}$ denotes the set of vertices of a high-resolution mesh, and the set $\mathcal{D}'$ contains samples near the manifolds that are obtained by perturbing samples $x\in \mathcal{D}$ according to $y=x+c\epsilon$, with $\epsilon\sim \mathcal{N}(0,~I_3)$ and $c=0.05$.
The first term in Eq.~\eqref{sdf_loss} imposes that $\xi$ is close to zero on vertices, whereas the second term serves as a regularization term and ensures that $\xi$ has non-vanishing gradient near the manifold.
The neural network is trained for $200000$ steps using Adam optimizer, with batch size $512$ and learning rate $10^{-4}$.
With the learned function $\xi$, we consider the manifold defined by
$\mathcal{M}=\{x\in \mathbb{R}^3 | \xi(x)=0\}$. The values of $\xi$ on the dataset are at the order $10^{-2}$. To ensure that the data is on $\mathcal{M}$ with high precision, we refine the dataset by solving the following ordinary differential equation (ODE):
\begin{equation}\label{refine_ode}
\frac{d x_t}{dt}=-\xi(x_t) \nabla\xi(x_t), \quad t\ge 0\,,
\end{equation}
starting from each point in the dataset until the condition $|\xi(x_t)|<10^{-5}$ is reached (notice that
Eq.~\eqref{refine_ode} is a gradient flow and $\lim_{t\rightarrow \infty} |\xi(x_t)|=0$). This ensures that the refined points conform to the manifold accurately.
For the Newton's method in generating paths, we set $\mathrm{tol}=10^{-4}$ and $n_{\mathrm{step}}=10$ in Algorithm~\ref{Newton_solver}. Due to the
complex geometry of the objects, there is a small portion of paths (less than $1.0\%$) that can not be successfully generated. Apart from these paths, the Newton's method reaches convergence within $3$ iteration steps.
\subsection{High-dimensional special orthogonal group}\label{SOn_details}
We view the group $\mathrm{SO}(10)$ as a $45$-dimensional submanifold of $\mathbb{R}^{100}$ that corresponds to (a connected component of) the zero level set of the map $\xi: \mathbb{R}^{100} \rightarrow \mathbb{R}^{55}$, whose components consist of the upper triangle portion of the matrix $S^\top S-I_{10}$, where $S$ is a $10\times 10$ matrix.
The dataset is constructed as a mixture of $m$ wrapped normal distributions, each of which is the image (under the exponential map) of a normal distribution in the tangent space of a center $S_i\in \mathrm{SO}(10)$, $1\le i\le m$.
To ensure multimodality, we define the centers $S_i$ as follows. We initially define a $2\times 2$ matrix $A_0 := \begin{bmatrix}\cos\frac{\pi}{3} & \sin\frac{\pi}{3} \\-\sin\frac{\pi}{3} & \cos\frac{\pi}{3}\end{bmatrix}$, which represents a rotation by $\frac{\pi}{3}$ radians. We then construct block diagonal matrices of order $10$ by incorporating $A_0$ and the identity matrix $I_2$ in various combinations:
\vspace{0.2cm}
\begin{equation}\label{SO10_center}
\begin{gathered}
X_1=\operatorname{diag}\{A_0, I_2, I_2, I_2, I_2\}, ~~
X_2 =\operatorname{diag}\{A_0, A_0, I_2, I_2, I_2\}, ~~
X_3=\operatorname{diag}\{A_0, A_0, A_0, I_2, I_2\}, \\
X_4=\operatorname{diag}\{A_0, A_0, A_0, A_0, I_2\}, ~~
X_5=\operatorname{diag}\{A_0, A_0, A_0, A_0, A_0\}.\\
\end{gathered}
\end{equation}
The centers $S_i$ of the $m$ wrapped normal distributions are chosen as $S_i=Q_i^\top X_i Q_i$, where $Q_i \in \mathrm{SO}(10)$ are randomly drawn from the uniform distribution. According to Eq.~\eqref{SO10_center}, the statistics $\eta(S)=(\operatorname{tr}(S), \operatorname{tr}(S^2), \operatorname{tr}(S^4),\operatorname{tr}(S^5))$ of the centers can be explicitly computed (using the trace identities $\operatorname{tr}(AB)=\operatorname{tr}(BA)$ and $\operatorname{tr}(Q_i^\top X_iQ_i)=\operatorname{tr}(X_i)$) as
\vspace{0.2cm}
\begin{equation}\label{SO10_center_eta}
\begin{gathered}
\eta(S_1)=(9,7,7,9), \quad
\eta(S_2)=(8,4,4,8), \quad
\eta(S_3)=(7,1,1,7), \\
\eta(S_4)=(6,-2,-2,6), \quad
\eta(S_5)=(5,-5,-5,5).
\end{gathered}
\end{equation}
To generate data in the dataset, we select a center $S_i$ with equal probability, sample tangent vectors $Y$ from the normal distribution (in the tangent space at $S_i$) with zero mean and standard deviation $0.05$, and then compute their images $S$ under the exponential map, that is, $S=S_i\mathrm{e}^{S_i^\top Y}$.
Due to our choice of $A_0$, the distribution of $\operatorname{tr}(S^3)$ is narrowly concentrated at its local peaks. To see this,
using $S_i^3=Q_i^\top X_i^3 Q_i$ and the fact that
$S_i^\top Y$ is anti-symmetric,
we can compute
\begin{equation}
\begin{aligned}
\operatorname{tr}(S^3) =& \operatorname{tr}(S_i\mathrm{e}^{S_i^\top Y}S_i\mathrm{e}^{S_i^\top Y}S_i\mathrm{e}^{S_i^\top Y}) \\
=& \operatorname{tr}\left(S_i^3 + 3(S_i^3) S_i^\top Y + O(|Y|^2)\right) \\
= &
\operatorname{tr}(S_i^3)
+ 3\operatorname{tr}\left(
X_i^3 Q_i S_i^\top YQ_i^\top\right)
+ O(|Y|^2)\\
=& \operatorname{tr}(S_i^3) + O(|Y|^2)\,,
\end{aligned}
\end{equation}
where we have used the trace identity $\operatorname{tr}(AB)=\operatorname{tr}(BA)$, the expansion of matrix exponential, and the fact that the diagonal elements of $X_i^3 Q_i S_i^\top YQ_i^\top$ are all zero (hence its trace equals zero), since $X_i^3$ is a diagonal matrix and $Q_i S_i^\top YQ_i^\top$ is anti-symmetric. Therefore, we omit $\operatorname{tr}(S^3)$ in our choice of statistical analysis for clear presentation.
For the Newton's method, we set $\mathrm{tol}=10^{-6}$ and $n_{\mathrm{step}}=10$ in Algorithm~\ref{Newton_solver}. For this example, the convergence is always reached within $3$ iteration steps.
We also examine the dataset with $m=3$ modes, in which case we use
centers defined as $S_i=Q_i^{\top}X_iQ_i$ for $i=1,2,3$ (the random matrices $Q_i$ are different from those in the case of $m=5$).
Figure \ref{fig:SOn_result_3} shows that the distributions of the learned reverse processes match the corresponding distributions of the forward process at different Markov chain jump steps.
\begin{figure}[htbp]
\centering
\includegraphics[width=0.95\linewidth]{figs/SOn_Statistics_M=3_smooth.pdf}
\caption{Results for $\mathrm{SO}(10)$ with $m=3$.
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. Colors black, red, green, and blue correspond to steps $k = 0, 50, 200, 500$, respectively.}\label{fig:SOn_result_3}
\end{figure}
\subsection{Alanine dipeptide}\label{subsec-dipeptide_details}
To generate the dataset, we initially perform a constrained molecular simulation of alanine dipeptide in water for $1$ns using the molecular dynamics package GROMACS~\citep{GROMACS} with step-size $1$fs. We apply the harmonic biasing method in COLVARS module~\citep{Fiorin2013}, where the collective variable is chosen as the dihedral angle $\phi$ and the harmonic potential is centered at $\phi=-70^\circ$ with the force constant $5.0$. Further simulation details are omitted since they are similar to those in~\cite{ae2024}. In total, $10^4$ configurations are obtained by recording every $100$ simulation steps. We exclude the hydrogen atoms and work with the coordinates of the $10$ non-hydrogen atoms in the system (see Figure~\ref{dipeptide_system}). In a final preparatory step, we apply the refinement technique in Appendix~\ref{Bunny_Spot_details} (see Eq.~\eqref{refine_ode}) to the recorded coordinates, so that the data in the dataset lives in the manifold
$\mathcal{M}=\{x\in \mathbb{R}^{30}| \phi(x)=-70^\circ\}$ up to a small numerical error of order $10^{-5}$.
Since $\mathcal{M}$ is unbounded, we
adopt a nonzero function $b$ in our model to make sure that the Markov chain processes stay in bounded region. To this end, we choose a reference configuration
$x^\mathrm{ref}$ from the dataset and define the potential function
\begin{equation}
V(x) = \frac{\kappa}{2} |R_x^*(x-w^*_x) - x^{\mathrm{ref}}|^2\,, \quad x\in \mathbb{R}^{30}\,,
\label{potential-v}
\end{equation}
with $\kappa = 50$, where $R_x^*$, $w^*_x$ are the optimal rotation and the optimal translation that minimize the RMSD (see Eq.~\eqref{rmsd}). The function $b$ is defined as (the negative gradient of $V$ in full space)
\begin{equation}
b=-\nabla V(x) = -\kappa \left(R_x^*(x-w^*_x) - x^{\mathrm{ref}}\right),
\end{equation}
where the second equality follows by differentiating $V$ in Eq.~\eqref{potential-v} and using the first order optimality equations satisfied by $R_x^*$ and $w^*_x$~(also see \citet{rmsd_quaternions}).
We also build our model to make sure that the generated distribution is
$\mathrm{SE}(3)$-invariant (i.e. invariant under rotations and translations).
For this, we rely on the theoretical results in \citet{xu2022geodiff} and in Appendix \ref{Equivariant_NN}.
One can check that $V(x)$ is $\mathrm{SE}(3)$-invariant and $b$
satisfies property (\ref{SE3_sym}) in Appendix~\ref{Equivariant_NN}, that is, $b$ is equivariant under rotations and invariant under translations. This guarantees that the prior distribution $p(x^{(N)})$, which we choose as the invariant distribution of the forward process, is $\mathrm{SE}(3)$-invariant as well.
We still need to make sure that the transition densities of the reverse Markov chain are $\mathrm{SE}(3)$-invariant. For this purpose, in the reverse process we set $s^{(k+1),\theta}(x) = (R^*_x)^\top f_\theta(R^*_x(x-w^*_x),\frac{(k+1)T}{N})$, where $f_{\theta}:\mathbb{R}^{30}\times \mathbb{R}\rightarrow\mathbb{R}^{30}$ is modeled by a single MLP with parameter $\theta$, and both $R^*_x$ and $b^*_x$ are computed by the Kabsch algorithm \citep{kabsch1976solution}.
With this choice,
$s^{(k+1),\theta}(x)$ satisfies property (\ref{SE3_sym}) by
Proposition \ref{equiv_nn} in Appendix~\ref{Equivariant_NN}, and the transition density of the reverse process is $\mathrm{SE}(3)$-invariant by
Proposition \ref{equiv_transition} in Appendix~\ref{Equivariant_NN}. Since the prior $p(x^{(N)})$ is also $\mathrm{SE}(3)$-invariant, we conclude that the learned distribution $p_{\theta}(x^{(0)})$ is $\mathrm{SE}(3)$-invariant~\citep{xu2022geodiff}.
Compared to the commonly used equivariant networks~\citep{satorras2021n}, our network fits our experiment better thanks to its lower computational cost and reduced memory usage.
For the Newton's method, we set $\mathrm{tol}=10^{-5}$ and $n_{\mathrm{step}}=10$ in Algorithm~\ref{Newton_solver}.
For all but one point, the convergence is reached within $2$ steps.
\section{Theoretical results on neural networks for molecular systems}\label{Equivariant_NN}
In this section, we present theoretical results for the neural network architecture we employed in studying alanine dipeptide.
Assume that the system consists of $M$ atoms, where the coordinates of thee $i$-th atom are denoted by $\boldsymbol{x}_i\in \mathbb{R}^{3}$, for $i=1,2,\dots, M$. Let $x\in \mathbb{R}^{3M}$ be the vector consisting of all the coordinates
$\boldsymbol{x}_1\,\boldsymbol{x}_2,\dots, \boldsymbol{x}_M\in \mathbb{R}^{3}$.
For simplicity, given a rotation matrix $R\in \mathrm{SO}(3)$ and a translation vector $w\in\mathbb{R}^3$, we use the conventional notation $Rx+w$ to denote the vector in $\mathbb{R}^{3M}$ that consists of the transformed coordinates $R\boldsymbol{x}_1+w, R\boldsymbol{x}_2+w, \dots, R\boldsymbol{x}_M+w\in \mathbb{R}^3$. We say that a function $f$ defined in $\mathbb{R}^{3M}$ is $\mathrm{SE}(3)$-invariant, if $f(Rx+w) = f(x)$, for all $R\in \mathrm{SO}(3)$, $w\in\mathbb{R}^3$, and for all $x\in \mathbb{R}^{3M}$. We say that function $f:\mathbb{R}^{3M}\rightarrow \mathbb{R}^{3M}$ possesses property (\ref{SE3_sym}), if it is both equivariant under rotations and invariant under translations, i.e.\
\begin{equation}\label{SE3_sym}
f(Rx+w) = Rf(x),\mbox{ for all }R\in \mathrm{SO}(3), w\in\mathbb{R}^3, \mbox{and all }x\in \mathbb{R}^{3M}. \tag{\#}
\end{equation}
Assume that a configuration $x^\textrm{ref}$ is chosen as reference. Given $x$, the optimal rotation matrix and the optimal translation vector, which minimize the RMSD
\begin{equation}
\mathrm{RMSD}(x;x^\textrm{ref}) = \Big(\frac{1}{M}|R (x - w) - x^\textrm{ref}|^2\Big)^{\frac{1}{2}}
\label{rmsd}
\end{equation}
from the reference $x^\textrm{ref}$, are denoted by $R^*_x$ and $w^*_x$, respectively.
The following result characterizes functions that are both equivariant under rotations and invariant under translations.
\begin{proposition}\label{equiv_nn}
The following two claims are equivalent.
\begin{itemize}
\item
Function $s:\mathbb{R}^{3M}\rightarrow\mathbb{R}^{3M}$ possesses property (\ref{SE3_sym}).
\item
There is a function $f:\mathbb{R}^{3M}\rightarrow\mathbb{R}^{3M}$, such that
$s(x)
= (R^*_x)^\top f(R^*_x(x-w^*_x))$, for all $x\in \mathbb{R}^{3M}$.
\end{itemize}
\label{prop-equivariant}
\end{proposition}
\begin{proof}
It is straightforward to verify that the first claim implies the second claim. In fact, setting $R=R_x^*$, $w=-R^*_xw_x^*$, and using the identity $R^\top R=I_3$, we obtain from the first claim that
$s(x)=(R_x^*)^\top s(R^*_x(x-w^*_x))$. Hence, the second claim holds with $f=s$. To show that the second claim also implies the first one, we use the fact that the optimal rotation $R_{Rx+w}^*$ and the optimal translation $w_{Rx+w}^*$, which minimize the RMSD of the state $Rx+w$ from the reference $x^{\mathrm{ref}}$, are given by
$R_{Rx+w}^*=R_{x}^*R^\top$ and $w_{Rx+w}^*=Rw^*_x+w$, respectively. This fact can be directly checked using Eq.~\eqref{rmsd}. In particular, we have
\begin{equation*}
R^*_{Rx+w} (Rx+w - w^*_{Rx+w}) =
R^*_x(x - w^*_x)\,.
\end{equation*}
Therefore, for the function $s$ defined in the second claim, we can compute, for any $R\in \mathrm{SO}(3)$, $w\in\mathbb{R}^3$, and any $x\in \mathbb{R}^{3M}$,
\begin{equation*}
\begin{aligned}
s(Rx + w)=&(R^*_{Rx+w})^\top f\big(R^*_{Rx+w}(Rx+w - w^*_{Rx+w})\big) \\
=& R (R^*_x)^\top f(R^*_x (x -w^*_x))\\
=& R s(x)\,,
\end{aligned}
\end{equation*}
which shows the first claim.
\end{proof}
The following result guarantees the $\mathrm{SE}(3)$-invariance of the transition densities of our diffusion model.
\begin{proposition}\label{equiv_transition}
Assume that $\xi$ is $\mathrm{SE}(3)$-invariant and $b$ possesses property (\ref{SE3_sym}). Then, the transition density of the forward process in Eq.~\eqref{joint-forward-projection} is $\mathrm{SE}(3)$-invariant. Further assume that the function $s^{(k+1),\theta}$
possesses property (\ref{SE3_sym}) for $0 \le k\le N-1$. Then, the transition density of the reverse process in Eq.~\eqref{joint-backward-projection} is also $\mathrm{SE}(3)$-invariant.
\end{proposition}
\begin{proof}
We consider the transition density in Eq.~\eqref{joint-forward-projection}. Recall that
$U_x\in \mathbb{R}^{n\times d}$ is a matrix whose columns form an orthonormal basis of $T_x\mathcal{M}$.
Since $\xi$ is $\mathrm{SE}(3)$-invariant, we have $\xi(Rx+w)=\xi(x)$, for all rotations $R$ and translation vectors $w$, which implies that $Rx+w\in \mathcal{M}$, if and only if $x \in\mathcal{M}$.
Differentiating the identity $\xi(Rx+w)=\xi(x)$ with respect to $x$, we obtain the relation
$\nabla\xi(Rx+w) = R\nabla\xi(x)$, from which we see that $U_{Rx+w}$ can be chosen such that $U_{Rx+w}=RU_{x}$. For the orthogonal projection matrix $P$ in Eq.~\eqref{exp-pij}, using the identity $R^\top R = I_3$, we can compute
\begin{equation*}
\begin{aligned}
P(Rx+w)=&I_n -
\nabla\xi(Rx+w) \big(\nabla\xi(Rx+w)^\top\nabla\xi(Rx+w)\big)^{-1} \nabla\xi(Rx+w)^\top \\
=& I_n -
R\nabla\xi(x) \big(\nabla\xi(x)^\top\nabla\xi(x)\big)^{-1} \nabla\xi(x)^\top R^\top \\
=& R P(x) R^\top\,.
\end{aligned}
\end{equation*}
Moreover, since both $b$ and $\nabla \xi$ satisfy the property (\ref{SE3_sym}), we also have $\epsilon_{Rx^{(k)}+w}^{(\sigma_k)}= \epsilon_{x^{(k)}}^{(\sigma_k)}$ (i.e.\ the probabilities of having no solution are the same). Therefore, for the transition density in Eq.~\eqref{joint-forward-projection}, we can derive, for any $R\in \mathrm{SO}(3)$ and $w\in\mathbb{R}^3$,
\begin{equation*}
\begin{aligned}
&q(Rx^{(k+1)}+w\,|\,Rx^{(k)}+w) \\
=& (2\pi\sigma_k^2)^{-\frac{d}{2}} \big(1-\epsilon_{Rx^{(k)}+w}^{(\sigma_k)}\big)^{-1}
|\det (U_{Rx^{(k)}+w}^\top U_{Rx^{(k+1)}+w})|
\mathrm{e}^{-\frac{\big|P(Rx^{(k)}+b)\big(Rx^{(k+1)}-Rx^{(k)}-\sigma_k^2 b(Rx^{(k)}+w)\big)\big|^2}{2\sigma_k^2}}
\\
=& (2\pi\sigma_k^2)^{-\frac{d}{2}} \big(1-\epsilon_{x^{(k)}}^{(\sigma_k)}\big)^{-1} \left|\det (U_{x^{(k)}}^\top R^\top R U_{x^{(k+1)}})\right|
\mathrm{e}^{-\frac{\big|RP(x^{(k)})R^\top\big(Rx^{(k+1)}-Rx^{(k)}-\sigma_k^2 Rb(x^{(k)})\big)\big|^2}{2\sigma_k^2}} \\
=& (2\pi\sigma_k^2)^{-\frac{d}{2}} \big(1-\epsilon_{x^{(k)}}^{(\sigma_k)}\big)^{-1} \left|\det (U_{x^{(k)}}^\top U_{x^{(k+1)}})\right| \mathrm{e}^{-\frac{\big|P(x^{(k)})\big(x^{(k+1)}-x^{(k)}-\sigma_k^2 b(x^{(k)})\big)\big|^2}{2\sigma_k^2}} \\
=& q(x^{(k+1)}\,|\,x^{(k)})\,,
\end{aligned}
\end{equation*}
which shows the $\mathrm{SE}(3)$-invariance of the transition density of the forward process. The invariance of the transition density of the reverse process in Eq.~\eqref{joint-backward-projection} can be proved using the same argument, assuming that $s^{(k+1),\theta}$ satisfies the relation $s^{(k+1),\theta}(Rx+w)=Rs^{(k+1),\theta}(x)$.
\end{proof}
\section{Supplementary Experimental Results}
This section presents additional experimental results. In Section \ref{Running_time}, we present the computation time for the experiments in the main text. Section \ref{details_scheme} analyzes the computational complexity of our projection scheme and the non-convergence rate of trajectories. Finally, in Section \ref{Ablation}, we conduct an ablation study on the Flood dataset, investigating the impact of the total number of steps $N$ and the trajectory update frequency (proportional to $l_{\mathrm{f}}^{-1}$).
\subsection{Computation time for the experiments in the main text}\label{Running_time}
In Table \ref{runtime}, we present the simulation time $T_{\mathrm{sim}}$ and training time $T_{\mathrm{train}}$, with the percentages of the total runtime $T_{\mathrm{total}}$ shown in parentheses. For the mesh, $\mathrm{SO}(10)$, and the dipeptide datasets, we update the dataset every 100 epochs (i.e. $l_f = 100$), resulting in the simulation time accounting for less than 2\% of the total runtime in mesh and dipeptide datasets, and the simulation time accountingfor approximately 11\% of the total runtime in the $\mathrm{SO}(10)$ datasets due to the high co-dimension of the manifold. In contrast, for the earth and climate science datasets, the dataset is updated at every epoch (i.e. $l_f = 1$), leading to a higher proportion of simulation time.
\begin{table}[htbp]
\centering
\caption{Detailed runtime metrics in our experiments. We report $T_{\mathrm{sim}}$, $T_{\mathrm{train}}$, and $T_{\mathrm{total}}$ as the time for path generation, time for training, and total runtime, respectively, with the percentages of the total runtime $T_{\mathrm{total}}$ shown in parentheses. The parameter $l_{\mathrm{f}}$ determines the frequency of trajectory updates. The final column $T_{\mathrm{epoch}}$ shows the training time per epoch, calculated as $T_{\mathrm{train}}/N_{\mathrm{epoch}}$. All time metrics are reported in seconds.}
\label{runtime}
\begin{tabular}{lccccccc}
\toprule
Datasets & $l_{\mathrm{f}}$ & $N_{\mathrm{epoch}}$ & $T_{\mathrm{sim}}$ & $T_{\mathrm{train}}$ & $T_{\mathrm{total}}$ & $T_{\mathrm{epoch}}$\\
\midrule
Volcano &1 & 20000& 3946(67.0\%)& ~1943(33.0\%) & ~5889 & 0.10\\
Earthquake &1 & 20000& 4591(24.2\%)& 14395(75.8\%) & 18986 & 0.72\\
Flood &1 & 20000 & 4414(29.0\%) & 10814(71.0\%) & 15228 & 0.54\\
Fire &1 & 20000 & 5666(14.0\%)& 34893(86.0\%) & 40559 & 1.74\\
\midrule
Bunny, $k=50$ & 100 & 2000 & 252(1.8\%) & 14120(98.2\%)& 14372 & 7.06\\
Bunny, $k=100$ & 100 & 2000 & 142(1.6\%) & ~8718(98.4\%) & ~8860 & 4.36\\
Spot, $k=50$ & 100 & 2000 & 113(1.3\%)& ~8751(98.7\%) & ~8864 & 4.38\\
Spot, $k=100$ & 100& 2000 & ~70(1.3\%)& ~5189(98.7\%) & ~5259 & 2.59\\
\midrule
$\mathrm{SO}(10)$,\,$m=3$ & 100 & 2000 & 2414(11.2\%) & 19042(88.8\%) & 21456 & 9.52\\
$\mathrm{SO}(10)$,\,$m=5$ & 100 & 2000 & 2426(11.2\%) & 19289(88.8\%) & 21715 & 9.64\\
\midrule
Alanine dipeptide & 100 & 5000 & 159(1.1\%)& 14299(98.9\%)& 14458 & 2.86\\
\bottomrule
\end{tabular}
\end{table}
\subsection{More details about the projection scheme}\label{details_scheme}
\paragraph{Computational complexity of the projection scheme.}
Let $k_{\mathrm{iter}}$ denote the number of Newton iterations, and $C_{\xi}$ be the computational cost of evaluating $\nabla \xi$. The complexity of solving the linear equations in Algorithm~\ref{Newton_solver} is $\mathcal{O}((n-d)^3)$ at maximum, where $n-d$ is the co-dimension of the manifold. Thus, the total complexity of Newton’s method is $\mathcal{O}(k_{\mathrm{iter}}(C_{\xi} +(n-d)^3))$.
In our examples, we have $C_{\xi}=\mathcal{O}(1)$, except for the mesh datasets, where neural network propagation is required. In most cases, $n-d = 1$, while for the $\mathrm{SO}(10)$ datasets, $n-d = 55$. When Newton’s method converges, the number of iterations $k_{\mathrm{iter}}$ is no more than $3$.
\paragraph{Failure rate of trajectory generation.}
As discussed in Section \ref{subsec-projection-scheme}, Eq.~\eqref{projection} may not have solutions for some vectors $v$. In such case, the corresponding trajectories are discarded.
Table~\ref{Newton} reports the failure rate of trajectory generation in our experiments. For the earth and climate science datasets, the general scheme in Eq.~\eqref{projection} is computed analytically using Eq.~\eqref{proj_scheme_S2}, and no trajectories are discarded with parameters in Table~\ref{Hyperparameters}. For the mesh dataset, due to the complex geometry of the manifolds, a small portion of trajectories (less than 1.0\%) fail. For $\mathrm{SO}(10)$, trajectory generation always succeeds. In the alanine dipeptide experiment, only one out of 8000 trajectories fails.
Furthermore, we investigate the impact of the step-size $\sigma$ on the failure rate of Newton’s method. Specifically, as shown in Table \ref{bunny100_fail}, we conduct experiments on the Bunny dataset with $k=100$, using various values of $N$, where $N$ and $\sigma$ are related by $\sigma \sim 1/\sqrt{N}$. As the total number of steps $N$ decreases (i.e., as the step-size $\sigma$ increases), the proportion of failed trajectories increases.
% Table \ref{Newton} reports the non-convergence rate of trajectories during sampling in our experiments. For the Earth and climate science datasets, the general scheme (\ref{projection}) is computed using Eq.~\eqref{proj_scheme_S2}, and no trajectories are dropped under the settings in Table \ref{Hyperparameters}. For the mesh dataset, due to the complex geometry of the objects, a small portion of trajectories (less than 1.0\%) fail to converge. Besides, Newton’s method achieves convergence within 3 iterations. For $\mathrm{SO}(10)$, convergence is always reached within 3 iterations. In the Alanine dipeptide experiment, only one out of 8000 trajectories fails to converge, while the remaining trajectories reach convergence within 2 iterations.
\begin{table}[htbp]
\centering
\caption{Failure rate of trajectory generation. $R_{\mathrm{fail\_fwd}}$ and $R_{\mathrm{fail\_bwd}}$ represent the percentages of discarded trajectories when sampling the forward and reverse process, respectively. $\sigma_{\max}$ denotes the maximum value of $(\sigma_k)_{0\leq k \leq N-1}$.}
\label{Newton}
\begin{tabular}{lccc}
\toprule
Datasets & $\sigma_{\max}$& $R_{\mathrm{fail\_fwd}}$ & $R_{\mathrm{fail\_bwd}}$ \\
\midrule
Volcano & 0.100 & 0.00\% & 0.00\% \\
Earthquake & 0.100 & 0.00\% & 0.00\% \\
Flood & 0.100 & 0.00\% & 0.00\% \\
Fire & 0.100 & 0.00\% & 0.00\% \\
\midrule
Bunny, $k=50$& 0.007 & 1.00\% & 0.82\%\\
Bunny, $k=100$& 0.007 & 0.65\% & 0.55\%\\
Spot, $k=50$& 0.010 & 0.15\% & 0.25\% \\
Spot, $k=100$& 0.010& 0.11\% & 0.10\% \\
\midrule
$\mathrm{SO}(10)$,\,$m=3$ &0.089 & 0.00\% & 0.00\% \\
$\mathrm{SO}(10)$,\,$m=5$ & 0.089 & 0.00\% & 0.00\% \\
\midrule
Alanine dipeptide & 0.022 & 0.01\% & 0.00\% \\
\bottomrule
\end{tabular}
\end{table}
% \begin{table}[htbp]
% \centering
% \caption{The total number of steps $N$ in the Markov chain for the Bunny dataset with $k=100$. The definitions of $R_{\mathrm{non\_c}}$ if the same as those in Table \ref{Ablation_flood_N}. }
% \label{bunny100_fail}
% \begin{tabular}{lccc}
% \toprule
% $N$ & $\sigma_{\max}$ & $R_{\mathrm{non\_c}}$ \\
% \midrule
% 100 & 1.57e-2 & 3.51\% \\
% 200 & 1.11e-2 & 2.30\% \\
% 300 & 0.90e-2 & 1.49\% \\
% 400 & 0.78e-2 & 0.97\% \\
% 500 & 0.70e-2 & 0.66\% \\
% 600 & 0.64e-2 & 0.46\% \\
% 700 & 0.59e-2 & 0.17\% \\
% \bottomrule
% \end{tabular}
% \end{table}
\begin{table}[htbp]
\centering
\caption{Failure rate of trajectory generation for the Bunny dataset with $k=100$ under various $N$. Here, $\sigma$ denotes the step-size and $R_{\mathrm{fail\_fwd}}$ represents the proportion of failed forward trajectories.}
\label{bunny100_fail}
\begin{tabular}{lccccccc}
\toprule
$N$ & 100 & 200 & 300 & 400 & 500 & 600 & 700 \\
\midrule
$\sigma$ & 1.57e-2 & 1.11e-2 & 0.90e-2 & 0.78e-2 & 0.70e-2 & 0.64e-2 & 0.59e-2 \\
$R_{\mathrm{fail\_fwd}}$ & 3.51\% & 2.30\% & 1.49\% & 0.97\% & 0.66\% & 0.46\% & 0.17\% \\
\bottomrule
\end{tabular}
\end{table}
\subsection{Ablation studies}\label{Ablation}
In Table \ref{Ablation_flood_N}, we investigate the impact of the total number of steps $N$ in the Markov chain on the runtime for the Flood dataset. In Table \ref{Ablation_flood_lf}, we analyze the effect of the trajectory update frequency during training on the Flood dataset.
As shown in Table \ref{Ablation_flood_N}, as the total number of steps $N$ decreases (i.e., as the step-size $\sigma$ increases), the Newton-based projection scheme may fail to find a solution in some cases. Specifically, when $N = 50$ and $N = 100$, the failure rates are 0.55\% and 0.0001\%, respectively. When $N = 200$ and $N = 400$, no failure is observed. Additionally, as $N$ decreases, both the training time and the simulation time reduce accordingly. The best NNL value is achieved when $N=100$. However, in the main text, we point out the issue of NNL value in this example and choose $N=400$ to ensure that all trajectories are successfully generated. Furthermore, as shown in Table \ref{Ablation_flood_lf}, increasing the frequency of trajectory updates (decreasing $l_f$) leads to longer computational time.
\begin{table}[htbp]
\centering
\caption{Ablation study on the total number of steps $N$ in the Markov chain for the Flood dataset. The parameter $R_{\mathrm{fail\_fwd}}$ denotes the percentage of discarded trajectories when simulating the forward process. The definitions of $T_{\mathrm{sim}}$, $T_{\mathrm{train}}$, $T_{\mathrm{total}}$, $T_{\mathrm{epoch}}$ can be found in the caption of Table \ref{runtime}.}
\label{Ablation_flood_N}
\begin{tabular}{lccccccc}
\toprule
$N$ & $\sigma_{\max}$ & $R_{\mathrm{fail\_fwd}}$ & $T_{\mathrm{sim}}$ & $T_{\mathrm{train}}$ & $T_{\mathrm{total}}$ & $T_{\mathrm{epoch}}$ & NLL \\
\midrule
~50 & 0.28 & 0.55\% & ~643(30.8\%) & ~1442(69.2\%) & 2085 & 0.07 & 0.48{\scriptsize $\pm$0.08} \\
100 & 0.20 & 0.00\% & 1151(29.1\%) & ~2801(70.9\%) & 3952 & 0.14 & 0.45{\scriptsize$\pm$0.07} \\
200 & 0.14 & 0.00\% & 2245(29.0\%) & ~5497(71.0\%) & 7742 & 0.27 & 0.46{\scriptsize $\pm$0.07} \\
400 & 0.10 & 0.00\% & 4414(29.0\%) & 10814(71.0\%) & 15228 & 0.54 & 0.49{\scriptsize $\pm$0.09}\\
\bottomrule
\end{tabular}
\end{table}
\begin{table}[htbp]
\centering
\caption{Ablation study on the trajectory update frequency (proportional to $l_{\mathrm{f}}^{-1}$) for the Flood dataset with $N=100$. The definitions of $R_{\mathrm{non\_c}}$, $T_{\mathrm{sim}}$, $T_{\mathrm{train}}$, $T_{\mathrm{total}}$, $T_{\mathrm{epoch}}$ can be found in the caption of Table \ref{runtime}.}
\label{Ablation_flood_lf}
\begin{tabular}{lccccccc}
\toprule
$l_{\mathrm{f}}$ & $T_{\mathrm{sim}}$ & $T_{\mathrm{train}}$ & $T_{\mathrm{total}}$ & $T_{\mathrm{epoch}}$ & NLL \\
\midrule
1 & 1151(29.1\%) & ~2801(70.9\%) & 3952 & 0.14 & 0.45{\scriptsize$\pm$0.07} \\
2 & 556(16.6\%) & 2792(83.4\%)& 3348 & 0.14 & 0.49{\scriptsize$\pm$0.09} \\
3 & 366(11.4\%) & 2838(88.6\%) & 3204 & 0.14 & 0.49{\scriptsize$\pm$0.08} \\
5 & 230(7.6\%) & 2787(92.4\%)& 3017 & 0.14 & 0.48{\scriptsize$\pm$0.09} \\
10 & 111(2.9\%) & 3682(97.1\%) & 3793 & 0.18 & 0.45{\scriptsize $\pm$0.07}\\
\bottomrule
\end{tabular}
\end{table}
% As shown in Table \ref{Ablation_bunny100}, as the total number of steps $N$ decreases (i.e., as the step size $\sigma$ increases), the proportion of non-convergent trajectories increases, while both the training time and the simulation time are correspondingly reduced.
% \begin{table}[htbp]
% \centering
% \caption{Ablation study on the total number of steps $N$ in the Markov chain for the Bunny dataset with $k=100$. The definitions of $R_{\mathrm{non\_c}}$, $T_{\mathrm{sim}}$, $T_{\mathrm{train}}$, $T_{\mathrm{total}}$, $T_{\mathrm{epoch}}$ can be found in the caption of Table \ref{runtime}. }
% \label{Ablation_bunny100}
% \begin{tabular}{lccccccccc}
% \toprule
% $N$ & $R_{\mathrm{non\_c}}$ & $\sigma$ & $T_{\mathrm{sim}}$ & $T_{\mathrm{train}}$ & $T_{\mathrm{total}}$ & $T_{\mathrm{epoch}}$ & NLL \\
% \midrule
% 100 & 3.51\% & 0.0157 & 39(2.41\%) & 1598 & 1637 & 0.80 & 1.25 \\
% 200 & 2.30\% & 0.0111 & 76(1.88\%) & 3980 & 4056 & 1.99 & 1.27 \\
% 300 & 1.49\% & 0.0090 & 111(1.85\%) & 5929 & 6040 & 2.96& 1.29 \\
% 400 & 0.97\% &0.0078 &139(1.95\%) & 6968 & 7107 & 3.48 & 1.30 \\
% 500 & 0.66\% & 0.0070 &150(1.44\%)& 10262 & 10412 &5.13 & 1.31 \\
% 600 & 0.46\% & 0.0064 & 163(1.31\%)& 12322 & 12485 &6.16 & 1.32 \\
% 700 & 0.17\% & 0.0059 & 175(1.20\%)& 14440 & 14615 &7.22 & 1.33 \\
% \bottomrule
% \end{tabular}
% \end{table}
% \subsection{Wasserstein distance}
% We evaluate the distribution on $\mathrm{SO}(10)$ using quantitative metrics. Specifically, we first project $\mathrm{SO}(10)$ onto its Lie algebra. For an orthogonal matrix $A$, its matrix logarithm $\log(A)$ is computed, resulting in an antisymmetric matrix. The upper triangular part of this matrix is then extracted to form a 45-dimensional vector, denoted as $S(A)$. Based on this transformation, we calculate the Wasserstein distance ($W_1$) and Maximum Mean Discrepancy (MMD).
% Table \ref{dist_dist} reports the following distances: the distance between two sets of samples drawn from the dataset (representing the ground truth distribution), the distance between samples from the dataset and those generated by our method, and the distance between samples from the dataset and those generated by a Euclidean diffusion model. Notably, when employing a Euclidean diffusion model \citep{song2021scorebased}, the generated distribution may not lie exactly on the manifold. In such cases, a projection step onto the manifold is applied after diffusion sampling.
% \begin{table}[htbp]
% \centering
% \caption{Quantitative evaluation of generated distributions on $\mathrm{SO}(10)$ using Maximum Mean Discrepancy (MMD) and Wasserstein distances ($W_1$ and $W_2$). Results are presented for two distributions: $\mathrm{SO}(10)$ with $m=3$. Each row represents a specific comparison: (1) the distance between two sets of samples drawn from the dataset (ground truth distribution), (2) the distance between dataset samples and those generated by RDDPM, and (3) the distance between dataset samples and those generated by a Euclidean diffusion model.}
% \label{dist_dist}
% \begin{tabular}{lccccccc}
% \toprule
% & \multicolumn{3}{c}{$\mathrm{SO}(10)$,\,$m=3$} \\
% & MMD & W1 & W2\\
% \midrule
% Dataset-Dataset & 0.82e-4(0.20e-4) & 0.207439(0.010189) & 0.288294(0.035330)\\
% Dataset-RDDPM(ours) & 1.76e-4(0.45e-4) & 0.294814(0.005744) & 0.403184(0.016028)\\
% Dataset-EuropeanDM{\scriptsize \citep{song2021scorebased}} & 5.78e-4(0.44e-4) & 0.295402(0.004302)& 0.380815(0.005938) \\
% \bottomrule
% \end{tabular}
% \end{table}
\end{document}
\begin{thebibliography}{38}
\providecommand{\natexlab}[1]{#1}
\providecommand{\url}[1]{\texttt{#1}}
\expandafter\ifx\csname urlstyle\endcsname\relax
\providecommand{\doi}[1]{doi: #1}\else
\providecommand{\doi}{doi: \begingroup \urlstyle{rm}\Url}\fi
\bibitem[Ben-Hamu et~al.(2022)Ben-Hamu, Cohen, Bose, Amos, Nickel, Grover, Chen, and Lipman]{BenHamu2022MatchingNF}
Ben-Hamu, H., Cohen, S., Bose, J., Amos, B., Nickel, M., Grover, A., Chen, R. T.~Q., and Lipman, Y.
\newblock Matching normalizing flows and probability paths on manifolds.
\newblock In \emph{International Conference on Machine Learning}, volume 162, pp.\ 1749--1763. PMLR, 2022.
\bibitem[Brakenridge(2017)]{data_flood}
Brakenridge, G.~R.
\newblock Global active archive of large flood events.
\newblock \href{http://floodobservatory.colorado.edu/Archives/index.html}{http://floodobservatory.colorado.edu/Archives/index.html}, 2017.
\newblock Dartmouth Flood Observatory, University of Colorado.
\bibitem[Chen \& Lipman(2024)Chen and Lipman]{chen2024flow}
Chen, R. T.~Q. and Lipman, Y.
\newblock Flow matching on general geometries.
\newblock In \emph{International Conference on Learning Representations}, 2024.
\bibitem[Chen et~al.(2018)Chen, Rubanova, Bettencourt, and Duvenaud]{chen2018neural}
Chen, R. T.~Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D.~K.
\newblock Neural ordinary differential equations.
\newblock In \emph{Advances in Neural Information Processing Systems}, volume~31, 2018.
\bibitem[Ciccotti et~al.(2008)Ciccotti, Leli{\`e}vre, and Vanden-Eijnden]{projection_diffusion}
Ciccotti, G., Leli{\`e}vre, T., and Vanden-Eijnden, E.
\newblock Projection of diffusions on submanifolds: Application to mean force computation.
\newblock \emph{Communications on Pure and Applied Mathematics: A Journal Issued by the Courant Institute of Mathematical Sciences}, 61\penalty0 (3):\penalty0 371--408, 2008.
\bibitem[Coutsias et~al.(2004)Coutsias, Seok, and Dill]{rmsd_quaternions}
Coutsias, E.~A., Seok, C., and Dill, K.~A.
\newblock Using quaternions to calculate {RMSD}.
\newblock \emph{Journal of Computational Chemistry}, 25\penalty0 (15):\penalty0 1849--1857, 2004.
\bibitem[Crane et~al.(2013)Crane, Pinkall, and Schr{\"o}der]{crane2013robust}
Crane, K., Pinkall, U., and Schr{\"o}der, P.
\newblock Robust fairing via conformal curvature flow.
\newblock \emph{ACM Transactions on Graphics (TOG)}, 32\penalty0 (4):\penalty0 1--10, 2013.
\bibitem[De~Bortoli et~al.(2022)De~Bortoli, Mathieu, Hutchinson, Thornton, Teh, and Doucet]{de2022riemannian}
De~Bortoli, V., Mathieu, E., Hutchinson, M., Thornton, J., Teh, Y.~W., and Doucet, A.
\newblock Riemannian score-based generative modelling.
\newblock In \emph{Advances in Neural Information Processing Systems}, volume~35, pp.\ 2406--2422, 2022.
\bibitem[EOSDIS(2020)]{data_fire}
EOSDIS.
\newblock Active fire data.
\newblock \href{https://earthdata.nasa.gov/earth-observation-data/near-real-time/firms/active-fire-data}{https://earthdata.nasa.gov/earth-observation-data/near-real-time/firms/active-fire-data}, 2020.
\newblock Land, Atmosphere Near real-time Capability for EOS (LANCE) system operated by NASA's Earth Science Data and Information System (ESDIS).
\bibitem[Fiorin et~al.(2013)Fiorin, Klein, and H\'enin]{Fiorin2013}
Fiorin, G., Klein, M.~L., and H\'enin, J.
\newblock Using collective variables to drive molecular dynamics simulations.
\newblock \emph{Molecular Physics}, 111\penalty0 (22-23, SI):\penalty0 3345--3362, 2013.
\bibitem[Gropp et~al.(2020)Gropp, Yariv, Haim, Atzmon, and Lipman]{gropp2020implicit}
Gropp, A., Yariv, L., Haim, N., Atzmon, M., and Lipman, Y.
\newblock Implicit geometric regularization for learning shapes.
\newblock In \emph{International Conference on Machine Learning}, volume 119, pp.\ 3569--3579. PMLR, 2020.
\bibitem[Ho et~al.(2020)Ho, Jain, and Abbeel]{NEURIPS2020_ddpm}
Ho, J., Jain, A., and Abbeel, P.
\newblock Denoising diffusion probabilistic models.
\newblock In \emph{Advances in Neural Information Processing Systems}, volume~33, pp.\ 6840--6851, 2020.
\bibitem[Hsu(2002)]{analysis_manifold}
Hsu, E.~P.
\newblock \emph{Stochastic analysis on manifolds}.
\newblock Number~38 in Graduate Studies in Mathematics. American Mathematical Society, 2002.
\bibitem[Huang et~al.(2022)Huang, Aghajohari, Bose, Panangaden, and Courville]{huang2022riemannian}
Huang, C.-W., Aghajohari, M., Bose, J., Panangaden, P., and Courville, A.
\newblock Riemannian diffusion models.
\newblock In \emph{Advances in Neural Information Processing Systems}, volume~35, pp.\ 2750--2761, 2022.
\bibitem[Hyv{\"a}rinen \& Dayan(2005)Hyv{\"a}rinen and Dayan]{hyvarinen2005estimation}
Hyv{\"a}rinen, A. and Dayan, P.
\newblock Estimation of non-normalized statistical models by score matching.
\newblock \emph{Journal of Machine Learning Research}, 6\penalty0 (4), 2005.
\bibitem[Jo \& Hwang(2024)Jo and Hwang]{jo2024generativemodelingmanifoldsmixture}
Jo, J. and Hwang, S.~J.
\newblock Generative modeling on manifolds through mixture of {Riemannian} diffusion processes.
\newblock In \emph{International Conference on Machine Learning}, volume 235, pp.\ 22348--22370. PMLR, 2024.
\bibitem[Kabsch(1976)]{kabsch1976solution}
Kabsch, W.
\newblock A solution for the best rotation to relate two sets of vectors.
\newblock \emph{Acta Crystallographica Section A: Crystal Physics, Diffraction, Theoretical and General Crystallography}, 32\penalty0 (5):\penalty0 922--923, 1976.
\bibitem[Leli{\`e}vre et~al.(2019)Leli{\`e}vre, Rousset, and Stoltz]{hmc-submanifold-tony}
Leli{\`e}vre, T., Rousset, M., and Stoltz, G.
\newblock {{Hybrid Monte Carlo methods for sampling probability measures on submanifolds}}.
\newblock \emph{Numerische Mathematik}, 143\penalty0 (2):\penalty0 379--421, 2019.
\bibitem[Leli{\`e}vre et~al.(2022)Leli{\`e}vre, Stoltz, and Zhang]{LelievreStoltzZhang2022}
Leli{\`e}vre, T., Stoltz, G., and Zhang, W.
\newblock Multiple projection {MCMC} algorithms on submanifolds.
\newblock \emph{IMA Journal of Numerical Analysis}, 43\penalty0 (2):\penalty0 737--788, 2022.
\bibitem[Leli{\`e}vre et~al.(2024)Leli{\`e}vre, Pigeon, Stoltz, and Zhang]{ae2024}
Leli{\`e}vre, T., Pigeon, T., Stoltz, G., and Zhang, W.
\newblock Analyzing multimodal probability measures with autoencoders.
\newblock \emph{The Journal of Physical Chemistry B}, 128\penalty0 (11):\penalty0 2607--2631, 2024.
\bibitem[Lou et~al.(2020)Lou, Lim, Katsman, Huang, Jiang, Lim, and De~Sa]{lou2020neural}
Lou, A., Lim, D., Katsman, I., Huang, L., Jiang, Q., Lim, S.~N., and De~Sa, C.~M.
\newblock Neural manifold ordinary differential equations.
\newblock In \emph{Advances in Neural Information Processing Systems}, volume~33, pp.\ 17548--17558, 2020.
\bibitem[Lou et~al.(2023)Lou, Xu, Farris, and Ermon]{NEURIPS2023_ScalingRDM}
Lou, A., Xu, M., Farris, A., and Ermon, S.
\newblock Scaling {Riemannian} diffusion models.
\newblock In \emph{Advances in Neural Information Processing Systems}, volume~36, pp.\ 80291--80305, 2023.
\bibitem[Mathieu \& Nickel(2020)Mathieu and Nickel]{mathieu2020riemannian}
Mathieu, E. and Nickel, M.
\newblock Riemannian continuous normalizing flows.
\newblock In \emph{Advances in Neural Information Processing Systems}, volume~33, pp.\ 2503--2515, 2020.
\bibitem[NOAA(2020{\natexlab{a}})]{data_earthquake}
NOAA.
\newblock Global significant earthquake database.
\newblock \href{https://data.nodc.noaa.gov/cgi-bin/iso?id=gov.noaa.ngdc.mgg.hazards:G012153}{https://data.nodc.noaa.gov/cgi-bin/iso?id=gov.noaa.ngdc.mgg.hazards:G012153}, 2020{\natexlab{a}}.
\newblock National Geophysical Data Center / World Data Service (NGDC/WDS): NCEI/WDS Global Significant Earthquake Database. NOAA National Centers for Environmental Information.
\bibitem[NOAA(2020{\natexlab{b}})]{data_volcano}
NOAA.
\newblock Global significant volcanic eruptions database.
\newblock \href{https://data.nodc.noaa.gov/cgi-bin/iso?id=gov.noaa.ngdc.mgg.hazards:G10147}{https://data.nodc.noaa.gov/cgi-bin/iso?id=gov.noaa.ngdc.mgg.hazards:G10147}, 2020{\natexlab{b}}.
\newblock National Geophysical Data Center / World Data Service (NGDC/WDS): NCEI/WDS Global Significant Volcanic Eruptions Database. NOAA National Centers for Environmental Information.
\bibitem[Polyak \& Juditsky(1992)Polyak and Juditsky]{polyak1992acceleration}
Polyak, B.~T. and Juditsky, A.~B.
\newblock Acceleration of stochastic approximation by averaging.
\newblock \emph{SIAM Journal on Control and Optimization}, 30\penalty0 (4):\penalty0 838--855, 1992.
\bibitem[Rozen et~al.(2021)Rozen, Grover, Nickel, and Lipman]{rozen2021moser}
Rozen, N., Grover, A., Nickel, M., and Lipman, Y.
\newblock Moser flow: Divergence-based generative modeling on manifolds.
\newblock \emph{Advances in Neural Information Processing Systems}, 34:\penalty0 17669--17680, 2021.
\bibitem[Satorras et~al.(2021)Satorras, Hoogeboom, and Welling]{satorras2021n}
Satorras, V.~G., Hoogeboom, E., and Welling, M.
\newblock E(n) equivariant graph neural networks.
\newblock In \emph{International Conference on Machine Learning}, volume 139, pp.\ 9323--9332. PMLR, 2021.
\bibitem[Sohl-Dickstein et~al.(2015)Sohl-Dickstein, Weiss, Maheswaranathan, and Ganguli]{pmlr-v37-sohl-dickstein15}
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S.
\newblock Deep unsupervised learning using nonequilibrium thermodynamics.
\newblock In \emph{International Conference on Machine Learning}, volume~37, pp.\ 2256--2265. PMLR, 2015.
\bibitem[Song \& Ermon(2019)Song and Ermon]{song2019generative}
Song, Y. and Ermon, S.
\newblock Generative modeling by estimating gradients of the data distribution.
\newblock In \emph{Advances in Neural Information Processing Systems}, volume~32, pp.\ 11895--11907, 2019.
\bibitem[Song et~al.(2021{\natexlab{a}})Song, Durkan, Murray, and Ermon]{NEURIPS2021_likelihood_training_sbdm}
Song, Y., Durkan, C., Murray, I., and Ermon, S.
\newblock Maximum likelihood training of score-based diffusion models.
\newblock In \emph{Advances in Neural Information Processing Systems}, volume~34, pp.\ 1415--1428, 2021{\natexlab{a}}.
\bibitem[Song et~al.(2021{\natexlab{b}})Song, Sohl-Dickstein, Kingma, Kumar, Ermon, and Poole]{song2021scorebased}
Song, Y., Sohl-Dickstein, J., Kingma, D.~P., Kumar, A., Ermon, S., and Poole, B.
\newblock Score-based generative modeling through stochastic differential equations.
\newblock In \emph{International Conference on Learning Representations}, 2021{\natexlab{b}}.
\bibitem[Turk \& Levoy(1994)Turk and Levoy]{turk1994zippered}
Turk, G. and Levoy, M.
\newblock Zippered polygon meshes from range images.
\newblock In \emph{Proceedings of the 21st Annual Conference on Computer Graphics and Interactive Techniques}, pp.\ 311--318, 1994.
\bibitem[Van Der~Spoel et~al.(2005)Van Der~Spoel, Lindahl, Hess, Groenhof, Mark, and Berendsen]{GROMACS}
Van Der~Spoel, D., Lindahl, E., Hess, B., Groenhof, G., Mark, A.~E., and Berendsen, H.~J.
\newblock {{G}{R}{O}{M}{A}{C}{S}: {F}ast, flexible, and free}.
\newblock \emph{Journal of Computational Chemistry}, 26\penalty0 (16):\penalty0 1701--1718, 2005.
\bibitem[Xu et~al.(2022)Xu, Yu, Song, Shi, Ermon, and Tang]{xu2022geodiff}
Xu, M., Yu, L., Song, Y., Shi, C., Ermon, S., and Tang, J.
\newblock Geodiff: A geometric diffusion model for molecular conformation generation.
\newblock In \emph{International Conference on Learning Representations}, 2022.
\bibitem[Zappa et~al.(2018)Zappa, Holmes-Cerfon, and Goodman]{goodman-submanifold}
Zappa, E., Holmes-Cerfon, M., and Goodman, J.
\newblock Monte {Carlo} on manifolds: Sampling densities and integrating functions.
\newblock \emph{Communications on Pure and Applied Mathematics}, 71\penalty0 (12):\penalty0 2609--2647, 2018.
\bibitem[Zhang(2020)]{zhang2020}
Zhang, W.
\newblock Ergodic {SDEs} on submanifolds and related numerical sampling schemes.
\newblock \emph{ESAIM: Mathematical Modelling and Numerical Analysis}, 54\penalty0 (2):\penalty0 391--430, 2020.
\bibitem[Zhang et~al.(2014)Zhang, Wang, Hartmann, Weber, and Sch{\"u}tte]{zwhws2014}
Zhang, W., Wang, H., Hartmann, C., Weber, M., and Sch{\"u}tte, C.
\newblock Applications of the cross-entropy method to importance sampling and optimal control of diffusions.
\newblock \emph{SIAM Journal on Scientific Computing}, 36\penalty0 (6):\penalty0 A2654--A2672, 2014.
\end{thebibliography}