# Diffusion Models for Constrained Domains

Nic Fishman  Leo Klarner  Valentin De Bortoli  Emile Mathieu  Michael
Hutchinson

###### Abstract

Denoising diffusion models are a recent class of generative models which
achieve state-of-the-art results in many domains such as unconditional image
generation and text-to-speech tasks. They consist of a noising process
destroying the data and a backward stage defined as the time-reversal of the
noising diffusion. Building on their success, diffusion models have recently
been extended to the Riemannian manifold setting. Yet, these Riemannian
diffusion models require geodesics to be defined for all times. While this
setting encompasses many important applications, it does not include manifolds
defined via a set of inequality constraints, which are ubiquitous in many
scientific domains such as robotics and protein design. In this work, we
introduce two methods to bridge this gap. First, we design a noising process
based on the _logarithmic barrier_ metric induced by the inequality
constraints. Second, we introduce a noising process based on the _reflected_
Brownian motion. As existing diffusion model techniques cannot be applied in
this setting, we derive new tools to define such models in our framework. We
empirically demonstrate the applicability of our methods to a number of
synthetic and real-world tasks, including the constrained conformational
modelling of protein backbones and robotic arms.

Machine Learning, ICML

Figure 1: The behaviour of different types of noising processes considered in
this work defined on the unit interval. $\mathbf{B}_{t}$: Euclidean
(unconstrained) Brownian motion. $\hat{\mathbf{B}}_{t}$: Log barrier forward
noising process. $\bar{\mathbf{B}}_{t}$: Reflected Brownian motion. All
sampled with the same initial point and driving noise. Black line indicates
the boundary.

##  1 Introduction

Diffusion models (Sohl-Dickstein et al., 2015; Song & Ermon, 2019; Song et
al., 2021; Ho et al., 2020; Dhariwal & Nichol, 2021) have recently been
introduced as a powerful new paradigm for generative modelling. They work as
follows: noise is progressively added to data following a Stochastic
Differential Equation (SDE)—the forward _noising_ process—until it is
approximately Gaussian. The generative model is given by an approximation of
the associated _time-reversed process_ called the backward _denoising_
process. This process is also an SDE whose drift depends on the gradient of
the logarithm of the densities of the forward process, referred as the Stein
score. This score is approximated by leveraging techniques from deep learning
and score matching (Hyvärinen, 2005; Vincent, 2011).

Recently, this framework has been extended to a broad class of Riemannian
manifolds, see De Bortoli et al. (2022); Huang et al. (2022), opening the door
to many applications in the natural sciences (e.g. Trippe et al., 2022; Corso
et al., 2022; Watson et al., 2022; Leach et al., 2022; Urain et al., 2022). In
particular, this framework has been applied to Lie groups such as the group of
rotations $\mathrm{SO}(3)$ (Trippe et al., 2022; Corso et al., 2022) or the
group of rigid body motions $\mathrm{SE}(3)$ (Urain et al., 2022). However,
both De Bortoli et al. (2022) and Huang et al. (2022) require that geodesics,
i.e length minimising curves on the manifold, are defined for all times. As a
consequence, this methodology cannot be applied to manifolds that are defined
by a set of inequality constraints. Indeed, geodesics might then not be
defined at all times, as illustrated on Figure 2. This rich class of
constrained manifolds comprises a large variety of settings such as polytopes
and convex sets of Euclidean spaces. These manifolds are ubiquitous in several
scientific fields such as computational statistics and biology (Thiele et al.,
2013; Morris, 2002; Lewis et al., 2012).

As sampling problems are also important on these manifolds (Kook et al., 2022;
Heirendt et al., 2019), a flurry of Markov chain based methods have been
developed to sample from unnormalised densities. Successful algorithms include
the reflected Brownian motion (Williams, 1987; Petit, 1997; Shkolnikov &
Karatzas, 2013), log-barrier methods (Kannan & Narayanan, 2009; Lee & Vempala,
2017; Noble et al., 2022; Kook et al., 2022; Gatmiry & Vempala, 2022; Lee &
Vempala, 2018) and hit-and-run approaches in the case of polytopes (Smith,
1984; Lovász & Vempala, 2006).

In this work, we study the generative modelling counterparts of these
algorithms through the lens of diffusion models. Among existing algorithms for
statistical sampling on constrained manifolds, the reflected Brownian motion
(Williams, 1987) and the geodesics Brownian motion (Lee & Vempala, 2017) are
continuous stochastic processes, thus well suited for extending the continuous
diffusion on manifolds framework developed by De Bortoli et al. (2022); Huang
et al. (2022). This is what we achieve in this work. In particular, we
introduce two principled diffusion models for generative modelling on
constrained manifolds: a) one based on the geodesic Brownian motion,
leveraging tools from the log-barrier methods, b) the other one based on the
reflected Brownian motion. In both cases, we show how one can extend the ideas
of time-reversal and score matching to these settings. We empirically
demonstrate the scalability and flexbility of these methods on high-
dimensional convex polytopes and symmetric positive definite matrices, with
applications from robotics and protein design.

##  2 Background

#### Riemannian manifolds.

A Riemannian manifold is a tuple $(\mathcal{M},\mathfrak{g})$ with
$\mathcal{M}$ a smooth manifold and $\mathfrak{g}$ a metric which defines an
inner product on tangent spaces. The metric $\mathfrak{g}$ induces key
quantities on the manifold, such as an exponential map
$\exp_{x}:\mathrm{T}_{x}\mathcal{M}\rightarrow\mathcal{M}$, the notion of
following straight lines on manifolds, a gradient operator $\nabla$111The
(Riemannian) gradient $\nabla$ is defined s.t. for any smooth
$f\in\mathrm{C}^{\infty}(\mathcal{M})$,
$x\in\mathcal{M},v\in\mathrm{T}_{x}\mathcal{M}$,
$\mathfrak{g}(x)(f(x),v)=\mathrm{d}f(x)(v)$. and a divergence operator
$\mathrm{div}$222The Riemannian divergence acts vector fields and can be
defined using the volume form of $\mathcal{M}$.. The metric induces the
Laplace-Beltrami operator $\Delta$ and consequently a Brownian motion with
density (w.r.t. the volume form333We assume that $\mathcal{M}$ is orientable
and therefore that a volume form exists.) whose density is given by the heat
equation ${\partial_{t}}p_{t}=\Delta p_{t}$. We refer the reader to Appendix B
for a brief introduction to differential geometry, to Lee (2013) for a
thorough treatment and to Hsu (2002) for details on stochastic analysis on
manifolds.

#### Constrained manifolds.

Figure 2:  A convex polytope defined by six constraints
$\\{f_{i}\\}_{i\in\mathcal{I}}$, along with the log barrier potential, and
‘straight trajectories’ under the log-barrier metric and under the Euclidean
metric with and without reflection at the boundary.

In this work, we are concerned with _constrained_ manifolds. More precisely,
given a manifold $(\mathcal{N},\mathfrak{h})$, we consider a family of real
functions $\\{f_{i}:\mathcal{N}\to\mathbb{R}\\}_{i\in\mathcal{I}}$ indexed by
$\mathcal{I}$. We then define

$\mathcal{M}=\\{x\in\mathcal{N}\,:\;f_{i}(x)<0,\ i\in\mathcal{I}\\}.$ (1)

In this scenario, $\\{f_{i}\\}_{i\in\mathcal{I}}$ is interpreted as a set of
constraints on $\mathcal{N}$. This setting encompasses a large class of
(topological) manifolds. For instance, using the Hahn-Banach theorem (Brezis &
Brézis, 2011, Chapter 1, Theorem 1.7), it can be shown that every closed
convex set can be described using (1)444Replacing $<$ by $\leq$.. Similarly,
with affine constraints $f_{i}(x)=\langle a_{i},x\rangle-b_{i}$,
$x\in\mathbb{R}^{d}$, we get a polytope as illustrated in Figure 2.

Additionally, since sampling from a distribution defined on a manifold
described by (1) can be understood as sampling from a distribution whose
supports is defined by _constraints_ , this setting naturally appears in many
areas of science (Han & Rudolph, 2006). While the two methods we introduce in
Section 3 can be applied to the general framework (1), in our applications, we
focus on two specific settings:  (a) polytopes—$\mathcal{N}=\mathbb{R}^{d}$,
(b) symmetric positive definite matrices (SPD) under trace
conditions—$\mathcal{N}=\mathcal{S}^{++}(d)$.

#### Continuous diffusion models.

We briefly recall the framework for constructing continuous diffusion
processes introduced by Song et al. (2021) in the context of generative
modelling over $\mathbb{R}^{d}$. First, we consider a forward noising process
$(\mathbf{X}_{t})_{t\in\left[0,T\right]}$ which progressively noise a data
distribution $p_{0}$ into a Gaussian $\mathrm{N}(0,\operatorname{Id})$. In
particular we consider an Ornstein–Uhlenbeck (OU) process which is defined by
the following stochastic differential equation (SDE)

$\mathrm{d}\mathbf{X}_{t}=-\tfrac{1}{2}\mathbf{X}_{t}\mathrm{d}t+\mathrm{d}\mathbf{B}_{t},\
\mathbf{X}_{0}\sim p_{0}.$ (2)

Under mild conditions on $p_{0}$, the time-reversed process
$(\overleftarrow{\mathbf{X}}_{t})_{t\in\left[0,T\right]}=(\mathbf{X}_{T-t})_{t\in\left[0,T\right]}$
also satisfies an SDE (Cattiaux et al., 2021; Haussmann & Pardoux, 1986) given
by

$\mathrm{d}\overleftarrow{\mathbf{X}}_{t}=\\{\tfrac{1}{2}\overleftarrow{\mathbf{X}}_{t}+\nabla\log
p_{T-t}(\overleftarrow{\mathbf{X}}_{t})\\}\mathrm{d}t+\mathrm{d}\mathbf{B}_{t},\
\overleftarrow{\mathbf{X}}_{0}\sim p_{T},$ (3)

where $p_{t}$ denotes the density of $\mathbf{X}_{t}$. The quantity
$\nabla\log p_{t}$ is referred as the Stein score and is unavailable in
practice. It can be approximated with a score network $s_{\theta}(t,\cdot)$
trained by minimising an implicit score matching ($\mathrm{ism}$) loss

$\textstyle{\mathcal{L}(\theta)=\mathbb{E}[\lambda_{t}\\{\tfrac{1}{2}\left\lVert
s_{\theta}(t,\mathbf{X}_{t})\right\rVert^{2}+\mathrm{div}(s_{\theta})(t,\mathbf{X}_{t})}\\}],$
(4)

where $\lambda_{t}>0$ is a weighting function, and the expectation is taken
over $t\sim\mathcal{U}([0,T])$ and $(\mathbf{X}_{0},\mathbf{X}_{t})$. For an
arbitrarily flexible score network, the global minimiser
$\theta^{\star}=\mathrm{argmin}_{\theta}\mathcal{L}(\theta)$ satisfies
$s_{\theta^{\star}}(t,\cdot)=\nabla\log p_{t}$.

##  3 Inequality constrained diffusion models

We are now ready to introduce our methodology to deal with manifolds defined
via inequality _constraints_ (1). We first propose in Section 3.1, a
Riemannian diffusion model endowed with a metric induced by a log-barrier
potential. In Section 3.2, we introduce a _reflected_ diffusion model. While
both models extend classical diffusion models to the inequality constrained
settings, they exhibit very different behaviour. We discuss their practical
differences in Section 5.1.

###  3.1 Log-barrier diffusion models

Figure 3: Convergence of the Barrier Langevin dynamics on the unit interval to
the uniform distribution. $(\mathcal{N},\mathfrak{h})$
$(\mathcal{M},\nabla^{2}\phi)$ Original spaceConstrained space Figure 4:
Illustrative diagram of the barrier method and the change of metric.

#### Barrier Langevin Dynamics.

Barrier methods work by constructing a smooth potential
$\phi:\mathcal{M}\rightarrow\mathbb{R}$ such that it blows up on the boundary
of a desired set, see Nesterov et al. (2018). Such potentials are at the basis
of interior methods in optimisation (Boyd et al., 2004). Among these
functions, the _logarithmic barrier_ is the most popular among practitioners
(Lee & Vempala, 2017). For a convex polytope $\mathcal{M}$ defined by the
constraints $\mathrm{A}x<b$, the logarithmic barrier is given for any
$x\in\mathcal{M}$ by

$\textstyle{\phi(x)=-\sum_{i=1}^{m}\log(\langle\mathrm{A}_{i},x\rangle-
b_{i}).}$ (5)

Assuming that $\|\mathrm{A}_{i}\|=1$, we have that for any $x\in\mathcal{M}$,
$\phi(x)=-\sum_{i=1}^{m}\log(d(x,\partial\mathcal{M}_{i}))$, where
$\partial\mathcal{M}_{i}=\\{x\in\mathbb{R}^{d}\,:\;\langle\mathrm{A}_{i},x\rangle=b\\}$.
The core idea of barrier methods is to ‘warp the geometry’ of the constrained
space, so that sampling on this warped space bypasses the need to explicitly
deal with the boundary, whist ensuring that the returned samples are
distributed according to the target distribution on the original space (and
w.r.t. the original metric)—see Figure 4. Assuming $\phi$ to be strictly
convex and smooth—which is for example verified for (5)—its Hessian
$\nabla^{2}\phi$ is positive definite thus it defines a valid Riemannian
metric. The formal approach to ’transforming the geometry’ of the convex space
with boundary, is to endow that space with the Riemannian metric
$\mathfrak{g}=\nabla^{2}\phi$, thus referred as Hessian manifolds. For (5), we
have that $\mathfrak{g}(x)=\mathrm{A}^{\top}S^{-2}(x)\mathrm{A}$ with
$S(x)=\text{diag}(b_{i}-\langle\mathrm{A}_{i},x\rangle)_{i}$.

Let’s consider the following Langevin dynamics as the forward process (Lee &
Vempala, 2017)

$\textstyle{\mathrm{d}\mathbf{X}_{t}=\tfrac{1}{2}\mathrm{div}(\mathfrak{g}^{-1})(\mathbf{X}_{t})\mathrm{d}t+\mathfrak{g}(\mathbf{X}_{t})^{-\frac{1}{2}}\mathrm{d}\mathbf{B}_{t},}$
(6)

with
$\mathrm{div}(F)(x)\triangleq\left(\mathrm{div}(F_{1})(x),\dots,\mathrm{div}(F_{d})(x)\right)^{\top}$,
for any smooth $F:\ \mathbb{R}^{d}\to\mathbb{R}^{d}$. Denoting $p_{t}$ the
density of $\mathbf{X}_{t}$, we have that
$\partial_{t}p_{t}=\mathrm{Tr}(\mathfrak{g}^{-1}\nabla^{2}p_{t})$, see
Appendix C for more details on this process. Hence, assuming that
$\mathcal{M}$ is compact, the uniform distribution on $\mathcal{M}$, is the
invariant measure and we have that the distribution of
$(\mathbf{X}_{t})_{t\geq 0}$ converges exponentially fast towards the uniform,
see Lee & Vempala (2017, Theorem 23).

#### Time-reversal.

Assuming that $\mathfrak{g}^{-1}$ and its derivative are bounded on
$\mathcal{M}$, the time-reversal of (6) is given by Cattiaux et al. (2021), in
particular we have

$\displaystyle\mathrm{d}\overleftarrow{\mathbf{X}}_{t}$
$\displaystyle=\textstyle{[-\frac{1}{2}\mathrm{div}(\mathfrak{g}^{-1})+\mathrm{div}(\mathfrak{g}^{-1})}$
(7) $\displaystyle\quad\quad\textstyle{+\mathfrak{g}^{-1}\nabla\log
p_{T-t}](\overleftarrow{\mathbf{X}_{t}})\mathrm{d}t+\mathfrak{g}(\overleftarrow{\mathbf{X}}_{t})^{-\frac{1}{2}}\mathrm{d}\mathbf{B}_{t},}$
(8)
$\displaystyle=\textstyle{\left[\frac{1}{2}\mathrm{div}(\mathfrak{g}^{-1})+\mathfrak{g}^{-1}\nabla\log
p_{T-t}\right](\overleftarrow{\mathbf{X}}_{t})\mathrm{d}t+\mathfrak{g}(\overleftarrow{\mathbf{X}}_{t})^{-\frac{1}{2}}\mathrm{d}\mathbf{B}_{t}.}$
(9)

$\overleftarrow{\mathbf{X}}_{0}$ is initialised with the uniform distribution
on $\mathcal{M}$ (which is close to the one of $\mathbf{X}_{T}$ for large
$T$). Practically, one can obtain samples from the uniform reference
distribution (w.r.t the original metric) using any MCMC converging to this
stationary distribution. In this work, we focus on the _hit-and-run_ algorithm
(Claude J. P. Bélisle et al., 1993), which we found to be converging very fast
to the uniform in our polytope settings. The estimation of the score term
$\nabla\log p_{t}$ is done by minimising the $\mathrm{ism}$ loss function (4).
We refer to Section 3.3 for details on the training and parameterisation.

#### Sampling.

Sampling from the forward (6) and backward (8) processes, once the score is
learnt, requires a discretisation scheme. We use Geodesic Random Walks (GRW)
(Jørgensen, 1975) for this purpose, see Algorithm 1. This discretisation is a
generalisation of the Euler-Maruyama discretisation of SDE in Euclidean
spaces, where the $+$ operator is replaced by the exponential mapping on the
manifold, computing the geodesics.

Algorithm 1 _Geodesic Random Walk._ Discretisation of the SDE
$\mathrm{d}\mathbf{X}_{t}=b(t,\mathbf{X}_{t})\mathrm{d}t+\mathrm{d}\mathbf{B}_{t}$.

Require: $T,N,X_{0}^{\gamma},b$

$\gamma=T/N$

for $k\in\\{0,\dots,N-1\\}$ do

$Z_{k+1}\sim\mathrm{N}(0,\operatorname{Id})$

$W_{k+1}=\gamma b(k\gamma,X_{k})+\sqrt{\gamma}Z_{k+1}$

$X_{k+1}^{\gamma}=\exp_{X_{k}}[W_{k+1}]\approx\mathrm{proj}_{\mathcal{M}}(X_{k}+W_{k+1})$

return $\\{X_{k}\\}_{k=0}^{N}$

However, contrary to (De Bortoli et al., 2022), we do not have access
explicitly to the exponential mapping. Instead, we rely on an approximation,
using a _retraction_ , see Ablin & Peyré (2022) for a definition and
alternative schemes. We also emphasise that the log-likelihood of the model
can be evaluated using the equivalent ODE to (8), see Appendix E for more
details.

###  3.2 Reflected diffusion models

Figure 5:  Convergence of the reflected Brownian motion on the unit interval
to the uniform distribution.

Another approach to deal with the geometry of $\mathcal{M}$ is to leave the
metric $\mathfrak{g}$ of $\mathcal{M}$ untouched, but instead to consider an
_unconstrained_ forward dynamics on $\mathcal{N}$, and then restraining it to
$\mathcal{M}$ by considering a _reflected_ version of the process. To simplify
the presentation we focus on the case $\mathcal{N}=\mathbb{R}^{d}$ and
$\partial M$ to be smooth555We refer to Appendix G for a definition of smooth
boundary.. We also assume that $\mathcal{M}$ is compact and convex. In the
next paragraph, we formally introduce the _reflected Brownian motion_.

#### Skorokhod problem.

We say that $(\bar{\mathbf{B}}_{t},\mathbf{k}_{t})_{t\geq 0}$ is a solution to
the _Skorokhod problem_ (Skorokhod, 1961) if $(\mathbf{k}_{t})_{t\geq 0}$ is a
bounded variation process and $(\bar{\mathbf{B}}_{t})_{t\geq 0}$ a continuous
adapted process such that for any $t\geq 0$,

$\bar{\mathbf{B}}_{t}=\bar{\mathbf{B}}_{0}+\mathbf{B}_{t}-\mathbf{k}_{t}\in\mathcal{M},$
(10)

and
$\textstyle{\left|\mathbf{k}\right|_{t}=\int_{0}^{t}\mathbf{1}_{\bar{\mathbf{B}}_{s}\in\partial\mathcal{M}}\mathrm{d}\left|\mathbf{k}\right|_{s}}$,
$\textstyle{\mathbf{k}_{t}=\int_{0}^{t}\mathbf{n}(\bar{\mathbf{B}}_{s})\mathrm{d}\left|\mathbf{k}\right|_{s},}$
where $(\left|\mathbf{k}\right|_{t})_{t\geq 0}$ is the total variation of
$(\mathbf{k}_{t})_{t\geq 0}$. The condition
$\textstyle{\left|\mathbf{k}\right|_{t}=\int_{0}^{t}\mathbf{1}_{\bar{\mathbf{B}}_{s}\in\partial\mathcal{M}}\mathrm{d}\left|\mathbf{k}\right|_{s}}$
can be interpreted as $\mathbf{k}_{t}$ being constant when
$(\bar{\mathbf{B}}_{t})_{t\geq 0}$ does not hit the boundary. When
$(\bar{\mathbf{B}}_{t})_{t\geq 0}$ hits the boundary, the condition
$\textstyle{\mathbf{k}_{t}=\int_{0}^{t}\mathbf{n}(\bar{\mathbf{B}}_{s})\mathrm{d}\left|\mathbf{k}\right|_{s},}$
tells us that $-\mathbf{k}_{t}$ “compensates” for $\bar{\mathbf{B}}_{t}$ by
pushing the process back into $\mathcal{M}$ along the inward normal
$-\mathbf{n}$. As a result $(\bar{\mathbf{B}}_{t})_{t\geq 0}$ can be
understood as the continuous-time counterpart to the reflected Gaussian random
walk. The process $(\mathbf{k}_{t})_{t\geq 0}$ can be related to the notion of
_local time_ (Revuz & Yor, 2013) and quantify the amount of time
$(\bar{\mathbf{B}}_{t})_{t\geq 0}$ spends at the boundary
$\partial\mathcal{M}$. Lions & Sznitman (1984, Theorem 2.1) ensure the
existence and uniqueness of a solution to the Skorokhod problem. One key
observation is that the event
$\\{\bar{\mathbf{B}}_{t}\in\partial\mathcal{M}\\}$ has probability zero
(Harrison & Williams, 1987, Section 7, Lemma 7). As in the _unconstrained_
setting, one can describe the dynamics of the density of
$\bar{\mathbf{B}}_{t}$.

######  Proposition 3.1 (Burdzy et al. (2004)).

For any $t>0$, the distribution of $\bar{\mathbf{B}}_{t}$ admits a density
w.r.t. the Lebesgue measure denoted $p_{t}$. In addition, we have for any
$x\in\mathrm{int}(\mathcal{M})$ and $x_{0}\in\partial\mathcal{M}$

$\partial_{t}p_{t}(x)=\tfrac{1}{2}\Delta
p_{t}(x),\quad\partial_{\mathbf{n}}p_{t}(x_{0})=0,$ (11)

where we recall that $\mathbf{n}$ is the outward normal.

Note that contrary to the unconstrained setting, the heat equation has
_Neumann_ boundary conditions. Similarly to the compact Riemannian setting
(Saloff-Coste, 1994) it can be shown that the reflected Brownian motion
converges to the uniform distribution on $\mathcal{M}$ exponentially fast
(Loper, 2020), see Figure 5. Hence, $(\bar{\mathbf{B}}_{t})_{t\geq 0}$ is a
candidate for a forward noising process in the context of diffusion models.

#### Time-reversal.

In order to extend the diffusion model approach to the reflected setting, we
need to derive a _time-reversal_ for
$(\bar{\mathbf{B}}_{t})_{t\in\left[0,T\right]}$. Namely, we need to
characterise the evolution of
$(\overleftarrow{\mathbf{X}}_{t})_{t\in\left[0,T\right]}=(\bar{\mathbf{B}}_{T-t})_{t\in\left[0,T\right]}$.
It can be shown that the time-reversal of
$(\bar{\mathbf{B}}_{t})_{t\in\left[0,T\right]}$ is also the solution to a
Skorokhod problem.

######  Theorem 3.2.

There exist $(\overleftarrow{\mathbf{k}}_{t})_{t\geq 0}$ a bounded variation
process and a Brownian motion $(\mathbf{B}_{t})_{t\geq 0}$ such that

$\textstyle{\overleftarrow{\mathbf{X}}_{t}=\overleftarrow{\mathbf{X}}_{0}+\mathbf{B}_{t}+\int_{0}^{t}\nabla\log
p_{T-s}(\overleftarrow{\mathbf{X}}_{s})\mathrm{d}s-\overleftarrow{\mathbf{k}}_{t}.}$
(12)

In addition, for any $t\in\left[0,T\right]$ we have

$\textstyle{\overleftarrow{\left|\mathbf{k}\right|}_{t}=\int_{0}^{t}\mathbf{1}_{\overleftarrow{\mathbf{X}}_{s}\in\partial\mathcal{M}}\mathrm{d}\overleftarrow{\left|\mathbf{k}\right|}_{s},\quad\overleftarrow{\mathbf{k}}_{t}=\int_{0}^{t}\mathbf{n}(\overleftarrow{\mathbf{X}}_{s})\mathrm{d}\overleftarrow{\left|\mathbf{k}\right|}_{s}.}$
(13)

The proof, see Appendix G, follows Petit (1997) which provides a time-reversal
in the case where $\mathcal{M}$ is the positive orthant. It is based on an
extension of Haussmann & Pardoux (1986) to the reflected setting, with a
careful handling of the boundary conditions. In particular, contrary to Petit
(1997), we do not rely on an explicit expression of $p_{t}$ but instead use
the intrinsic properties of $(\mathbf{k}_{t})_{t\geq 0}$. Informally, Theorem
3.2 means that the process
$(\overleftarrow{\mathbf{X}}_{t})_{t\in\left[0,T\right]}$ satisfies

$\mathrm{d}\overleftarrow{\mathbf{X}}_{t}=\nabla\log
p_{T-t}(\overleftarrow{\mathbf{X}}_{t})\mathrm{d}t+\mathrm{d}\mathbf{B}_{t}-\mathrm{d}\overleftarrow{\mathbf{k}}_{t},$
(14)

which echoes the usual time-reversal formula (3). In practice, in order to
sample from $(\overleftarrow{\mathbf{X}}_{t})_{t\in\left[0,T\right]}$, one
need to consider the _reflected_ version of the _unconstrained_ dynamics
$\mathrm{d}\overleftarrow{\mathbf{X}}_{t}=\nabla\log
p_{T-t}(\overleftarrow{\mathbf{X}}_{t})\mathrm{d}t+\mathrm{d}\mathbf{B}_{t}$.
The next proposition ensures that the $\mathrm{ism}$ loss is still valid in
this reflected setting.

######  Proposition 3.3.

Let
$s\in\mathrm{C}^{\infty}(\left[0,T\right]\times\mathbb{R}^{d},\mathbb{R}^{d})$
such that for any $x\in\partial\mathcal{M}$ and $t\geq 0$, $s_{t}(x)=0$. Then,
there exists $C>0$ such that

$\textstyle{\mathbb{E}[\|\nabla\log
p_{t}-s_{t}\|^{2}]=\mathbb{E}[\|s_{t}\|^{2}+2~{}\mathrm{div}(s_{t})]}+C,$ (15)

where $\mathbb{E}$ is taken over $\mathbf{X}_{t}\sim p_{t}$ and
$t\sim\mathcal{U}([0,T])$.

The estimation of the score term $\nabla\log p_{t}$ is also done by minimising
the $\mathrm{ism}$ loss function (4), see Section 3.3 for details on the
training and parameterisation of the score.

#### Sampling.

In practice, we approximately sample the reflected dynamics by considering the
Markov chain given by Algorithm 3. We refer to Pacchiarotti et al. (1998);
Bossy et al. (2004) for weak convergence results on this numerical scheme in
the Euclidean setting.

0: $x\in\mathcal{M}$, $\bm{v}\in\mathrm{T}_{x}\mathcal{M}$,
$\\{f_{i}\\}_{i\in\mathcal{I}}$

$\ell\leftarrow\left\lVert\bm{v}\right\rVert_{\mathfrak{g}}$

$\bm{s}\leftarrow\bm{v}/\left\lVert\bm{v}\right\rVert_{\mathfrak{g}}$

while $\ell\geq 0$ do

$d_{i}=\arg\operatorname{intersect}_{t}\left[\exp_{\mathfrak{g}}(x,t\bm{s}),f_{i}\right]$

$i\leftarrow\mathrm{argmin}_{i}\;d_{i}\;s.t.\;d_{i}>0$

$\alpha\leftarrow\min\left(d_{i},\ell\right)$

$x^{\prime}\leftarrow\exp_{\mathfrak{g}}\left(x,\alpha\bm{s}\right)$

$\bm{s}\leftarrow\operatorname{paralleltransport}_{\mathfrak{g}}\left(x,\bm{s},x^{\prime}\right)$

$\bm{s}\leftarrow\operatorname{reflect}\left(\bm{s},f_{i}\right)$

$\ell\leftarrow\ell-\alpha$

$x\leftarrow x^{\prime}$

return $x$

Algorithm 2 _Reflected step algorithm_. The algorithm operates by repeatedly
taking geodesic steps until one of the constraint is violated, or the step is
fully taken. Upon hitting the boundary we parallel transport the tangent
vector to the boundary and then reflect it against the boundary. We then start
a new geodesic from this point in the new direction. The
${\arg}\operatorname{intersect}_{t}$ function computes the distance one must
travel along a geodesic in direction $\bm{s}$ til constraint $f_{i}$ is
intersected. For a discussion of $\operatorname{paralleltransport}$,
$\exp_{\mathfrak{g}}$ and $\operatorname{reflect}$ please see Appendix B
Algorithm 3 _Reflected Random Walk_. Discretisation of the SDE
$\mathrm{d}\mathbf{X}_{t}=b(t,\mathbf{X}_{t})\mathrm{d}t+\mathrm{d}\mathbf{B}_{t}-\mathrm{d}\mathbf{k}_{t}$.

Require: $T,N,X_{0}^{\gamma},\\{f_{i}\\}_{i\in\mathcal{I}}$

$\gamma=T/N$

for $k\in\\{0,\dots,N-1\\}$ do

$Z_{k+1}\sim\mathrm{N}(0,\operatorname{Id})$

$X_{k+1}^{\gamma}=\text{ReflectedStep}[X_{k}^{\gamma},\sqrt{\gamma}Z_{k+1},\\{f_{i}\\}_{i\in\mathcal{I}}]$

return $\\{X_{k}^{\gamma}\\}_{k=0}^{N}$

Figure 6: _Left:_ Reflection against a linear boundary. For a step $s$ with
magnitude $|s|$ and direction $\hat{s}$ the distance to the boundary described
by the normal $\hat{n}$ and offset $b$ is
$d=\frac{\left\langle\hat{s},x_{t}\right\rangle-b}{\left\langle\hat{s},\hat{n}\right\rangle}$.
The reflected direction is given by
$\hat{s}^{\prime}=\hat{s}-2\left\langle\hat{s},\hat{n}\right\rangle\hat{n}$.
_Right:_ Reflection against a spherical boundary. For a sphere of radius $r$,
the distance to the boundary from $x_{t}$ in direction $s$ is given by
$d=\frac{1}{2}(\left\langle\hat{s},x_{t}\right\rangle^{2}+4(r^{2}-\left\lVert
x_{t}\right\rVert^{2}))^{1/2}-\frac{1}{2}\left\langle\hat{s},x_{t}\right\rangle$.
The normal at the intersection can be computed as the unit vector in the
direction $-2(d\hat{s}+x_{t})$, and then $\hat{s}^{\prime}$ as above.

###  3.3 Score network parametrisation and training

The score term $\nabla\log p_{T-t}$ appearing in both time-reversal processes
(8) and (14) is intractable. It is thus approximated with a _score network_
$s_{\theta}(t,\cdot)\approx\nabla\log p_{t}$. We use a multi-layer perceptron
architecture with $\sin$ activation functions, see Appendix J for more details
on the experimental setup. Due to boundary condition (11), the normal
component of the score is zero at the boundary in the reflected case. A
similar result holds in the log-barrier setting. We include this in our score
parameterisation by scaling the output of the neural network as
$s_{\theta}(t,x)=\textrm{ReLU}(d(x,\partial\mathcal{M})-\delta)\cdot\textrm{NN}_{\theta}(t,x)$
with $\delta>0$ a threshold, see Appendix J for an illustration. The weighting
function in (4) is set to $\lambda_{t}=(t+1)$.

The forward processes (6) and (10) for the barrier and reflected methods
cannot be sampled in closed form, so at training time samples from the
conditional marginals $p(\mathbf{X}_{t}|\mathbf{X}_{0})$ are obtained by
discretising these processes. As to take the most of this computational
overhead, we use several samples from the discretised forward trajectory
$(\mathbf{X}_{t_{1}},\cdots,\mathbf{X}_{t_{k}}|\mathbf{X}_{0})$ instead of
only using the last sample.

##  4 Related Work

#### Sampling on constrained manifolds.

Sampling from a distribution on a space defined by a set of constraints is an
important ingredient in several computational tasks, such as computing the
volume of a polytope (Lee & Vempala, 2017). Incorporating such constraints
inside MCMC algorithms while preserving fast convergence properties is an
active field of research (Kook et al., 2022; Lee & Vempala, 2017; Noble et
al., 2022). In this work, we are interested in sampling from the uniform
distribution defined on the constrained set in order to define a proper
_forward process_ for our diffusion model. Log-barrier methods such as the
Dikin walk or Riemannian Hamiltonian Monte Carlo (Kannan & Narayanan, 2009;
Lee & Vempala, 2017; Noble et al., 2022) change the geometry of the underlying
space and define stochastic processes which never violate the constraints, see
(Kannan & Narayanan, 2009; Noble et al., 2022) for instance. If we keep the
Euclidean metric, then a classical random walk discretising the Brownian
motion will almost surely violate the constraints. The manifold is said to be
geodesically not complete. To counter this effect, it has been proposed to
_reflect_ the Brownian motion (Williams, 1987; Petit, 1997; Shkolnikov &
Karatzas, 2013). Finally, we also highlight hit-and-run approaches (Smith,
1984; Lovász & Vempala, 2006) which generalise Gibbs’ algorithm and enjoy fast
convergence properties provided that one knows how to sample from the one-
dimensional marginals.

#### Diffusion models on manifolds.

De Bortoli et al. (2022) extended the work of Song et al. (2021) to Riemannian
manifolds by defining forward and backward stochastic processes in this
setting. Concurrently, a similar framework was introduced by Huang et al.
(2022) extending the maximum likelihood approach of Huang et al. (2021).
Existing applications of denoising diffusion models on Riemannian manifolds
have been focused on well-known manifolds for which one can find metrics so
that the framework of De Bortoli et al. (2022) applies. In particular, on
compact Lie groups, geodesics and Brownian motions can be defined in a
canonical manner. Leach et al. (2022) define diffusion models on
$\mathrm{SO}(3)$ for rotational alignment, while Jing et al. (2022) consider
the product of tori for molecular conformer generation. Corso et al. (2022)
use diffusion models on
$\mathbb{R}^{3}\times\mathrm{SO}(3)\times\mathrm{SO}(2)$ for protein docking
applications. Finally, Urain et al. (2022) introduce a methodology for
$\mathrm{SE}(3)$ diffusion models with applications to robotics.

##  5 Experimental results

To illustrate the modelling ability of the constrained diffusion models
introduced in Section 3, we first construct a series of increasingly difficult
synthetic tasks on convex polytopes such as the hypercube and simplex in
Section 5.1. We then demonstrate their practical applicability for two real-
world problems from robotics and molecular biology. In particular in Section
5.2, we show that these models are able to model distributions on the space of
$d\times d$ symmetric positive definite (SPD) matrices $S_{++}^{d}$ under
trace constraints, a setting that is essential to describing and controlling
the motions and exerted forces of robotic platforms (Yoshikawa, 1985). Our
dataset is taken from (Jaquier et al., 2021), see Appendix I for details. In
Section 5.3, we map the problem of modelling the conformational ensembles of
proteins under positional constraints on their endpoints, to the product
manifold of a convex polytope and a torus (Han & Rudolph, 2006). We refer the
read to Appendix J for more details on the experimental setup. Additionally,
the codebase will be made public upon publication.

###  5.1 Convex polytopes and methods comparison

#### Synthetic data.

First, we aim to assess the ability of our proposed methods to scale with the
space dimensionality. We focus on polytopes, and construct synthetic datasets
from bimodal distributions. We recall that a polytope is an open set in
$\mathbb{R}^{d}$ with affine constraints, i.e. such that $\mathrm{A}x<b$,
$\mathrm{A}\in\mathbb{R}^{p\times d}$ and $b\in\mathbb{R}^{p}$. In particular
we investigate two specific instances of polytopes: hypercubes and simplices.
In Section J.1 we also present results on the Birkhoff polytope. We measure
the performance of trained models via the Maximum Mean Discrepancy (MMD)
(Gretton et al., 2012), which is a kernel based metric between distributions.

#### Comparison.

On Figure 7, we qualitatively observe that both methods recover the data
distribution on the two-dimensional hypercube and simplex, although the
reflected method seem to have a better fit. In Table 1, we report the MMD
between the data distribution and the learnt diffusion models. We similarly
observe that the reflected method consistently gives better results than the
log-barrier one. One key drawback of the log-barrier method is the slow
convergence of the forward, which can already be seen in Figure 3. It should
also be highlighted that in order to implement the log-barrier method, one
needs in general to invert a $d\times d$ matrix, which can be implemented
using $O(d^{3})$ operations. This prevents the use of the log-barrier method
for very high dimensional applications. On the other hand, the likelihood
evaluation tools from Riemannian diffusion models (De Bortoli et al., 2022,
Appendix D) can readily be adapted to the log-barrier method, while no
straightforward equivalent exists for the reflected method (see Appendix E).

Square

Simplex

DataLog-barrierReflected

Figure 7:  Histograms of samples from the data distribution and from trained
constrained diffusion models.

Figure 8:  (a-c) 100 polypeptides samples—modelled as chain of $C_{\alpha}$ atoms—from the training dataset, a trained reflected model and from the uniform distribution over the polytope and torus product $\mathbb{P}^{3}\times\mathbb{T}^{4}$ ((a) groundtruth, (b) output of our method, (c) sample from the uniform distribution). (d) Pairwise comparison of the polytope components of the target distribution over cyclic peptides and samples from the reflected method. Additional results and similar plots for the log barrier method are postponed to Appendix J.  Table 1:  MMD metrics between synthetic data of varying dimension and samples from trained constrained diffusion models. Bold indicates best results.  Space | Dimension | Log-barrier | Reflected
---|---|---|---
$\Delta^{d}$ | 2 | $0.05$ | $\bm{0.02}$
3 | $0.10$ | $\bm{0.09}$
10 | $0.17$ | $\bm{0.07}$
$[-1,1]^{d}$ | 2 | $0.03$ | $\bm{0.03}$
3 | ${0.05}$ | $0.06$
10 | $0.26$ | $\bm{0.06}$

###  5.2 Learning constrained SPD matrices for robotic arms modelling

Accurately determining and controlling the movement and exerted forces of a
robotic platform is a fundamental problem in many real-world robotics
applications (Yoshikawa, 1985). A kinetostatic descriptor that is commonly
used to quantify the ability of a robotic arm to move and apply forces along
certain dimensions is the so-called manipulability ellipsoid
$E\in\mathbb{R}^{d}$ (Yoshikawa, 1985) which is naturally described as a
symmetric positive definite (SPD) matrix $\mathrm{M}\in\mathbb{R}^{d\times d}$
(Jaquier et al., 2021). The manifold of such $d\times d$ SPD matrices, denoted
as $S_{++}^{d}$, is defined as the set of matrices $\\{x^{\top}\mathrm{M}x\geq
0,\ x\in\mathbb{R}^{d}:\mathrm{M}\in\mathbb{R}^{d\times d}\\}$.

In many practical settings, it may be desirable to constrain the volume of $E$
to retain flexibility or limit the magnitude of an exerted force, which can be
expressed as an upper bound on the trace of $\mathrm{M}$, i.e. as the
inequality constraint $\sum_{i=1}^{d}\mathrm{M}_{ii}<C$ with $C>0$.
Constraining the rest of the entries of the matrix to ensure it is SPD is non-
trivial.

Alternatively, we can parameterise the SPD matrices via their Cholesky
decomposition. Each SPD matrix has a unique decomposition of the form
$\mathrm{M}=\mathrm{L}\mathrm{L}^{\top}$, where $\mathrm{L}$ is a lower
triangular matrix with strictly positive diagonal (Golub & Van Loan, 2013,
p.143). Constraining the entries of this matrix simply requires ensuring the
diagonal is positive. The trace of the SPD matrix is given by
$\mathrm{Tr}(\mathrm{M})=\mathrm{Tr}(\mathrm{L}\mathrm{L}^{\top})=\sum_{ij}\mathrm{L}_{ij}^{2}$,
and results in the constraint on the entries of $\mathrm{L}$ to live in a ball
of radius $C$.

We additionally model the two-dimensional position of the arm. In summary, the
space over which we parameterise the diffusion models is defined as
$\\{\mathrm{L}\in\mathbb{R}^{d(d+1)/2}:\mathrm{L}_{i,i}>0,\sum_{i,j}\mathrm{L}_{i,j}^{2}<C\\}\times\mathbb{R}^{2}$.
Under the Euclidean metric we can apply both our log-barrier and reflected
approaches. The positive diagonal (linear) constraint is handled similarly as
in the polytope setting. The reflection on the ball boundary is defined and
illustrated on Figure 6. In Figure 9, we plot the joint distribution over the
SPD matrices (represented as ellipsoids) and their positions. We qualitatively
observe that the reflected method is able to fit the joint data distribution.
Quantitatively, the reflected method achieves an MMD of 0.161 as opposed to
0.247 for the log-barrier one.

(a) Data

(b) Reflected

(c) Uniform

Figure 9: Samples of joint distributions over SPD matrices (illustrated as
ellipsoids) and their location, see Appendix J for additional results and pair
plots.

###  5.3 Modelling protein backbones with anchored endpoints

Modelling the conformational ensembles of proteins is an important task in the
field of molecular biology, particularly in the context of bioengineering and
drug design. In many data-scarce settings, it may be unnecessary to model the
structure of an entire protein, as researchers are primarily interested in
specific functional sites. However, only generating conformations for a subset
of a protein requires positional constraints on its endpoints to ensure that
it can still be accommodated by the remaining scaffold. More generally, this
is the problem of how to model a chain with fixed end points.

We follow the framework outlined in Han & Rudolph (2006) for parameterising an
end-points-constrained chain with $d$ links as the product of a convex
polytope $\mathbb{P}$ and torus $\mathbb{T}$:
$\mathbb{P}^{(d-2)}\times\mathbb{T}^{(d-1)}$. The essential idea in this
parameterisation is to fix one end point as an “anchor” and model the chain as
the series of triangles formed by the anchor and each pair of sequential
joints in the chain. A point in the polytope corresponds to the lengths of the
diagonals of these triangles, and the point in the torus to the angles between
each subsequent triangle. See Section H.1 for more details.

Using this framework, we model the (smoothed) energetic landscape of the
conformations of cyclic peptides of length $6$, i.e. circular polypeptide
chains with coinciding endpoints, as a distribution over
$\mathbb{P}^{3}\times\mathbb{T}^{4}$. We then draw samples on this product
manifold from a Boltzmann distribution parametrised by molecular dynamics
potentials (using OpenMM (Eastman et al., 2017) and AMBER (Hornak et al.,
2006)), which we subsequently divide into training and evaluation datasets,
see Section H.2. The space on which we effectively parameterise the reflected
diffusion model is thus the product space
$\mathbb{P}^{3}\times\mathbb{T}^{4}$, leveraging the methodology introduced in
Section 3.2 for the polytope $\mathbb{P}^{3}$ component and in De Bortoli et
al. (2022) for the torus $\mathbb{T}^{4}$ part.

Figure 8 shows samples from the data distribution, along with generated
samples from our learnt reflected diffusion model. From (d) we can see that on
the polytope component we successfully learn the modes of the distribution.
This translates into the visualisation of the loops in (a-c) where we can see
a pronounced gap in the conformational space occupied by samples from the data
distribution and the reflected method that is not apparent for uniform
samples.

##  6 Discussion

Learning complex distributions supported on constrained spaces is of
importance in several fields of the natural sciences, such as computational
biology, see Han & Rudolph (2006). In this work we extend continuous diffusion
models to this setting, proposing two complementary approaches. One
constructing a reflected Brownian motion as the noising process, and another
one based on barrier methods. For both approaches, we derive the time-reversal
formula, propose discretisation schemes and extend the score matching toolbox.
In practice, we show that reflected type methods, while enjoying less
theoretical guarantees than their log-barrier counterparts, are more
efficient. We demonstrate the efficiency of our method in two real-world
scenarios: conformational modelling and robotic arms.

We conclude by highlighting important future directions. First, the
computational cost of the performing the reflection when discretising the
reflected Brownian motion is high. Finding numerically efficient
approximations of the reflected process is therefore necessary to extend this
methodology to very high dimensional settings. Second, the retraction used in
place of the exponential map for the barrier method leads to a high number of
discretisation steps to ensure a good approximation. Designing faster forward
process for the log-barrier method is key to target more complex
distributions. Finally, quantum Bayesian applications (Lukens et al., 2020)
are a challenging research direction requiring the addition of _sparsity_
constraints to our framework.

## References

  * Ablin & Peyré (2022) Ablin, P. and Peyré, G.  Fast and accurate optimization on the orthogonal manifold without retraction.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 5636–5657. PMLR, 2022.
  * Bossy et al. (2004) Bossy, M., Gobet, E., and Talay, D.  A symmetrized euler scheme for an efficient approximation of reflected diffusions.  _Journal of applied probability_ , 41(3):877–889, 2004.
  * Boyd et al. (2004) Boyd, S., Boyd, S. P., and Vandenberghe, L.  _Convex optimization_.  Cambridge university press, 2004.
  * Bradbury et al. (2018) Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q.  JAX: composable transformations of Python+NumPy programs, 2018\.  URL http://github.com/google/jax.
  * Brezis & Brézis (2011) Brezis, H. and Brézis, H.  _Functional analysis, Sobolev spaces and partial differential equations_ , volume 2.  Springer, 2011.
  * Burdzy et al. (2004) Burdzy, K., Chen, Z.-Q., and Sylvester, J.  The heat equation and reflected brownian motion in time-dependent domains.  _The Annals of Probability_ , 32(1B):775–804, 2004.
  * Cattiaux et al. (2021) Cattiaux, P., Conforti, G., Gentil, I., and Léonard, C.  Time reversal of diffusion processes under a finite entropy condition.  _arXiv preprint arXiv:2104.07708_ , 2021.
  * Claude J. P. Bélisle et al. (1993) Claude J. P. Bélisle, Romeijn, H. E., and Smith, R. L.  Hit-and-Run Algorithms for Generating Multivariate Distributions.  _Mathematics of Operations Research_ , 18(2):255–266, 1993.  ISSN 0364765X, 15265471.  URL http://www.jstor.org/stable/3690278.
  * Corso et al. (2022) Corso, G., Stärk, H., Jing, B., Barzilay, R., and Jaakkola, T.  DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking, October 2022.  URL http://arxiv.org/abs/2210.01776.
  * De Bortoli et al. (2022) De Bortoli, V., Mathieu, E., Hutchinson, M., Thornton, J., Teh, Y. W., and Doucet, A.  Riemannian score-based generative modeling, 2022.
  * Dhariwal & Nichol (2021) Dhariwal, P. and Nichol, A.  Diffusion models beat GAN on image synthesis.  _arXiv preprint arXiv:2105.05233_ , 2021.
  * Dougherty et al. (2019) Dougherty, P. G., Sahni, A., and Pei, D.  Understanding cell penetration of cyclic peptides.  _Chemical Reviews_ , 119(17):10241–10287, 2019\.
  * Eastman et al. (2017) Eastman, P., Swails, J., Chodera, J. D., McGibbon, R. T., Zhao, Y., Beauchamp, K. A., Wang, L.-P., Simmonett, A. C., Harrigan, M. P., Stern, C. D., et al.  Openmm 7: Rapid development of high performance algorithms for molecular dynamics.  _PLoS computational biology_ , 13(7):e1005659, 2017.
  * Gatmiry & Vempala (2022) Gatmiry, K. and Vempala, S. S.  Convergence of the riemannian langevin algorithm.  _arXiv preprint arXiv:2204.10818_ , 2022.
  * Golub & Van Loan (2013) Golub, G. H. and Van Loan, C. F.  _Matrix computations_.  JHU press, 2013.
  * Gretton et al. (2012) Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., and Smola, A.  A kernel two-sample test.  _Journal of Machine Learning Research_ , 13(null):723–773, March 2012.  ISSN 1532-4435.
  * Han & Rudolph (2006) Han, L. and Rudolph, L.  Inverse kinematics for a serial chain with joints under distance constraints.  In _Robotics: Science and systems_ , 2006.
  * Harrison & Williams (1987) Harrison, J. M. and Williams, R. J.  Brownian models of open queueing networks with homogeneous customer populations.  _Stochastics: An International Journal of Probability and Stochastic Processes_ , 22(2):77–115, 1987.
  * Haussmann & Pardoux (1986) Haussmann, U. G. and Pardoux, E.  Time reversal of diffusions.  _The Annals of Probability_ , pp. 1188–1205, 1986.
  * Heirendt et al. (2019) Heirendt, L., Arreckx, S., Pfau, T., Mendoza, S. N., Richelle, A., Heinken, A., Haraldsdóttir, H. S., Wachowiak, J., Keating, S. M., Vlasov, V., et al.  Creation and analysis of biochemical constraint-based models using the cobra toolbox v. 3.0.  _Nature protocols_ , 14(3):639–702, 2019.
  * Ho et al. (2020) Ho, J., Jain, A., and Abbeel, P.  Denoising diffusion probabilistic models.  _Advances in Neural Information Processing Systems_ , 2020.
  * Hornak et al. (2006) Hornak, V., Abel, R., Okur, A., Strockbine, B., Roitberg, A., and Simmerling, C.  Comparison of multiple amber force fields and development of improved protein backbone parameters.  _Proteins: Structure, Function, and Bioinformatics_ , 65(3):712–725, 2006.
  * Hsu (2002) Hsu, E. P.  _Stochastic analysis on manifolds_.  Number 38. American Mathematical Soc., 2002.
  * Huang et al. (2021) Huang, C.-W., Lim, J. H., and Courville, A. C.  A variational perspective on diffusion-based generative models and score matching.  _Advances in Neural Information Processing Systems_ , 34:22863–22876, 2021.
  * Huang et al. (2022) Huang, C.-W., Aghajohari, M., Bose, A. J., Panangaden, P., and Courville, A.  Riemannian Diffusion Models, August 2022.
  * Hyvärinen (2005) Hyvärinen, A.  Estimation of non-normalized statistical models by score matching.  _Journal of Machine Learning Research_ , 6(4), 2005.
  * Ingraham et al. (2022) Ingraham, J., Baranov, M., Costello, Z., Frappier, V., Ismail, A., Tie, S., Wang, W., Xue, V., Obermeyer, F., Beam, A., et al.  Illuminating protein space with a programmable generative model.  _bioRxiv_ , pp. 2022–12, 2022.
  * Jaquier et al. (2021) Jaquier, N., Rozo, L., Caldwell, D. G., and Calinon, S.  Geometry-aware manipulability learning, tracking, and transfer.  _The International Journal of Robotics Research_ , 40(2-3):624–650, 2021.
  * Jing et al. (2022) Jing, B., Corso, G., Chang, J., Barzilay, R., and Jaakkola, T.  Torsional Diffusion for Molecular Conformer Generation, June 2022\.  URL http://arxiv.org/abs/2206.01729.
  * Jørgensen (1975) Jørgensen, E.  The central limit problem for geodesic random walks.  _Zeitschrift für Wahrscheinlichkeitstheorie und verwandte Gebiete_ , 32(1-2):1–64, 1975.
  * Kang & Ramanan (2017) Kang, W. and Ramanan, K.  On the submartingale problem for reflected diffusions in domains with piecewise smooth boundaries.  2017\.
  * Kannan & Narayanan (2009) Kannan, R. and Narayanan, H.  Random walks on polytopes and an affine interior point method for linear programming.  In _Proceedings of the Forty-First Annual ACM Symposium on Theory of Computing_ , STOC ’09, pp. 561–570, New York, NY, USA, 2009. Association for Computing Machinery.  ISBN 9781605585062.  doi: 10.1145/1536414.1536491.  URL https://doi.org/10.1145/1536414.1536491.
  * Kingma & Ba (2014) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization, 2014.  URL http://arxiv.org/abs/1412.6980.
  * Kook et al. (2022) Kook, Y., Lee, Y. T., Shen, R., and Vempala, S. S.  Sampling with riemannian hamiltonian monte carlo in a constrained space.  _arXiv preprint arXiv:2202.01908_ , 2022.
  * Krivov & Karplus (2004) Krivov, S. V. and Karplus, M.  Hidden complexity of free energy surfaces for peptide (protein) folding.  _Proceedings of the National Academy of Sciences_ , 101(41):14766–14770, 2004.
  * Lane (2023) Lane, T. J.  Protein structure prediction has reached the single-structure frontier.  _Nature Methods_ , pp. 1–4, 2023.
  * Leach et al. (2022) Leach, A., Schmon, S. M., Degiacomi, M. T., and Willcocks, C. G.  Denoising Diffusion Probabilistic Models on SO(3) for Rotational Alignment.  _ICLR 2022 Workshop on Geometrical and Topological Representation Learning_ , pp. 8, 2022.
  * Lee (2013) Lee, J. M.  Smooth manifolds.  In _Introduction to Smooth Manifolds_ , pp. 1–31. Springer, 2013\.
  * Lee & Vempala (2017) Lee, Y. T. and Vempala, S. S.  Geodesic walks in polytopes.  In _Proceedings of the 49th Annual ACM SIGACT Symposium on Theory of Computing_ , pp. 927–940, Montreal Canada, June 2017. ACM.  ISBN 978-1-4503-4528-6.  doi: 10.1145/3055399.3055416.
  * Lee & Vempala (2018) Lee, Y. T. and Vempala, S. S.  Convergence rate of riemannian hamiltonian monte carlo and faster polytope volume computation.  In _Proceedings of the 50th Annual ACM SIGACT Symposium on Theory of Computing_ , pp. 1115–1121, 2018.
  * Lewis et al. (2012) Lewis, N., Nagarajan, H., and Palsson, B.  Constraining the metabolic genotype-phenotype relationship using a phylogeny of in silico methods.  _Nature Reviews Microbiology_ , 10(4):291–305, April 2012.  ISSN 1740-1526.  doi: 10.1038/nrmicro2737.
  * Lions & Sznitman (1984) Lions, P.-L. and Sznitman, A.-S.  Stochastic differential equations with reflecting boundary conditions.  _Communications on pure and applied Mathematics_ , 37(4):511–537, 1984.
  * Loper (2020) Loper, J.  Uniform ergodicity for brownian motion in a bounded convex set.  _Journal of Theoretical Probability_ , 33(1):22–35, 2020.
  * Lovász & Vempala (2006) Lovász, L. and Vempala, S.  Hit-and-run from a corner.  _SIAM Journal on Computing_ , 35(4):985–1005, 2006.  doi: 10.1137/S009753970544727X.
  * Lukens et al. (2020) Lukens, J. M., Law, K. J., Jasra, A., and Lougovski, P.  A practical and efficient approach for bayesian quantum state estimation.  _New Journal of Physics_ , 22(6):063038, 2020\.
  * Miolane et al. (2020) Miolane, N., Guigui, N., Brigant, A. L., Mathe, J., Hou, B., Thanwerdas, Y., Heyder, S., Peltre, O., Koep, N., Zaatiti, H., Hajri, H., Cabanes, Y., Gerald, T., Chauchat, P., Shewmake, C., Brooks, D., Kainz, B., Donnat, C., Holmes, S., and Pennec, X.  Geomstats: A python package for riemannian geometry in machine learning.  _Journal of Machine Learning Research_ , 21(223):1–9, 2020.  URL http://jmlr.org/papers/v21/19-027.html.
  * Morris (2002) Morris, B. J.  Improved bounds for sampling contingency tables.  _Random Structures & Algorithms_, 21(2):135–146, 2002.  doi: https://doi.org/10.1002/rsa.10049.  URL https://onlinelibrary.wiley.com/doi/abs/10.1002/rsa.10049.
  * Nesterov et al. (2018) Nesterov, Y. et al.  _Lectures on convex optimization_ , volume 137.  Springer, 2018.
  * Noble et al. (2022) Noble, M., De Bortoli, V., and Durmus, A.  Barrier Hamiltonian Monte Carlo, October 2022.
  * Pacchiarotti et al. (1998) Pacchiarotti, B., Costantini, C., and Sartoretto, F.  Numerical approximation for functionals of reflecting diffusion processes.  _SIAM Journal on Applied Mathematics_ , 58(1):73–102, 1998.
  * Petit (1997) Petit, F.  Time reversal and reflected diffusions.  _Stochastic Processes and their Applications_ , 69(1):25–53, July 1997.  ISSN 03044149.  doi: 10.1016/S0304-4149(97)00035-5.
  * Revuz & Yor (2013) Revuz, D. and Yor, M.  _Continuous martingales and Brownian motion_ , volume 293.  Springer Science & Business Media, 2013.
  * Rotkiewicz & Skolnick (2008) Rotkiewicz, P. and Skolnick, J.  Fast procedure for reconstruction of full-atom protein models from reduced representations.  _Journal of computational chemistry_ , 29(9):1460–1465, 2008.
  * Saloff-Coste (1994) Saloff-Coste, L.  Precise estimates on the rate at which certain diffusions tend to equilibrium.  _Mathematische Zeitschrift_ , 217(1):641–677, 1994.
  * Shkolnikov & Karatzas (2013) Shkolnikov, M. and Karatzas, I.  Time-reversal of reflected Brownian motions in the orthant, July 2013\.  URL http://arxiv.org/abs/1307.4422.
  * Skorokhod (1961) Skorokhod, A. V.  Stochastic equations for diffusion processes in a bounded region.  _Theory of Probability & Its Applications_, 6(3):264–274, 1961.
  * Smith (1984) Smith, R. L.  Efficient Monte Carlo Procedures for Generating Points Uniformly Distributed over Bounded Regions.  _Operations Research_ , 32(6):1296–1308, December 1984.  doi: 10.1287/opre.32.6.1296.  URL https://ideas.repec.org/a/inm/oropre/v32y1984i6p1296-1308.html.
  * Sohl-Dickstein et al. (2015) Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S.  Deep unsupervised learning using nonequilibrium thermodynamics.  In _International Conference on Machine Learning_ , pp. 2256–2265. PMLR, 2015.
  * Song & Ermon (2019) Song, Y. and Ermon, S.  Generative modeling by estimating gradients of the data distribution.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Song et al. (2021) Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B.  Score-based generative modeling through stochastic differential equations.  In _International Conference on Learning Representations_ , 2021.
  * Thiele et al. (2013) Thiele, I., Swainston, N., Fleming, R., Hoppe, A., Sahoo, S., Aurich, M., Haraldsdottir, H., Mo, M., Rolfsson, O., Stobbe, M., Thorleifsson, S., Agren, R., Bölling, C., Bordel, S., Chavali, A., Dobson, P., Dunn, W., Endler, L., Hala, D., Hucka, M., Hull, D., Jameson, D., Jamshidi, N., Jonsson, J., Juty, N., Keating, S., Nookaew, I., Le Novère, N., Malys, N., Mazein, A., Papin, J., Price, N., Selkov, E., Sigurdsson, M., Simeonidis, E., Sonnenschein, N., Smallbone, K., Sorokin, A., van Beek, J., Weichart, D., Goryanin, I., Nielsen, J., Westerhoff, H., Kell, D., Mendes, P., and Palsson, B.  A community-driven global reconstruction of human metabolism.  _Nature Biotechnology_ , March 2013.  ISSN 1087-0156.  doi: 10.1038/nbt.2488.
  * Trippe et al. (2022) Trippe, B. L., Yim, J., Tischer, D., Broderick, T., Baker, D., Barzilay, R., and Jaakkola, T.  Diffusion probabilistic modeling of protein backbones in 3D for the motif-scaffolding problem, June 2022.
  * Urain et al. (2022) Urain, J., Funk, N., Chalvatzaki, G., and Peters, J.  Se (3)-diffusionfields: Learning cost functions for joint grasp and motion optimization through diffusion.  _arXiv preprint arXiv:2209.03855_ , 2022.
  * Vincent (2011) Vincent, P.  A connection between score matching and denoising autoencoders.  _Neural Computation_ , 23(7):1661–1674, 2011\.
  * Watson et al. (2022) Watson, J. L., Juergens, D., Bennett, N. R., Trippe, B. L., Yim, J., Eisenach, H. E., Ahern, W., Borst, A. J., Ragotte, R. J., Milles, L. F., Wicky, B. I. M., Hanikel, N., Pellock, S. J., Courbet, A., Sheffler, W., Wang, J., Venkatesh, P., Sappington, I., Torres, S. V., Lauko, A., Bortoli, V. D., Mathieu, E., Barzilay, R., Jaakkola, T. S., DiMaio, F., Baek, M., and Baker, D.  Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models, December 2022.
  * Williams (1987) Williams, R. J.  Reflected Brownian motion with skew symmetric data in a polyhedral domain.  _Probability Theory and Related Fields_ , 75(4):459–485, August 1987.  ISSN 0178-8051, 1432-2064.  doi: 10.1007/BF00320328.
  * Wu et al. (2022) Wu, K. E., Yang, K. K., Berg, R. v. d., Zou, J. Y., Lu, A. X., and Amini, A. P.  Protein structure generation via folding diffusion.  _arXiv preprint arXiv:2209.15611_ , 2022.
  * Yoshikawa (1985) Yoshikawa, T.  Manipulability of robotic mechanisms.  _The international journal of Robotics Research_ , 4(2):3–9, 1985.

##  Appendix A Introduction

In this supplementary, we first recall some key concepts from Riemannian
geometry in Appendix B. In Appendix C we remind the expression of the Brownian
motion in local coordinates, and use it to derive the SDE that the Langevin
dynamics satisfy. in Appendix D we derive the implicit score matching loss. We
give details about the likelihood evaluation in Appendix E. Then in Appendix G
we prove the time-reversal formula of reflected Brownian motion. In Appendix H
we give some background on the conformational modelling of proteins backbone
for the experiment in Section 5.3. In Appendix I we detail the geometrical
constraint arising from the configurational robotics arm modelling experiment
from Section 5.2. Additional results, training and miscalleneous experimental
details are reported in Appendix J.

##  Appendix B Manifold concepts

For readers unfamiliar with Riemannian geometry here we give a brief overview
of some of the key concepts. This is not a technical introduction, but a
conceptual one for the understanding of terms. For a technical introduction,
we refer readers to Lee (2013). A Riemannian manifold is a tuple
$(\mathcal{M},\mathfrak{g})$ with $\mathcal{M}$ a smooth manifold and
$\mathfrak{g}$ a metric which defines an inner product on tangent spaces.

Figure 10: Example charts of the 2D manifold $\mathcal{S}^{2}$.

A _smooth manifold_ is a topological space which locally can be described by
Euclidean space. It is characterised by a family of _charts_
$\\{U\subset\mathcal{M},\phi:U\to\mathbb{R}^{d}\\}$, homeomorphic mappings
between subsets of the manifold and Euclidean space. The collection of charts
must cover the whole manifold. For the manifolds to be smooth the charts must
be smooth, and the transition between charts where their domains overlap must
also be smooth.

The _metric_ on a Riemannian manifold gives the manifold a notion of distance
and curvature. The same underlying smooth manifold with different metrics can
look wildly different. The metric is defined as a smooth choice of positive
definite inner product on each of the tangent spaces of the manifold. That is
we have at every point a symmetric bilinear map

$\mathfrak{g}(p):\mathrm{T}_{p}\mathcal{M}\times\mathrm{T}_{p}\mathcal{M}\to\mathbb{R}$

The tangent space of a point on a manifold is the generalisation of the notion
of tangent planes and can be thought of as the space of derivatives of scalar
functions on the manifold at that point. The collection of all tangent spaces
is written as
$\mathrm{T}\mathcal{M}=\bigcup_{p\in\mathcal{M}}\mathrm{T}_{p}\mathcal{M}$ and
is called the tangent bundle. It is also manifold. Vector fields on manifolds
are defined by making a choice of tangent vector at every point on the
manifold in a smooth fashion. The space of vector fields is written at
$\Gamma\left(\mathrm{T}\mathcal{M}\right)$ and is more technically the space
of _sections_ of the tangent bundle.

One thing that the metric itself does not immediately define is how different
tangent spaces at points on the manifold relate to one another. For this, we
need additionally the concept of a _connection_. A connection is a map that
takes two vector fields and produces a derivative of the first with respect to
the second, that is a function
$\nabla:\Gamma(\mathrm{T}\mathcal{M})\times\Gamma(\mathrm{T}\mathcal{M})\to\Gamma(\mathrm{T}\mathcal{M})$
and it typically written as $\nabla(X,Y)=\nabla_{X}Y$. Such a connection must
for $X,Y,Z\in\Gamma(\mathrm{T}\mathcal{M})$ and smooth functions on the
manifold $a,b:\mathcal{M}\to\mathbb{R}$ obey the following conditions:

  1. (a)

$\nabla_{aX+bY}Z=a\nabla_{X}Z+b\nabla_{Y}Z$,

  2. (b)

$\nabla_{X}(Y+Z)=\nabla_{X}Y+\nabla_{X}Z$,

  3. (c)

$\nabla_{X}(aY)=\partial_{X}aY+a\nabla_{X}Y$,

where $\partial_{X}aY$ is the regular _directional derivative_ of $aY$ in the
direction $X$. These conditions ensure the connection is a well-defined
derivative. In what follows, we focus on a specific choice of connection which
is particularly well-suited for Riemannian manifolds.

On a given manifold, there are infinitely many connections. Fortunately, there
is a natural choice, called the _Levi-Cevita_ connection if we impose two
additional conditions:

  1. (a)

$X\cdot(\mathfrak{g}(Y,Z))=\mathfrak{g}(\nabla_{X}Y,Z)+\mathfrak{g}(Y,\nabla_{X}Z)$,

  2. (b)

$\left[X,Y\right]=\nabla_{X}Y-\nabla_{Y}X$,

where $\left[\cdot,\cdot\right]$ is the Lie bracket. The first condition
ensures that the metric is preserved by the connection. That is to say, the
_parallel transport_ (to be defined shortly) using the connection leaves inner
products unchanged on the manifold. The second ensures the connection is
_torsion-free_. The change in tangent space along a geodesic (again to be
defined shortly) can be described in two parts, the _curvature_ , how the
tangent space rotates perpendicular to the direction of travel, and the
torsion, how the tangent space rotates around the axis of the direction of
travel. The curvature of the connection is uniquely fixed by the other 4
conditions (the well-defined derivative and preservation of the metric). The
torsion however is not fixed. By requiring it to be zero we ensure a unique
connection. The requirement of zero torsion also has implications for ensuring
integrability on the manifold.

With the metric and the Levi-Cevita connection in hand, we can define a number
of key concepts.

Figure 11: Parallel transport of vectors along a path on the sphere.

We say that a vector field $X$ is _parallel_ to a curve
$\gamma:(0,1)\to\mathcal{M}$ if

$\nabla_{\gamma^{\prime}}X=0$

where $\gamma^{\prime}:(0,1)\to\mathrm{T}_{\gamma(t)}\mathcal{M}$ is the
derivative of the path. For two points on the manifold $p,q\in\mathcal{M}$ and
a curve between them, $\gamma$, $\gamma(0)=p$, $\gamma(1)=q$, for an initial
vector $X_{0}\in\mathrm{T}_{p}\mathcal{M}$ there is a unique vector field $X$
that is parallel to $\gamma$ such that $X(p)=X_{0}$. This induces a map
between the tangent spaces at $p$ and $q$

$\tau_{\gamma}:\mathrm{T}_{p}\mathcal{M}\to\mathrm{T}_{q}\mathcal{M}$

This map is referred as the _parallel transport_ of tangent vectors between
$p$ and $q$, and this satisfies the condition that for
$\bm{v},\bm{u}\in\mathrm{T}_{p}\mathcal{M}$

$\mathfrak{g}(p)(\bm{v},\bm{u})=\mathfrak{g}(q)(\tau_{\gamma}(\bm{v}),\tau_{\gamma}(\bm{u})).$

A _geodesic_ on a manifold is the unique path on the manifold
$\gamma:(0,1)\to\mathcal{M}$ such that
$\nabla_{\gamma^{\prime}}\gamma^{\prime}=0$. It is also the shortest path
between two points on a manifold in the sense that

$\textstyle{L(\gamma)=\int_{0}^{1}\sqrt{\mathfrak{g}(\gamma(t))(\gamma^{\prime}(t),\gamma(t))},}$

is minimal out of any path between the start and end of the geodesic.
Geodesics give the notion of ‘straight lines’ on manifolds. We define the
_exponential map_ on a manifold as the mapping between an element of the
tangent space at point $p$, $\bm{v}\in T_{p}\mathcal{M}$ and the endpoint of
the unique geodesic $\gamma$ with $\gamma(0)=p$ and
$\gamma^{\prime}(0)=\bm{v}$. In the tangent space of a manifold, we require
the notion of a reflection in order to reflect geodesics off of boundary
constraints. If at a point in the manifold with have $\bm{v}\in
T_{p}\mathcal{M}$ and a unit vector normal to the constraint $\bm{n}\in
T_{p}\mathcal{M}$, then the reflection of this vector is given by
$\bm{v}^{\prime}=\bm{v}-2\mathfrak{g}(\bm{v},\bm{n})\bm{n}$.

##  Appendix C Brownian motion in local coordinates

We consider a smooth function $f\in\mathrm{C}^{\infty}(\mathcal{M})$. The
Laplace-Beltrami operator $\Delta_{\mathcal{M}}$ is given by
$\Delta_{\mathcal{M}}(f)=\mathrm{div}(\mathrm{grad}(f)))$. In local
coordinates we have

$\textstyle{\mathrm{div}(X)=(\det(\mathfrak{g})^{-1/2})\sum_{i=1}^{d}\partial_{i}(\det(\mathfrak{g})^{1/2}X_{i}),}$
(16)

as well as

$\mathrm{grad}(f)=\mathfrak{g}^{-1}\nabla f.$ (17)

Therefore, the Laplace-Beltrami operator is given by

$\textstyle{\Delta_{\mathcal{M}}(f)=\sum_{i,j=1}^{d}\mathfrak{g}_{i,j}^{-1}\partial_{i,j}f+(\det(\mathfrak{g})^{-1/2})\sum_{i,j=1}^{d}\partial_{i}(\det(\mathfrak{g})^{1/2}\mathfrak{g}_{i,j}^{-1})\partial_{j}f.}$
(18)

Therefore, in local coordinates the infinitesimal generator associated with
the Laplace-Beltrami operator is given by

$\textstyle{\mathcal{A}(f)=\sum_{i,j=1}^{d}\mathfrak{g}_{i,j}^{-1}\partial_{i,j}f+\langle
b^{i},\nabla f\rangle,}$ (19)

with

$\textstyle{b^{i}=(\det(\mathfrak{g})^{-1/2})\sum_{j=1}^{d}\partial_{j}(\det(\mathfrak{g})^{1/2}\mathfrak{g}_{i,j}^{-1}).}$
(20)

Therefore, the dual operator associated with $\mathcal{A}$ is given by

$\textstyle{\mathcal{A}^{\star}(f)=\sum_{i,j=1}^{d}\partial_{i,j}(\mathfrak{g}_{i,j}^{-1}f)-\sum_{i=1}^{d}\partial_{i}(b^{i}f).}$
(21)

Note that by letting $f=\det(\mathfrak{g})^{1/2}$ we get that
$\mathcal{A}^{\star}(f)=0$ and therefore we recover that
$p\propto\det(\mathfrak{g})$ is the invariant distribution of the Brownian
motion.

#### Langevin dynamics on $\mathcal{M}$.

We know that the Brownian motion targets $\det(\mathfrak{g})^{1/2}$. Therefore
in order to correct and sample from the uniform distribution we consider the
Langevin dynamics

$\mathrm{d}\mathbf{X}_{t}=-\mathrm{grad}\log(\det(\mathfrak{g})^{1/2})(\mathbf{X}_{t})\mathrm{d}t+\sqrt{2}\mathrm{d}\mathbf{B}_{t}^{\mathcal{M}}.$
(22)

Note that in the previous equation $\mathrm{grad}$ and
$\mathbf{B}_{t}^{\mathcal{M}}$ are defined w.r.t. the metric of the manifold.
In local coordinates we have

$\mathrm{d}\mathbf{X}_{t}=\\{b-\mathrm{grad}\log(\det(\mathfrak{g})^{1/2})\\}(\mathbf{X}_{t})\mathrm{d}t+\sqrt{2}\mathfrak{g}(\mathbf{X}_{t})^{-1/2}\mathrm{d}\mathbf{B}_{t}.$
(23)

where $b=\\{b^{i}\\}_{i=1}^{d}$ is given in (20). In addition, we have

$\mathrm{grad}\log(\det(\mathfrak{g})^{1/2})=\det(\mathfrak{g})^{-1/2}\mathfrak{g}^{-1}\nabla\det(\mathfrak{g}).$
(24)

Using (20) we have

$b^{i}=\textstyle{(\det(\mathfrak{g})^{-1/2})\sum_{j=1}^{d}\partial_{j}(\det(\mathfrak{g})^{1/2}\mathfrak{g}_{i,j}^{-1})}=\textstyle{\sum_{j=1}^{d}\partial_{j}\mathfrak{g}_{i,j}^{-1}+\mathrm{grad}\log(\det(\mathfrak{g})^{1/2})_{i}}$
(25)

This can also be rewritten as

$\mathrm{div}_{\mathcal{M}}(\mathfrak{g}^{-1})=\textstyle{\mu+\mathrm{grad}\log(\det(\mathfrak{g})^{1/2}),}$
(26)

with

$\textstyle{\mu_{i}=\sum_{j=1}^{d}\partial_{j}\mathfrak{g}_{i,j}^{-1}.}$ (27)

Combining this result and (24) we get that (23) can be rewritten as

$\mathrm{d}\mathbf{X}_{t}=\mu(\mathbf{X}_{t})+\sqrt{2}\mathrm{d}\mathbf{B}_{t}.$
(28)

Note that (up to a factor $2$) this is the same SDE as the one considered in
Lee & Vempala (2017).

##  Appendix D Implicit Score Matching Loss

Using the divergence theorem, we have

$\displaystyle\textstyle{(1/2)\int_{\mathcal{M}}\|\mathbf{s}_{\theta}(x)-\nabla\log
p_{t}(x)\|^{2}p_{t}(x)\mathrm{d}\mu(x)}$ (29)
$\displaystyle=\textstyle{(1/2)\int_{\mathcal{M}}\|\mathbf{s}_{\theta}(x)\|^{2}p_{t}(x)\mathrm{d}\mu(x)-\int_{\mathcal{M}}\langle\mathbf{s}_{\theta}(x),\nabla\log
p_{t}(x)\rangle p_{t}(x)\mathrm{d}\mu(x)+(1/2)\int_{\mathcal{M}}\|\nabla\log
p_{t}(x)\|^{2}p_{t}(x)\mathrm{d}\mu(x)}$ (30)
$\displaystyle=\textstyle{(1/2)\int_{\mathcal{M}}\|\mathbf{s}_{\theta}(x)\|^{2}p_{t}(x)\mathrm{d}\mu(x)-\int_{\partial\mathcal{M}}\langle\mathbf{s}_{\theta}(x),\mathbf{n}\rangle
p_{t}(x)\mathrm{d}\nu(x)}$ (31)
$\displaystyle\qquad\textstyle{+\int_{\mathcal{M}}\mathrm{div}(\mathbf{s}_{\theta})(x)p_{t}(x)\mathrm{d}\mu(x)+(1/2)\int_{\mathcal{M}}\|\nabla\log
p_{t}(x)\|^{2}p_{t}(x)\mathrm{d}\mu(x)}.$ (32)

Using that $\mathbf{s}_{\theta}(x)=0$ for all $x\in\partial\mathcal{M}$, we
get that

$\displaystyle\textstyle{(1/2)\int_{\mathcal{M}}\|\mathbf{s}_{\theta}(x)-\nabla\log
p_{t}(x)\|^{2}p_{t}(x)\mathrm{d}\mu(x)}$ (33)
$\displaystyle\qquad\qquad=\textstyle{(1/2)\int_{\mathcal{M}}\|\mathbf{s}_{\theta}(x)\|^{2}p_{t}(x)\mathrm{d}\mu(x)}\textstyle{+\int_{\mathcal{M}}\mathrm{div}(\mathbf{s}_{\theta})(x)p_{t}(x)\mathrm{d}\mu(x)+(1/2)\int_{\mathcal{M}}\|\nabla\log
p_{t}(x)\|^{2}p_{t}(x)\mathrm{d}\mu(x)},$ (34)

which concludes the proof.

##  Appendix E Likelihood evaluation

One main advantage of constructing a continuous noising process, is that
similarly to Song et al. (2021), we can evaluate the model’s likelihood via
the following _probability flow_ Ordinary Differential Equation (ODE). In
particular for the Langevin dynamics Equation 6 which we recall

$\textstyle{\mathrm{d}\mathbf{X}_{t}=\tfrac{1}{2}\mathrm{div}(\mathfrak{g}^{-1})(\mathbf{X}_{t})\mathrm{d}t+\mathfrak{g}(\mathbf{X}_{t})^{-\frac{1}{2}}\mathrm{d}\mathbf{B}_{t},}$
(35)

the following ODE has the same marginal density

$\textstyle{\mathrm{d}{\mathbf{Y}}_{t}=\left[\tfrac{1}{2}\nabla\cdot\mathfrak{g}^{-1}(\mathbf{Y}_{t})-\tfrac{1}{2}\mathfrak{g}^{-1}(\mathbf{Y}_{t})\nabla\log
p_{t}(\mathbf{Y}_{t})\right]\mathrm{d}t.}$ (36)

##  Appendix F Reflected Brownian Motion and Skorokhod problems

In this section, we provide the basic definitions and results to derive the
time-reversal of the reflected Brownian motion in Appendix G. We follow
closely the presentation of (Lions & Sznitman, 1984) and (Burdzy et al.,
2004). We first define the _Skorokhod problem_ for deterministic problems. We
consider $\mathcal{M}$ to be a smooth open bounded domain. We recall that the
normal vector $n$ is defined on $\partial\mathcal{M}$ and we set $n(x)=0$ for
any $x\notin\partial\mathcal{M}$.

Before giving the definition of the _Skorokhod problem_ , we recall what is
the space of functions of _bounded variations_.

######  Definition F.1.

Let $a,b\in\left(-\infty,+\infty\right)$ and
$f:\in\mathrm{C}(\left[a,b\right],\mathbb{R})$. We define the _total
variation_ of $f$ as

$\textstyle\mathrm{V}_{a,b}(f)=\sup\\{\sum_{i=0}^{n-1}\|f(x_{i+1})-f(x_{i})\|\,:\;(x_{i})_{i=0}^{n-1},\
a=x_{0}\leq x_{1}\leq\dots\leq x_{n-1}\leq x_{n}=b,\ n\in\mathbb{N}\\}.$ (37)

$f$ has bounded variations over $\left[a,b\right]$ if
$\mathrm{V}_{a,b}(f)<+\infty$. Let
$f\in\mathrm{C}(\left[0,+\infty\right),\mathbb{R})$. $f$ has bounded
variations over $\left[0,+\infty\right)$ if for any $b>0$, $f$ has bounded
variations over $\left[0,b\right]$.

The notion of bounded variation is a relaxation of the differentiability
requirement. In particular, if
$f\in\mathrm{C}^{1}(\left[a,b\right],\mathbb{R})$, we have
$\mathrm{V}_{a,b}(f)=\int_{a}^{b}\|f^{\prime}(t)\|\mathrm{d}t$. In the
definition of the _Skorokhod problem_ , we will see that this relaxation is
necessary, even in the deterministic setting.

For any function of bounded variation
$f\in\mathrm{C}(\left[a,b\right],\mathbb{R})$ on $\left[a,b\right]$, we define
$\left|f\right|:\ \left[a,b\right]\to\left[0,+\infty\right)$ given for any
$t\in\left[a,b\right]$ by $\left|f\right|_{t}=\mathrm{V}_{a,t}(f)$. Note that
$\left|f\right|$ is non-decreasing and right-continuous. Therefore, we can
define the measure $\mu_{\left|f\right|}$ on $\left[a,b\right]$, given for any
$s,t\in\left[a,b\right]$ with $t\geq s$ by
$\mu_{\left|f\right|}([s,t])=\left|f\right|(t)-\left|f\right|(s)$. In
particular, for any $\varphi:\left[a,b\right]\to\mathbb{R}_{+}$, we define

$\textstyle\int_{a}^{b}\varphi(t)\mathrm{d}\left|f\right|_{t}=\int_{a}^{b}\varphi(t)\mathrm{d}\mu_{\left|f\right|}(t).$
(38)

In addition, $f$ can be decomposed in a difference of two non-decreasing
processes right continuous processes $g_{1}$, $g_{2}$, where for any
$t\in\left[a,b\right]$, $f(t)=g_{1}(t)-g_{2}(t)$,
$g_{1}(t)=\left|f\right|_{t}$ and $g_{2}(t)=\left|f\right|_{t}-f(t)$. Hence,
for every $\varphi$ bounded on $\left[a,b\right]$, we can define

$\textstyle\int_{a}^{b}\varphi(t)\mathrm{d}f(t)=\int_{a}^{b}\varphi(t)\mathrm{d}g_{1}(t)-\int_{a}^{b}\varphi(t)\mathrm{d}g_{2}(t).$
(39)

Note that these definitions can be extended to the setting where $f:\
\left[a,b\right]\to\mathbb{R}^{d}$.

We begin with the following result, see (Lions & Sznitman, 1984).

######  Theorem F.2.

Let $(x_{t})_{t\geq 0}\in\mathrm{C}(\left[0,+\infty\right),\mathbb{R})$. Then,
there exists a unique couple $(\bar{x}_{t},k_{t})_{t\geq 0}$ such that

  1. (a)

$(k_{t})_{t\geq 0}$ has bounded variation over $\left[0,+\infty\right)$.

  2. (b)

$(\bar{x}_{t})_{t\geq
0}\in\mathrm{C}(\left[0,+\infty\right),\overline{\mathcal{M}})$.

  3. (c)

For any $t\geq 0$, $x_{t}+k_{t}=\bar{x}_{t}$.

  4. (d)

For any $t\geq 0$,
$\left|k\right|_{t}=\int_{0}^{t}\mathbf{1}_{\bar{x}_{s}\in\partial\mathcal{M}}(\bar{x}_{s})\mathrm{d}\left|k\right|_{s}$
and $k_{t}=\int_{0}^{t}n(\bar{x}_{s})\mathrm{d}\left|k\right|_{s}$.

Let us discuss Theorem F.2. First, Theorem F.2-(c) states the original
(unconstrained) process $(x_{t})_{t\geq 0}$ can be decomposed into a
constrained version $(\bar{x}_{t})_{t\geq 0}$ and a bounded variation process
$(k_{t})_{t\geq 0}$. The process $(\left|k\right|_{t})_{t\geq 0}$ counts the
number of times the constrained process $(\bar{x}_{t})_{t\geq 0}$ hits the
boundary. More formally, we have
$\left|k\right|_{t}=\int_{0}^{t}\mathbf{1}_{x\in\partial\mathcal{M}}(\bar{x}_{s})\mathrm{d}\left|k\right|_{s}$.
When, we hit the boundary, we reflect the process. This condition is expressed
in $k_{t}=\int_{0}^{t}n(\bar{x}_{s})\mathrm{d}\left|k\right|_{s}$.

We now consider the extension to stochastic processes. We are given
$(\mathbf{X}_{t})_{t\geq 0}$ such that

$\mathrm{d}\mathbf{X}_{t}=b(\mathbf{X}_{t})\mathrm{d}t+\sigma(t)\mathrm{d}\mathbf{B}_{t},$
(40)

where $(\mathbf{B}_{t})_{t\geq 0}$ is a $d$-dimensional Brownian motion. We
also assume that $b$ and $\sigma$ are Lipschitz which implies the existence
and strong uniqueness of $(\mathbf{X}_{t})_{t\geq 0}$. We have the following
result (Lions & Sznitman, 1984).

######  Theorem F.3.

There exists a unique process $(\bar{\mathbf{X}}_{t},\mathbf{k}_{t})_{t\geq
0}$ such that

  1. (a)

$(\mathbf{k}_{t})_{t\geq 0}$ has bounded variation over
$\left[0,+\infty\right)$ almost surely.

  2. (b)

$(\bar{\mathbf{X}}_{t})_{t\geq
0}\in\mathrm{C}(\left[0,+\infty\right),\overline{\mathcal{M}})$.

  3. (c)

For any $t\geq 0$,
$\bar{\mathbf{X}}_{t}=\bar{\mathbf{X}}_{0}+\int_{0}^{t}b(\bar{\mathbf{X}}_{s})\mathrm{d}s+\int_{0}^{t}\sigma(\bar{\mathbf{X}}_{s})\mathrm{d}\mathbf{B}_{s}-\mathbf{k}_{t}$.

  4. (d)

For any $t\geq 0$,
$\left|\mathbf{k}\right|_{t}=\int_{0}^{t}\mathbf{1}_{\bar{x}_{s}\in\partial\mathcal{M}}(\bar{x}_{s})\mathrm{d}\left|\mathbf{k}\right|_{s}$
and
$\mathbf{k}_{t}=\int_{0}^{t}n(\bar{x}_{s})\mathrm{d}\left|\mathbf{k}\right|_{s}$.

The The process $(\mathbf{X}_{t})_{t\geq 0}$ is almost surely continuous, so
we could apply the previous theorem almost surely for all the realizations of
the process,. However, this does not tell us if the obtained solutions
$(\bar{\mathbf{X}}_{t},\mathbf{k}_{t})_{t\geq 0}$ form themselves a process.
The main difference with Theorem F.2 is in Theorem F.3-(c) which differs from
Theorem F.3-(c). Note that in the case where $b=0$ and
$\sigma=\operatorname{Id}$ we recover Theorem F.3-(c). This is not true in the
general case. However, it can be seen that for any realization of the process
$(\bar{\mathbf{X}}_{t})_{t\geq 0}$, we have that
$(\bar{\mathbf{X}}_{t},\mathbf{k}_{t})_{t\geq 0}$ is solution of the
_deterministic_ Skorokhod problem by letting
$x_{t}=\bar{\mathbf{X}}_{0}+\int_{0}^{t}b(\bar{\mathbf{X}}_{s})\mathrm{d}s+\int_{0}^{t}\sigma(\bar{\mathbf{X}}_{s})\mathrm{d}\mathbf{B}_{s}$.
The backward and forward Kolmogorov equations can be found in (Burdzy et al.,
2004). Note that the presence of the process $(\mathbf{k}_{t})_{t\geq 0}$
incurs notable complications compared to unconstrained processes. In
particular, there is no martingale problem associated with weak solutions of
reflected SDEs but only sub-martingale problems, see (Kang & Ramanan, 2017)
for instance.

##  Appendix G Time-reversal for reflected Brownian motion

We start with the following definition.

######  Definition G.1.

Let $\mathcal{M}\subset\mathbb{R}^{d}$ be an open set. $\mathcal{M}$ has a
smooth boundary if for any $x\in\partial\mathcal{M}$, there exists
$\mathsf{U}\subset\mathbb{R}^{d}$ open and
$f\in\mathrm{C}^{\infty}(\mathsf{U},\mathbb{R})$ such that $x\in\mathsf{U}$
and  (a)
$\mathrm{cl}(\mathcal{M})\cap\mathsf{U}=\\{x\in\mathsf{U}\,:\;f(x)\leq 0\\}$,
(b) $\nabla f(x)\neq 0$for any $x\in\mathsf{U}$ where
$\mathrm{cl}(\mathcal{M})$ is the closure of $\mathcal{M}$.

We will make the use of the following lemma which is a straightforward
extension of Burdzy et al. (2004, Theorem 2.6).

######  Lemma G.2.

Let $u$ such that $s\mapsto
u(s,x)\in\mathrm{C}^{1}(\left(0,T\right),\mathbb{R})$, for any
$s\in\left(0,T\right)$, $x\mapsto
u(s,x)\in\mathrm{C}^{2}(\mathrm{\mathcal{M}},\mathbb{R})$ and
$u\in\mathrm{C}^{1}(\mathrm{cl}(\mathrm{\mathcal{M}}),\mathbb{R})$. Then for
any $T\geq 0$, $s,t\in\left[0,T\right]$, we have

$\textstyle{\mathbb{E}[\int_{s}^{t}u(w,\bar{\mathbf{B}}_{w})\mathrm{d}\mathbf{k}_{w}]=\tfrac{1}{2}\int_{s}^{t}\int_{\partial\mathcal{M}}u(w,x)p_{w}(x)\mathrm{d}\sigma(x)\mathrm{d}w.}$
(41)

Note that we recover Burdzy et al. (2004, Theorem 2.6) if we set $u=1$. We
also emphasize that the result of Burdzy et al. (2004, Theorem 2.6) is
stronger than Lemma G.2 as it holds not only in expectation but also in
$\mathrm{L}^{2}$ and almost surely.

We are now ready to prove Theorem 3.2. We follow the approach of (Petit, 1997)
which itself is based on an extension of (Haussmann & Pardoux, 1986). We refer
to Cattiaux et al. (2021) for recent entropic approaches of time-reversal.
Recall that $(\bar{\mathbf{B}}_{t},\mathbf{k}_{t})_{t\geq 0}$ is a solution to
the _Skorokhod problem_ (Skorokhod, 1961) if $(\mathbf{k}_{t})_{t\geq 0}$ is a
bounded variation process and $(\bar{\mathbf{B}}_{t})_{t\geq 0}$ a continuous
adapted process such that for any $t\geq 0$,
$\mathbf{B}_{t}=\bar{\mathbf{B}}_{t}+\mathbf{k}_{t}\in\mathcal{M}$,
$(\bar{\mathbf{B}}_{t})_{t\geq 0}$ and

$\textstyle{\left|\mathbf{k}\right|_{t}=\int_{0}^{t}\mathbf{1}_{\bar{\mathbf{B}}_{s}\in\partial\mathcal{M}}\mathrm{d}\left|\mathbf{k}\right|_{s},\quad\mathbf{k}_{t}=\int_{0}^{t}\mathbf{n}(\bar{\mathbf{B}}_{s})\mathrm{d}\left|\mathbf{k}\right|_{s},}$
(42)

In what follows, we define $(\mathbf{Y}_{t})_{t\in\left[0,T\right]}$ such that
for any $t\in\left[0,T\right]$, $\mathbf{Y}_{t}=\bar{\mathbf{B}}_{T-t}$. Let
us consider the process $(\tilde{\mathbf{B}}_{t})_{t\in\left[0,T\right]}$
defined for any $t\in\left[0,T\right]$ by

$\textstyle{\tilde{\mathbf{B}}_{t}=-\bar{\mathbf{B}}_{T}+\bar{\mathbf{B}}_{T-t}+\mathbf{k}_{T}-\mathbf{k}_{T-t}-\int_{T-t}^{T}\nabla\log
p_{s}(\bar{\mathbf{B}}_{s})\mathrm{d}s.}$ (43)

First, note that $t\mapsto\tilde{\mathbf{B}}_{t}$ is continuous. Denote by
$\mathcal{F}$, the filtration associated with
$(\bar{\mathbf{B}}_{T-t})_{t\in\left[0,T\right]}$. We have that
$(\tilde{\mathbf{B}}_{t})_{t\in\left[0,T\right]}$ is adapted to
$(\bar{\mathbf{B}}_{T-t})_{t\in\left[0,T\right]}$. Even more so, we have that
$(\tilde{\mathbf{B}}_{t})_{t\in\left[0,T\right]}$ satisfies the strong Markov
property since $(\bar{\mathbf{B}}_{t})_{t\in\left[0,T\right]}$ also satisfies
the strong Markov property. Let
$g\in\mathrm{C}_{c}^{\infty}(\mathrm{cl}(\mathcal{M}))$ and consider for any
$0\leq s\leq t\leq T$,
$\mathbb{E}[(\tilde{\mathbf{B}}_{t}-\tilde{\mathbf{B}}_{s})g(\bar{\mathbf{B}}_{T-t})]$.
For any $0\leq s\leq t\leq T$ we have

$\mathbb{E}[(\tilde{\mathbf{B}}_{t}-\tilde{\mathbf{B}}_{s})g(\bar{\mathbf{B}}_{T-t})]=\textstyle{\mathbb{E}[(-\bar{\mathbf{B}}_{T-s}+\bar{\mathbf{B}}_{T-t}+\mathbf{k}_{T-s}-\mathbf{k}_{T-t}-\int_{T-t}^{T-s}\nabla\log
p_{u}(\bar{\mathbf{B}}_{u})\mathrm{d}u)g(\bar{\mathbf{B}}_{T-t})].}$ (44)

In what follows, we prove that for any $0\leq s\leq t\leq T$ we have
$\mathbb{E}[(\tilde{\mathbf{B}}_{t}-\tilde{\mathbf{B}}_{s})g(\bar{\mathbf{B}}_{T-t})]=0$.
Therefore, we only need to prove that for any $0\leq s\leq t\leq T$ we have

$\textstyle{\mathbb{E}[(-\bar{\mathbf{B}}_{t}+\bar{\mathbf{B}}_{s}+\mathbf{k}_{t}-\mathbf{k}_{s}-\int_{s}^{t}\nabla\log
p_{u}(\bar{\mathbf{B}}_{u})\mathrm{d}u)g(\bar{\mathbf{B}}_{t})]=0.}$ (45)

Let $t\in\left(0,T\right]$. We introduce $u:\
\left[0,t\right]\times\mathcal{M}$ such that for any $s\in\left[0,t\right]$
and $x\in\mathcal{M}$,
$u(s,x)={\mathbb{E}}[g(\bar{\mathbf{B}}_{t})|\bar{\mathbf{B}}_{s}=x]$. Using
Burdzy et al. (2004, Theorem 2.8) we get that for any $x\in\mathcal{M}$,
$s\mapsto u(s,x)\in\mathrm{C}^{1}(\left(0,t\right),\mathbb{R})$ and for any
$s\in\left(0,t\right)$, $x\mapsto
u(s,x)\in\mathrm{C}^{2}(\mathrm{\mathcal{M}},\mathbb{R})$ and $x\mapsto
u(s,x)\in\mathrm{C}^{1}(\mathrm{cl}(\mathcal{M}),\mathbb{R})$. In addition, we
have that for any $s\in\left(0,t\right)$ and for any $x\in\mathcal{M}$ and
$x_{0}\in\partial\mathcal{M}$

$\textstyle{\partial_{s}u(s,x)+\tfrac{1}{2}\Delta u(s,x)=0,\qquad\langle\nabla
u(s,x_{0}),\mathbf{n}(x_{0})\rangle=0.}$ (46)

This equation is called the backward Kolmogorov equation (note that in the
case of a bounded domain, the solutions satisfy the Neumann boundary
conditions). Using (46), $\bar{\mathbf{B}}_{t}=\mathbf{B}_{t}-\mathbf{k}_{t}$
for any $t\geq 0$ and the Itô formula for semimartingale (Revuz & Yor, 2013,
Chapter IV, Theorem 3.3) we have that for any $s\in\left(0,t\right)$

$\displaystyle\mathbb{E}[u(t,\bar{\mathbf{B}}_{t})\bar{\mathbf{B}}_{t}]$
$\displaystyle\textstyle{=\mathbb{E}[u(s,\bar{\mathbf{B}}_{s})\bar{\mathbf{B}}_{s}]+\mathbb{E}[\tfrac{1}{2}\int_{s}^{t}\bar{\mathbf{B}}_{w}\Delta
u(w,\bar{\mathbf{B}}_{w})\mathrm{d}w]}\textstyle{+\mathbb{E}[\int_{s}^{t}\nabla
u(w,\bar{\mathbf{B}}_{w})\mathrm{d}w]}$ (47)
$\displaystyle\qquad\textstyle{-\mathbb{E}[\int_{s}^{t}\bar{\mathbf{B}}_{w}\langle\nabla
u(w,\bar{\mathbf{B}}_{w}),\mathbf{n}(\bar{\mathbf{B}}_{w})\rangle\mathrm{d}\left|\mathbf{k}\right|_{w}]}$
(48)
$\displaystyle\qquad\textstyle{-\mathbb{E}[\int_{s}^{t}u(w,\bar{\mathbf{B}}_{w})\mathbf{n}(\bar{\mathbf{B}}_{w})\mathrm{d}\left|\mathbf{k}\right|_{w}]}$
(49)
$\displaystyle\qquad\textstyle{+\mathbb{E}[\int_{s}^{t}\bar{\mathbf{B}}_{w}\partial_{w}u(w,\bar{\mathbf{B}}_{w})\mathrm{d}w]}$
(50)
$\displaystyle\textstyle{=\mathbb{E}[u(s,\bar{\mathbf{B}}_{s})\bar{\mathbf{B}}_{s}]\textstyle{+\mathbb{E}[\int_{s}^{t}\nabla
u(w,\bar{\mathbf{B}}_{w})\mathrm{d}w]}}\textstyle{-\mathbb{E}[\int_{s}^{t}u(w,\bar{\mathbf{B}}_{w})\mathbf{n}(\bar{\mathbf{B}}_{w})\mathrm{d}\left|\mathbf{k}\right|_{w}]}.$
(51)

In addition, using the Fubini theorem and the definition of $\mathbf{k}_{t}$
we have that for any $s\in\left(0,t\right)$

$\textstyle{\mathbb{E}[\int_{s}^{t}u(w,\bar{\mathbf{B}}_{w})\mathbf{n}(\bar{\mathbf{B}}_{w})\mathrm{d}\left|\mathbf{k}\right|_{w}]=\mathbb{E}[\int_{s}^{t}{\mathbb{E}}[g(\bar{\mathbf{B}}_{t})|\bar{\mathbf{B}}_{w}]\mathbf{n}(\bar{\mathbf{B}}_{w})\mathrm{d}\left|\mathbf{k}\right|_{w}]=\mathbb{E}[g(\bar{\mathbf{B}}_{t})(\mathbf{k}_{t}-\mathbf{k}_{s})].}$
(52)

Finally, using the divergence theorem and Burdzy et al. (2004, Theorem 2.6) we
have that for any $s\in\left(0,t\right)$

$\displaystyle\textstyle{\mathbb{E}[\int_{s}^{t}\nabla
u(w,\bar{\mathbf{B}}_{w})\mathrm{d}w]}$
$\displaystyle=\textstyle{\int_{s}^{t}\int_{\mathcal{M}}\nabla
u(w,x)p_{w}(x)\mathrm{d}x\mathrm{d}w}$ (53)
$\displaystyle=-\textstyle{\int_{s}^{t}\int_{\mathcal{M}}u(w,x)\nabla\log(p_{w}(x))p_{w}(x)\mathrm{d}x\mathrm{d}w+\int_{s}^{t}\int_{\partial\mathcal{M}}u(w,x)p_{w}(x)\mathrm{d}x\mathrm{d}\sigma(w)},$
(54)

where $\sigma$ is the surface area measure on $\partial\mathcal{M}$, see
Burdzy et al. (2004). Using Lemma G.2 and the Fubini theorem we get that

$\displaystyle\textstyle{\mathbb{E}[\int_{s}^{t}\nabla
u(w,\bar{\mathbf{B}}_{w})\mathrm{d}w]}$
$\displaystyle=-\textstyle{\int_{s}^{t}\int_{\mathcal{M}}u(w,x)\nabla\log(p_{w}(x))p_{w}(x)\mathrm{d}x\mathrm{d}w+\mathbb{E}[\int_{s}^{t}u(w,\bar{\mathbf{B}}_{w})\mathrm{d}\mathbf{k}_{w}]}$
(55)
$\displaystyle=-\textstyle{\mathbb{E}[\int_{s}^{t}g(\bar{\mathbf{B}}_{t})\nabla\log(p_{w}(\bar{\mathbf{B}}_{w}))\mathrm{d}w]+2\mathbb{E}[g(\bar{\mathbf{B}}_{t})(\mathbf{k}_{t}-\mathbf{k}_{s})]}$
(56)

Combining (51), (52) and (56) we get that

$\textstyle{\mathbb{E}[u(t,\bar{\mathbf{B}}_{t})]=\mathbb{E}[u(s,\bar{\mathbf{B}}_{s})]-\mathbb{E}[g(\bar{\mathbf{B}}_{t})\int_{s}^{t}\nabla\log
p_{w}(\bar{\mathbf{B}}_{w})\mathrm{d}w]}+\textstyle{\mathbb{E}[g(\bar{\mathbf{B}}_{t})(\mathbf{k}_{t}-\mathbf{k}_{s})]}.$
(57)

Therefore, we get (45) and (44). Hence,
$(\tilde{\mathbf{B}}_{t})_{t\in\left[0,T\right]}$ is a continuous martingale.
In addition, we have that for any $t\in\left[0,T\right]$,
$\mathbb{E}[\tilde{\mathbf{B}}_{t}\tilde{\mathbf{B}}_{t}^{\top}]=t\operatorname{Id}$
and therefore, $(\tilde{\mathbf{B}}_{t})_{t\in\left[0,T\right]}$ is a Brownian
motion using the Lévy characterisation of Brownian motion (Revuz & Yor, 2013,
Chapter IV, Theorem 3.6). Denote
$(\mathbf{j}_{t})_{t\in\left[0,T\right]}=(\mathbf{k}_{T}-\mathbf{k}_{T-t})_{t\in\left[0,T\right]}$.
Using (44), we have that for any $t\in\left[0,T\right]$

$\textstyle{\bar{\mathbf{Y}}_{t}=\bar{\mathbf{Y}}_{0}+\tilde{\mathbf{B}}_{t}+\int_{0}^{t}\nabla\log
p_{T-s}(\mathbf{Y}_{s})\mathrm{d}s-\mathbf{j}_{t}.}$ (58)

Using (42), we have for any $t\in\left[0,T\right]$

$\textstyle{\left|\mathbf{j}\right|_{t}=\int_{0}^{t}\mathbf{1}_{\mathbf{Y}_{s}\in\partial\mathcal{M}}\mathrm{d}\left|\mathbf{j}\right|_{s},\quad\mathbf{j}_{t}=\int_{0}^{t}\mathbf{n}(\bar{\mathbf{Y}}_{s})\mathrm{d}\left|\mathbf{j}\right|_{s},}$
(59)

which concludes the proof.

##  Appendix H Conformational modelling of polypeptide backbones under end
point constraints

Polypeptides and proteins constitute an important class of biogenic
macromolecules that underpin most aspects of organic life. Accurately
modelling their conformational ensembles, i.e. the set of three-dimensional
structures they assume under physiological conditions, is essential to both
understanding the biological function of existing and designing the enzymatic
or therapeutic properties of novel proteins (Lane, 2023). Motivated by the
success of diffusion models in computer vision and natural language
processing, there has been considerable interest in applying them to learn and
sample from distributions over the conformational space of protein structures
(Watson et al., 2022; Trippe et al., 2022; Wu et al., 2022).

###  H.1 Problem parameterisation

Proteins are biopolymers in which a sequence of $N$ amino acids is joined
together through $N-1$ peptide bonds, resulting in a so-called polypeptide
backbone with protruding amino acid residues. As the deviation of chemical
bond lengths and angles from their theoretical optimum is generally
negligible, the problem of modelling the three-dimensional structure of this
polypeptide chain is often reframed in the space of the internal torsion
angles $\Phi$ and $\Psi$ (see Figure 12(a) for an illustration), which can be
modelled on a $(2N-2)$-dimensional torus $\mathbb{T}^{2N-2}$.

(a) A commonly-used approximate parameterisation of backbone geometry only
considers the $C_{\alpha}$ torsion angles $\Phi$ and $\Psi$.

(b) As peptide bond orientations can be inferred relatively reliably,
researchers often only model the $C_{\alpha}$ traces.

Figure 12: Widely-used approaches to modelling the conformations of
polypeptide backbones.

In many data-scarce practical settings such as antibody or enzyme design, it
is often unnecessary or even undesirable to model the structure of an entire
protein, as researchers are primarily interested in specific functional sites
with distinct biochemical properties. However, generating conformational
ensembles for such substructural elements necessitates positional constraints
on their endpoints to ensure that they can be accommodated by the remaining
scaffold. While it is conceivable that a diffusion model could derive such
constraints from limited experimental data, we argue that it is much more
efficient and precise to encode them explicitly.

For this purpose, we adopt the distance constraint formulation from Han &
Rudolph (2006) and interpret the backbone as a spatial chain with $N$
spherical joints and fixed-length links (see Figure 13(a) for an
illustration). After selecting an arbitrary anchor point, the geometry of the
polypeptide chain is fully specified by  (a) the set of link lengths
$\mathbf{\ell}=\\{\ell_{j}\\}_{j=1}^{N}$,  (b) the set of vectors
$\mathbf{r}=\\{r(1,j)\\}_{j=2}^{N}$ between the anchor point and each atom in
the chain,  (c) the set of dihedral angles

$\mathbf{T}=\left\\{\arccos\left(\frac{|\langle r(1,j)\times
r(1,j+1),r(1,j+1)\times r(1,j+2)\rangle|}{|r(1,j)\times
r(1,j+1)||r(1,j+1)r(1,j+2)|}\right)\right\\}_{j=2}^{N-2}\in\mathbb{T}^{N-3}$

defined by any three consecutive vectors. Given an arbitrary anchor point
distance $d_{\text{anchor}}=\ell_{N}=r(1,N)$, the set of valid vectors
$\mathbf{r}$ is given by the convex polytope
$\mathbb{P}\subseteq\mathbb{R}^{3}$ defined by the following linear
constraints:

$\begin{array}[]{rcl}r(1,3)&\leq&\ell_{1}+\ell_{2},\\\
-r(1,3)&\leq&-\left|\ell_{1}-\ell_{2}\right|,\\\ \end{array}$
$\left.\begin{array}[]{rcl}r(1,j)-r(1,j+1)&\leq&\ell_{j},\\\
-r(1,j)+r(1,j+1)&\leq&\ell_{j},\\\
-r(j)-r(j+1)&\leq&-\ell_{j},\end{array}\right\\}3\leq j\leq N-2,\\\ $
$\begin{array}[]{rcl}r(1,N-1)&\leq&\ell_{N-1}+d_{\text{anchor}},\\\
-r(1,N-1)&\leq&-\left|\ell_{N-1}-d_{\text{anchor}}\right|.\end{array}$

(a) An illustrative diagram of the proposed parameterisation for modelling the
$C_{\alpha}$ trace geometry of the cyclic peptide c-AAGAGG.

(b) The convex polytope constraining the diagonals of the triangles for the
given bond lengths in the illustrated molecule. The total design space is the
product of this polytope with the 4D flat torus.

Figure 13: Parameterising the conformational space of polypeptide backbones
under anchor point distance constraints.

This means that the set of all valid polypeptide backbone conformations is
defined by the product manifold $\mathbb{P}\times\mathbb{T}^{N-3}$, enabling
us to train diffusion models that exclusively generate conformations with a
fixed anchor point distance $d_{\text{anchor}}$.

###  H.2 Data generation and model training

As a proof-of-concept for the practicality of our methods, we chose to model
the conformational distribution of the cyclic peptide c-AAGAGG (see Figure
13). Cyclic peptides are an increasingly important drug modality with
therapeutic uses ranging from antimicrobials to oncology, exhibiting circular
polypeptide backbones that confer a range of desirable pharmacodynamic and
pharmacokinetic properties (Dougherty et al., 2019). To reduce the
dimensionality of the problem, we only consider the $C_{\alpha}$ traces (with
fixed $C_{\alpha}$-$C_{\alpha}$ link distances of 3.6 Å) instead of the full
polypeptide backbone (see Figure 12(b)), although we note that our framework
is fully applicable to both settings.

To derive a suitable dataset, the product manifold describing the
conformations of cyclic $C_{\alpha}$ traces of length $N=6$, i.e.
$\mathbb{P}\times\mathbb{T}^{3}$, was constructed and used to generate
$10^{7}$ uniform samples satisfying the anchor point distance constraint
$d_{\text{anchor}}=0$. Subsequently, an estimate of the free energy $E_{i}$ of
each sample $i$ was obtained by (1) reconstructing the full-atom peptide from
each $C_{\alpha}$ trace using the PULCHRA algorithm (Rotkiewicz & Skolnick,
2008), (2) relaxing all non-$C_{\alpha}$ backbone and side-chain atoms
(keeping the $C_{\alpha}$ positions fixed), and (3) modelling the potential
energy of each of the resulting conformations using the OpenMM suite of
molecular dynamics tools (Eastman et al., 2017), and the AMBER force field
(Hornak et al., 2006). These free energy estimates were then used to
approximate the Boltzmann distribution over conformational states

$\textstyle{p_{B}(i)\propto\exp\left(-\frac{E_{i}}{k_{B}T}\right),}$

where temperature was set to $T=273.15K$ and $k_{B}=1.380649\cdot
10^{-23}J\cdot K^{-1}$ is the Boltzmann constant. As the potential energy
surfaces of polypeptides are known to be relatively rough (Krivov & Karplus,
2004), the resulting distribution was moderately smoothened by running forward
Brownian motion on both the polytope and the torus for 100 steps, using a
small step size of $10^{-2}$ and the respective metrics. Finally, a sample of
$10^{6}$ $C_{\alpha}$ traces was drawn from this distribution and used for
training and evaluating our models.

##  Appendix I Configurational modelling of robotic arms under manipulability
constraints

Accurately determining and specifying the movement of a robotic arm and the
forces it exerts is a fundamental problem in many real-world robotics
applications. A widely-used set of descriptors for modelling the flexibility
of a given joint configuration are so-called manipulability ellipsoids
(Yoshikawa, 1985), which are kinetostatic performance measures that quantify
the ability to move or exert forces along different directions. Jaquier et al.
(2021) present a geometric framework to learn trajectories of manipulability
ellipsoids by making use of the fact any ellipsoid $M\in\mathbb{R}^{N}$ is
defined by the set of points
$\\{\mathbf{x}|\mathbf{x}^{T}\mathbf{A}\mathbf{x}=1\\}$ where $\mathbf{A}$
lies on the manifold of $N\times N$ symmetric positive definite matrices
$S^{N}_{++}$.

In many practical settings, it is desirable to constrain the minimal or
maximal volume of a manipulability ellipsoid to retain motional flexibility or
limit the magnitude of the exerted force. This necessitates lower or upper
limits on the determinant of A, translating into constraints on $S^{N}_{++}$.
To model this, we make use of one of the datasets introduced by Jaquier et al.
(2021), containing demonstrations of a robotic arm drawing different letters
in the plane, providing the respective positional trajectories
($\mathbb{R}^{2}$) and velocity manipulability ellipsoids ($S^{2}_{++}$).

We use the processing routines provided by Jaquier et al. (2021) to
interpolate the trajectories into $10^{4}$ distinct points, for each of which
we derive the position $\mathbf{x}\in\mathbb{R}^{2}$ and the PSD matrix

$\mathbf{A}=\begin{pmatrix}a&b\\\ b&c\end{pmatrix}\in S^{2}_{++}.$

parametrizing the velocity manipulability ellipsoid $M\in\mathbb{R}^{2}$. The
resulting data is split into training, validation, and test sets by trajectory
and visualized in Figures 14 and 15.

(a)

(b)

Figure 14: Positional trajectories $\mathbf{x}\in\mathbb{R}^{2}$ (a) and the
parameters $a,b,c$ of the the SPD matrix $\mathbf{A}\in S^{2}_{++}$ (b) for
the letter C.

(a)

(b)

Figure 15: Positional trajectories $\mathbf{x}\in\mathbb{R}^{2}$ (a) and the
parameters $a,b,c$ of the the SPD matrix $\mathbf{A}\in S^{2}_{++}$ (b) for
the letter L.

##  Appendix J Experimental details

In what follows we describe the experimental settings used to generate results
introduced in Section 5. The models and experiments have been implemented in
Jax (Bradbury et al., 2018), using a modified version of the Riemannian
geometry library Geomstats (Miolane et al., 2020).

#### Architecture.

The architecture of the score network $\bm{s}_{\theta}$ is given by a
multilayer perceptron with $6$ hidden layers with $512$ units each. We use
sinusoidal activation functions.

#### Training.

All models are trained by the stochastic optimizer Adam (Kingma & Ba, 2014)
with parameters $\beta_{1}=0.9$, $\beta_{2}=0.999$, batch-size of $256$ data-
points. The learning rate is annealed with a linear ramp from $0$ to $1000$
steps, reaching the maximum value of $2e-4$, and from then with a cosine
schedule down to $0$ after $100k$ iterations in total.

#### Diffusion.

Following Song et al. (2021), the diffusion models diffusion coefficient is
parametrized as $g(t)=\sqrt{\beta(t)}$ with $\beta:\
t\mapsto\beta_{\min}+(\beta_{\max}-\beta_{\min})\cdot t$, where we found
$\beta_{\min}=0.001$ and $\beta_{\max}=6$ to work best.

#### Metrics.

We measure the performance of trained models via the Maximum Mean Discrepancy
(MMD) (Gretton et al., 2012), which is a kernel based metric between two
distributions $P$ and $Q$. The MMD can be empirically approximated with the
following U-statistics
$\mathrm{MMD}^{2}(P,Q)=\tfrac{1}{m(m-1)}\sum_{i}\sum_{j\neq
i}k(x_{i},x_{j})+\tfrac{1}{m(m-1)}\sum_{i}\sum_{j\neq
i}k(y_{i},y_{j})-2\tfrac{1}{m^{2}}\sum_{i}\sum_{j}k(x_{i},y_{j})$ with
$x_{i}\sim P$ and $y_{i}\sim Q$, where $k$ is a kernel.

#### Low temperature sampling.

As in Ingraham et al. (2022), we use low temperature sampling in order to
trade-off fidelity for diversity. It combines to ingredients:  1. an
approximation of the corrector sampling Song et al. (2021),  2. a low
temperature approximation of the target $p$. Namely, instead of targeting
$p_{0}$ we target $p_{0}^{\lambda_{0}}$ (unnormalised density).  We refer to
Ingraham et al. (2022, Appendix B) for details about the derivation. Recall
that we consider the following SDE in the backward process

$\mathrm{d}\overleftarrow{\mathbf{X}}_{t}=\beta_{T-t}\\{\tfrac{1}{2}\overleftarrow{\mathbf{X}}_{t}+s_{\theta}(T-t,\overleftarrow{\mathbf{X}}_{t})\\}\mathrm{d}t+\beta_{t}^{1/2}\mathrm{d}\mathbf{B}_{t}.$
(60)

The updated scheme is given by

$\mathrm{d}\overleftarrow{\mathbf{X}}_{t}=\beta_{T-t}\\{\tfrac{1}{2}\overleftarrow{\mathbf{X}}_{t}+(\lambda_{t}+\tfrac{\lambda_{0}\Psi}{2})s_{\theta}(T-t,\overleftarrow{\mathbf{X}}_{t})\\}\mathrm{d}t+(1+\Psi)^{1/2}\beta_{t}^{1/2}\mathrm{d}\mathbf{B}_{t},$
(61)

where $\lambda_{t}=\lambda_{0}/(\alpha_{t}+(1-\alpha_{t})\lambda_{0})$ and
$\alpha_{t}=\exp[-2\int_{0}^{t}\beta_{s}\mathrm{d}s]$. In practice
$\lambda_{0}$ controls the concentration of the measure while $\Psi$ controls
the strength of the corrector. We use $\Psi=0$ as we do not consider a
corrector term and perform a sweep over the parameters
$\lambda_{0}\in\\{1,2,3,5,10\\}$.

###  J.1 Synthetic data on polytopes

#### Hypercube $[-1,1]^{n}$.

The hypercube is a specific instance of a convex polytope where the affine
constraints are given by the following coefficients:

$A=\begin{pmatrix}1&\dots&0\\\ -1&\dots&0\\\ \vdots&\ddots&\vdots\\\
0&\dots&1\\\ 0&\dots&-1\\\ \end{pmatrix},\;\;b=\begin{pmatrix}1\\\ 1\\\
\vdots\\\ 1\\\ 1\end{pmatrix}.$ (62)

Where $A$ is a $2n\times n$ matrix and $b$ an $n$ dimensional vector.

We construct the training and test datasets by sampling for both $100000$
points from a mixture of ‘wrapped normal’ distributions illustrated in Figure
16(a) and which density is given by

$p_{0}(x)=0.7~{}\text{ReflectedStep}[(0.5,0.5),\cdot,\\{f_{i}\\}_{i\in\mathcal{I}}]\\#\mathcal{N}(0,0.25)+0.3~{}\text{ReflectedStep}[(-0.5,-0.5),\cdot,\\{f_{i}\\}_{i\in\mathcal{I}}]\\#\mathcal{N}(0,0.25).$

(a) On the hypercube.

(b) On the simplex.

Figure 16: Pairwise and marginals samples from the synthetic data
distribution.

#### Simplex $\Delta^{n}$.

Similarly, to parameterise the simplex as a convex polytope we set the matrix
and constraints to be given by

$A=\begin{pmatrix}-1&0&\dots&0\\\ \vdots&\vdots&\ddots&\vdots\\\
0&0&\dots&-1\\\ 1&1&1&1\end{pmatrix},\;\;b=\begin{pmatrix}0\\\ \vdots\\\ 0\\\
1\end{pmatrix}.$ (63)

Where $A$ will be a $n-1\times n$ matrix. Essentially we perform diffusion
over the first $n-1$ components of the simplex, allowing the last component to
be determined by the one minus the sum of the first $n-1$.

Similarly than for the hypercube, we construct the training and test datasets
from generated data points which are illustrated in Figure 16(b). The score
network at different times is illustrated in Figure 17.

(a) $t=0.01$.

(b) $t=0.5$.

(c) Generated distribution.

Figure 17: Evolution of the score on the simplex and generated distribution.

#### The Birkhoff polytope.

The Birkhoff polytope is the space of doubly stochastic matrices, i.e.
$B_{n}=\\{\mathrm{P}\in[0,1]^{n\times
n}:\sum_{i}^{n}P_{i,j}=1,\sum_{j}^{n}P_{i,j}=1\\}$. It is a convex polytope in
$\mathbb{R}^{n^{2}}$ and has dimension $d=(n-1)^{2}$.

(a) Data.

(b) Log-barrier.

(c) Reflected.

Figure 18: Pairwise and marginals samples on the Birkhoff polytope from
synthetic data distribution and from trained constrained diffusion models.