Diffusion Models for Constrained Domains
Authors: Nic Fishman
Leo Klarner
Valentin De Bortoli
Emile Mathieu
Michael John Hutchinson
Venue: Accepted by TMLR ()

Abstract: Denoising diffusion models are a novel class of generative algorithms that achieve state-of-the-art performance across a range of domains, including image generation and text-to-image tasks. Building on this success, diffusion models have recently been extended to the Riemannian manifold setting, broadening their applicability to a range of problems from the natural and engineering sciences. However, these Riemannian diffusion models are built on the assumption that their forward and backward processes are well-defined for all times, preventing them from being applied to an important set of tasks that consider manifolds defined via a set of inequality constraints. In this work, we introduce a principled framework to bridge this gap. We present two distinct noising processes based on (i) the logarithmic barrier metric and (ii) the reflected Brownian motion induced by the constraints. As existing diffusion model techniques cannot be applied in this setting, we proceed to derive new tools to define such models in our framework. We then empirically demonstrate the scalability and flexibility of our methods on a number of synthetic and real-world tasks, including applications from robotics and protein design.

Published in Transactions Machine Learning Research (07 (07/2023)
Diffusion Models for Constrained Domains
Nic Fishman
nigfish@gmmiil.. com
Deportment of Statistics, University of Oxford
Leo Klarner
woo.klornersssatss... oz. ac
Deportment of Statistics, University of Onford
Valentin De Bortoli
CNRS, KNS Ubm
Emile Mathicu
ebm320com
Department of Engineering, Unáversity of Cambridge
Michael Hutchinson
michael hal.biaee.aatccciinsooooooeoe or ac.
Deportment of Statistics, University of Onford
Reviewed on OpenReview: https: // openrevsew. net/ forum? id- zulTFFQ VGO
Abstract
Denoising diffusion models are a novel class of generative algorithms that achieve state-of the- art performance across a range of domains, including image generation and lexl-to-image lasks. Building this success, diffusion models have recently been extended to the Riemannian manifold setting, broademing their applicability to a range of problems from the natural and engincoring sciences. However, these Ricmannian diffusion models are built on the assumption that their forward and backward processes are well-defined for all times, preventing them from being applied to an important set of tasks that consider manifolds defined via a set inequality constraints. In this work, we introduce a principled framework bridge this gap, We present two distinct noising processes based on (i) the logurithmic burrier metric and (it) the reflected Brownian motion induced by the constraints. As existing diffusion model techniques cannot be applied in this sctting, derive new tools to define such models in our framework. We then demonstrate the practical utility of our methods on a number of synthetic and woal-workd tasks, including applications from robotics and protein design.
Introduction
Diffusion models (Sohl-Diae Dickstein et al., 2015; Song & Ermon, 2019; Song et al., 2021; Ho et al., 2020) have recently been introduced as powerful new paradigm for generative modelling. They work as follows: noise is progressively added to data following a Stochastic Differcntial Equation (SDE)-the forward noising process until it is approximately Gaussian. The generative model given by an approximation of the associated fime-reversed process called the backward denoising process. This is also SDE whose drift depends on the gradient of the logarithmic densities of the forward process, referred to as the Stein score. This score approxinated by leveraging techniques from deep learning and score matching (Hyvärinen, 2005; Vincent, 2011). Building on the success of diffusion models in domains such images and text, this framework has recently been extended to a wide range of Riemannian manifolds (De Bortoli et al., 2022; Huang et al., 2022), broadening their applicability and engincering scienees sciences-including including Lie groups such as the group of rotations SO(3), the group of rigid body motions SE(3), and many others (see e.g. Trippe et al., 2022; Corso et al., 2022; Watson et al., 2022; al., 2022; Urain al., 2022; Yim et al., 2023).
various important modelling domains from the natural
Leach

---

Published in Transactions on Machine Learning Research (07 (7/22023)
By: Unconstrained
B₁ Log-barrier
By Reflected
Time
Figure 1: The behaviour of different types of noising processes considered in this work defined on the unit interval. B,: Unconstrained (Euclidean) Brownian motion. B,: Log-barde vog-barrier forward noising process. B, Reflected Brownian motion. All sampled with the same initial point and driving noise. Black line indicates the boundary.
However, a key assumption of the Riemannian diffusion models introduced in De Bortoli et al. (2022) and Huang et al. (2022) is that the stochastic processes they consider are defined for all times. While this holds for a large class of stochastic processes, it is not the case for most. manifolds defined via a set of inequality constraints, For instance, the case of the hypercube (-1, 1)a equipped with the Euclidean metric, the Riemannian Brownian motion coincides with the Euclidean d-dimensional Brownian motion (Br)uejo.1g as long as B, E (-1,1). With probability one, (Br)c2 escapes from (-1,1), meaning that the Riemannian Brownian motion is not defined for all times and the frameworks introduced in De Bortoli et al. (2022) and Huang et al. (2022) do not apply. Such constrained manifolds comprise a wide variety of settings-incd including polytopes and convex sets of Euclidean spaces and are studied across a large nunber of disciplines, ranging from computational statisties (Morris, 2002), over roboties (Han & Rudolph, 2006) and quantum physics (Lukens et al., 2020), to computational biology (Thiele et al., 2013). Deriving principled diffusion models that are able operate directly on these manifolds thus of significant practical importance, as they enable generative modelling in data-scarce and safety-critical settings in which constraints the modelled domain may reduce the number of degrees of freedom or prevent unwanted behaviour.
As sampling problems such manifolds are important (Kook et al., 2022; Heirendt et al,, 2019), fiurry of Markov chain based methods have been developed to sample from unnormalised densities. Successful algorithms include the reflected Brownian motion (Williams, 1987; Potit, 1997: Shkolnikov & Karatzas, 2013), barrber 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 methods for statistical sampling on constrained manifolds, the geodesic Brownian motion (Lee & Vempala, 2017) and the reflected Brownian motion (Williams, 1987) are continuous stochastic processes, and thus well suited for extending the continuous Riemannian diffusion framework developed by De Bortoli et al. (2022) and Huang et al. (2022). In particular, we introduce two principled diffusion models for generative modelling on constrained domains basod on (i) the geodesic Brownian motion, leveraging tools from the log-barrier methods, and (ii) 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 demonstrate the practical utility of these methods on range of tasks defined on convex polytopes and the space of symmetric positive definite matrices, including the constrained conformational modelling of proteins and robotic arms. The code for all of our experiments is available here.
2 Background
Riemannian manifolds. A Riemannian manifold is tuple (M.g) with M smooth manifold and metric which defines an inner product on tangent spaces. The metric g induces key quantities on the manifold, such exponential map exp, : T&M + M, defining the notion of following straight lines on manifolds, gradient operator V and a divergence operator div2. It also induces the Laplace-Beitrami Beltrami operator A and
TTe (Riemannian) gradient V is defined s.t. for any smooth f E C (M), x E M, E T.M, g(=)(f(2),)) = (f())(o). "The Riemaniian divergence acts vector fields and can be defined using the volume form of M.

---

Published in Transactions Machine Learning Research (07 (07/2023)
Unconstrained Log-barrth
Figure A convex polytope defined by six con- straints (feliez. along with the log barrier po- tential, and 'straight trajectories' under the log- barrier metric and under the Euclidean metric with and without reflection at the boundary.
Reflected
consequently a Brownian motion with density (w.r.t. the volume form"), whose density is given by the heat oquation 0.p: Apt. We refer the reader to Appendix for a brief introduction to differential geometry, Lee (2013) for thorough treatment and to Hsu (2002) for details stochastic analysis manifolds.
Constrained manifolds. In this work, we are concerned with construined manifolds. More preciscly, given a Riemannian manifold (N,6), we consider a family of real functions {ʃi : - RRer indexed by We then define
M M f f(() <0, iET.
(I)
In this scenario, fffieez is interpreted a of constraints on N. For example, choosing = Rd and affine constraints f((x) (a,, r) b, E Rd, we get that .M is an open polytope as illustrated in Figure 2. This setting naturally appears many areas of engineering, biology, and physics (Boyd al., 2004; Han Rudolph, 2006; Lukens et al., 2020).
While the two methods we introduce in Section can be applied to the general framework (1), in our applications, focus two specific settings: (a) Polytopes- - R2 and (b) symmetrie positive-delinite (SPD) matrices under trace constraints N S#
Continuous diffusion models. We briefly recall the framework for constructing continuous diffusion processes introduced by Song al. (2021) in the context of generative modelling over IRd. At minimum diffusion models need four things: (i) A forward noising process converging to an invariant distribution. (ii) time reversal for the reverse process. (iii) discretization of contimuous-time process for forward/reverss process. (iv) A score matching loss - in this paper will focus on the implicit score matching loss, Song et al. (2021) consider forward noising process (Xx)sejo.r) which progressively noises data distribution Po into Gaussian N(O,, Id). More precisely (X.) an Ornstein -ueeinnnchlebkkk (OU) process which is given by the following stochastic differential equation (SDE)
dX, - dB,
X Po-
Under mild conditions Po. the time-reversed process (X Kieeo. = (Xr-theeom) also satisfies SDE (Cattiaux et al., 2021; Haussmann & Pardoux, 1986) given by
dx. 1 loggr- dB,
X pr.
where pe denotes the density of X,, This construction allows direct sampling of the forward process and leverages the Enler-Maruyane discretisation to facilitate sampling of the reverse process. Finally, the quantity log pr referred as the Stein score and is unavailable practice. can be approximated with a score network se(t,-) trained by minimising denoising score matching (dsm) loss
(()) = E A, V log Prjo (X, sa(1, x.)
(3)
We assurme that orientable therefore volume existas.

---

Published in Transactions Machine Learning Research (07 (07/2023)
or an oquivalent implicit score matching (ism)
loss
(()) = IE A (l (t X.)I div(so)(t, X.)H+
(4)
where C 2 O and X > 0 is weighting function, and the expectation is taken over U( TJ) and (X,, X,).
For arbitrarily flexible score network, the global minimiser 0* = argming((()) satisfies (f,-) V log Pt-
Inequalty-CCnstraiied Diffusion Models
We are now ready to introduce our methodology to deal with manifolds defined via inequality constraints (1). In Section 3.1, we propose Riemannian difflusion model endowed with metric induced by log-barrier potential. In Scction 3.2, we introduce a reflected diffusion model. While both models extend classical diffusion models inenuality-constraineed settings, they exhibit very different behaviours. We discuss their practical differences in Section 5.1.
3.1 Log-barrier diffusion models
Original space
Constrained space
(M,6)
(M, (2
dx;- = (ddiv(g")((X))) + a (X,)dB;
Figure 3: Convergence of the Barrier Langevin dynam- Figure 4: Illustraiive diagram of the barrier method ics on the unit interval to the uniform distribution. and the change of metric.
Barrier Langevin dynamics. Barrier methods work by constructing smooth potential Φ M + I& such that it blows up on the boundary of a desired set, see Nesterov et al. (2018). Such potentials form the basis of interior point methods in convex optimisation (Boyd et al., 2004). Of these functions, the logarithmic barrier is the most popular among practitioners (Lee & Vempala, 2017). For convex polytope M defined by the constraints Ax < b, with Rgadd and be IR"", the logarithmic barrier .M + IR, is given for any x M by
d(x)= -- Log
bi).
(5)
Assumning that || = 1, we have that for any x E M, Ø(x) = Em log(d(a, (M))), where OM, fo Rd (A,,2) 2) b}. More generally we can define for any x E M
((x) - log(d(z, 0M:))
where d(z, a.,)) computes the minimum goodesic distance from to the boundary ŒM. In general this is a highly non-trivial optimization problem, contrary to the polytope case which admits a simple closed form.
While developed and most commonly used in optimisation, barrier methods can also be used for sampling (Lee & Vempala, 2017). The core idea of barrier methods is to 'warp the geometry' of the constrained space, stretching it the process approaches the boundary explicitly deal with the boundary. Assuming ø to be strictly convex and smooth, its Hossian V is positive definite and thus defines valid Riemannian metric on M. The formal approach to 'warping the geometry' of the convex space with the boundary is to endow M with the Hessian Riemannian metric g V making it into Hessian Manifold, see Shima & Yagi (1997). In the special case where the barrier is given by (5), we get for any M
that it never hits it, hence bypassing the need to
g (xx) ATS with S(x) diag(b, (Ai,)))

---

Published in Transactions Machine Learning Research (07 (07/2023)
Equipped with this Riemannian metric, we consider the following Langevin dynamics as a forward process
dx, = =dii(g=")(xx)dd + o(X.)=
(6)
with (x) & (div(Ff)@o....... dir)).....di(((aa)))) (za))) for any smooth F: IR- - IRd. Under mild assumptions on M, div(g-)) and g² we get that (X,)22 is well-defined and for any 2 0, X, E M. In particular, for any I 2 0, X, does not reach the boundary. This stochastic process was first proposed by Lee & Vempala (2017)
in the context of efficient sampling from the uniform distribution over a polytope. Under similar conditions, X, admits density Pt w.r.t. the Lebesgue measure and have that 0,p, (T Pp))) In addition, (Xi)ee is irreducible. Hence, assuming that M is compact, the uniform distribution on M is the unique invariant measure the process (Xe) and (X)) converges to the uniform distribution some sense. We refer the reader to Appendix C for proof of these results.
Time-reversal. Assuming that g and derivative are bounded on M, the Lime-reversal of (6) is given by Cattiaux et al. (2021), in particular we have
ax,- | ddv(( div(g ")+ V logprr- (X, jt a AdB.,
= divi !) + g ' loggrr- (X.)dt + g x,) -x)xxddl dB,
X, is initialised with the uniform distribution on MM (which is close to the one of Xy for large T))
The estimation of the score term lo is done by minimising the 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 (7) processes, once the score learnt, requires discretisation scheme. We use Geodesic Random Walks (GRW) (Jergensen, 1975) for this purpose, see Amgorithm 1. This discretisation is generalisation the Fuler-Maruyama discretisation of SDE Euclidean spaces, where the operator is replaced by the exponential mapping the manifold, computing the geodesics.
Algorithm 1 Geodesic Random Walk. Discretisation of the SDE dX, d(l, X, + dB,-
Require: T (simulation time), N (number of steps), X3 (initial point), ol (drift function) = T/N
for E (O,...,N-- N.1.N-133 do
Zk+1 ~ Id),Id)
WK+1 = 7d((?y, Xx) + 4+Z
X expx, [Wk+1] proju Ws+1)
return LN
However, contrary to De Bortoli et al. (2022), we will not have access explicitly to the exponential mapping of the Hessian manifold. Instead, we rely on an approximation, using a retraction (see Absil & Malick (2012); Boumal (2023) for definition and alternative schemes).
3.2 Reflected diffusion models
Another approach to deal with the geometry of M is use the standard metric h and forward dynamics of / and constraining it to M by reflecting the process whenever it would encounter boundary. We will first assume that M is compact and convex. To simplify the presentation, we focus on the Euclidean case N == Rd with smooth boundary OM
The key difference between this approach and the barrier approach is that in the reflected case we leave the goometry unchanged, so all we need to do show that the dynamics induced by reflecting the forward process whenever it hits the boundary leads to an invariant distribution and admits time-reversal.
We refer Appendix for definition of smooth boundary.

---

Published in Transactions Machine Learning Research (07 (7/2202)
dB, = dB, dk.
[k/c
Figure 5: Left: Convergence of the rellected Brownian motion on the unit interval to the uniform distribution. Right: Value of [kl, for the trajectory samples on the left through time.
It is worth noting that while barrier approaches have received considerable theoretical attention in the sampling literature (Lee Vempala, 2017; Noble et al., 2022), reflected methods have remained comparatively undeveloped from the methodological and practical point of view. In the next section, we recall the basies of reflected stochastic processes.
Skorokhod problem. The reflected Brownian motion is defined as the solution to the Skorakhod problem. Roughly speaking a solution to the Skorokhod problem consists of two coupled processes, k,,,k) such that. (Bz)220 acts locally as Euclidean Brownian motion (B2)20 and k, compensates for the excursion of (Bz)c20 so that (Br) remains in M. We say that (B,, kr) is a solution to the Skorokhod problem (Skorokhod, 1961) if (kr) and (Br)220 are two processes satisfying mild conditions, see Appendix for a rigorous introduction, such that for any 2 0,
B By B. - k. E M,
(8)
and [k|, - Ju I, ck, k, = So n(B, (BBB)dJkk where (Ikl)leo is the total variation of (k,)120 and we recall that n is the outward normal to M When (Bz)z2 hits the boundary, the condition k, f0 (BB)ddlk), tells us that -k, "compensates" for B, by pushing the process back into M along the inward normal -n, while the condition |k], = S ea, can be interpreted as k, being constant when (Bz)r20 does not hit the boundary, As a result (B,)20 can be understood as the continuous- Gaussian random walk. The process (k:) can be related to the notion of local time (Revuz & Yor, 2013) and quantifies the amount of time (Bz)22 spends at the boundary a.M Lions & Sznitman (1984, Theorem 2.1) ensure the existence and uniqueness of a solution to the Skorokhod problem. One key observation that the event (B, E 3M} has probability zero (Harrison & Williams, 1987, Section 7, Lemma 7). As in the unconstrained selting, one can describe the dynamics of the density of B,-
counterpart to the reflected
Proposition 3.1 (Burday et al. (2004)). For any >>> O, the distribution of B, admils densily w.r.t. the Lebesgue measure denoted Pe- In addition, we have for any x E int(M) and No E aM
@x)(()) [)pp()) OuPe(2o) 0,
where we recall that n is the outward normal to M.
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 M exponentially fast (Loper, 2020; Burdzy et al., 2006), see Section 3.2. Hence, (Bz)220 is a candidate for a forward noising process 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-renersal for (B.)e Namely, we need to characterise the evolution of (Xx)eejo.1]= = (Br=thejoo7) It can be shown that the time-reversal of (B,)sejorr; is also the solution to a Skorokhod problem.
We extend the normal n the whole space by letting n(ar) = Of x a..

---

Published in Transactions Machine Learning Research (07 (2/220
Algorithm Reflected step algorithm. The algorithm operates by repeatedly taking geodesic steps until one of the constraints is violated, or the step is fully taken. Upon hitting the boundary we parallel transport the tangent vector to the boundary and then reflect against the boundary. We then start a new geodesic from this point in the new direction. The arg intersect, function computes the distance must travel along geodesic in direction til constraint f. is intersected. For a discussion of paralleltrnsport,, exp, and reflect please see Appendix
Input: x E M, v E T,M, (fihez
|v||_
2/ e//lell
while f do
d: arg intersoct, exp, (a,ts), ts), Ss]
i + min; de s.t. d, > 0
min(d;,
z' exp, (I, as)
8 < paralleltransport (1,, 8,2)
reflect(s, f()
f f
I + x²
return of
Algorithm 3 Reflected Random Walk, Discretisation of the SDE dX, b(t, )tt + dB, dk.
Require: T (simulation time), N (number of steps), X6 (initial point), (fikez (boundary functions) 7 = T/NN
for E (O,...,N--1) do
Zk41 N(0, Id)
X 1+1 ReflcctedStep/X7, Z (fihez)
return {
Theorem 3.2. There exist (k k)220 bounded veriation process and Brownian motion (Br)s20 such that
c, xo B, + So V Pl (x,)ds kr
In addition, for any t e [0,T] we have
kk| = 1 d/k/. k- = So (xxxx kkkk. ((a
The proof, see Appendix G, follows Petit (1997) which provides a timo-reversal in the case where M the positive orthant. It is based careful handling of the boundary conditions, In particular, contrary to Petit (1997), we do not rely on an explicit expression of p. but instead the intrinsic properties of (k,)20- Informally, Theorem 3.2 means that the process (X.)eejo.r) satisfies
extension of Haussmann & Pardoux (1986) to the reflected setting, with
dx, V log Lo (X) + dB; ddk..
(10)
which ochocs the usual Lime-reversal formula (2). practice, in order sample from (X.)e one needs to consider the reflected version of the unconstrained dynamics d',, V ( dB,
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) and Bossy et al. (2004) for weak convergence results this numerical scheme in the Fuclidean setting.

---

Published in Transactions Machine Learning Research (07 (07/2023)
Figure Rellection against a linear boundary. For step with magnit [s] and direction 8 the distance to the boundary described by the normal At and offset b is d = (8.m2) The reflected direction is given 8 = 22, A/ú. (3,8)
20
Figure 7: Reflection against a spherical boundary. For a sphere of radius Y, the distance the boundary from Ht in direction is given by d [((8, 22)² 4(r² - ||,,,1 1)) / 2 (8, z²). The normal at the intersoction can be computed as the unit vector in the direction - -2 + 2z), and then s' as above.
f
27
d
Likelihood evaluation. In the case of a reflected Brownian motion, it is possible to compute an equivalent ODE in order perform likelihood evaluation. The associated ODE was first derivod in Lou Ermon (2023). The form of the ODE and the proof that it remains in M are postponed to Appendix F.
3.3 Loss and score network parameterisation
In order to train log-barrier and reflected difffusion models we prove that can Iraclable matching loss in constrained manifolds. We will prove that the implicit score-matching loss leads to the recovery of the correct score when we have boundary, so long as we enforce that the score zero on the boundary (sce Appendix E.1). This proof holds for both the log- barrier and the reflected process.
Proposition 3.3. Let E ( T] x 1R Rdd) such that for any 3M and 1 2 0, se(x) = 0. Then, there caists > such that
E
V lo st E[||s, +2 div(s:) +C,
where E is laken over X, ~ Ps and (((O (]).
This result immediately implies can optimise the network using the ism loss function long enforce a Neumann boundary condition. The estimation of the score term V log ps is also done by minimising the ism loss function (4).
The score term V log Pr- appearing in both time-reversal processes (7) and (10) is intractable. It is thus approximated with a score network sa(t,-) z Vlog Pt. We use multi-layer perceptron architecture with sin activation functions, see Appendix J for more details on the experimental setup. Due to boundary condition (9), the normal component of the score is zero at the boundary in the reflected case. A similar result holds in the log- log-barrier setting. This additionally required for our proof of the ism loss.
Following Liu et al. (2022), we can accommodate this in our score parameterisation by additionally scaling the score output of the noural network by a monotone function h(d(z, a.))) where d is the distance from x the boundary and f(0)= 0. In particular,
use a clipped ReI.U function: st(t, min(1, Rel.U(d(x, 0.) 5)) NNo(tt,) z) with 5 > a threshold so the model is forced to be zero "close" to the boundary as well as exactly on the boundary, see Appendix for illustration The inclusion of this scaling function is necessary to produce reasonable results as we show in Appendix E.2. The weighting function in (4) is set to A (f+1)
The forward processes (6) and (8) for the barrier and reflected methods cannot be sampled in closed form, so at training time samples from the conditional marginals (X)) are obtained by discrctising these processcs. As to take the most of this computational overhead, we use several samples from the discrotised forward trajectory (X.,...- X., Xx) instead of only using the last sample.

---

Published in Transactions on Machine Learning Research (07 (07/2023)
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 within 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 & Vompala, 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) and Noble et al. (2022) for instance. If we keep the Euclidean metric, then unconstrained stochastic processes might not be well-defined for all times. To counter this elfect, it has been proposed to reflecl 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 manifolds. De Bortoli et al. (2022) extended the work of Song 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 maximun 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. Their specific structure can be leveraged to define efficient diffusion models (Yim et al., 2023). Leach et al. (2022) define diffusion models on SO(3) for rotational aligument, while Jing et al. (2022) consider the product of tori for molecular conformer generation. Corso el. al. (2022) use diff'usion models on R² x SO(3) x SO(2) for protein docking applications. RFDiffusion (Watson et al., 2022) and FrameDiff also incorporate SE(3) diffusion models. Finally, Urain et al. (2022) introduce a methodology for SE(3) diffusion models with applications to robotics.
Comparison with Lou Ermon (2023). We now discuss few key differences between our work and the reflected diffusion models presented in Lou & Ermon (2023). First, in the hypercube setting, methodologies are identical from a theoretical viewpoint. The main difference is that Lou & Ermon (2023)
use an approxiamted version of the DSM loss (3), whereas we rely on the ISM loss (4). Lou & Ermon (2023) exploit the specific factorised structure of the hypereube to make the DSM loss tractable, leading to significant practical advantages. These can however be directly employed our framework for reflected models, and also in the log-barrier setting. Second, in our work, we target scientific applications where the underlying geometry is not Euclidean, whereas Lou & Ermon (2023) focus on the case where the constrained domain of interest is the hypercube (a subeet of Euclidean space) with image applications, or where the domain can be easily projected into the hypercube, such as the simplex. Our approach is designed to handle wider range of settings, such as non-convex polytopes in Euclidean space, or non-Euclidean geometries.
5 Experimental results
To demonstrate the practical utility of the constrained diffusion models introduced in Section 3, we evaluate them on serics of increasingly difficult synthotic tasks on convex polytopes, including the hypercube and the simplex, in Section 5.1. We then highlight their applicability to real-world settings by considering two problems from robotics and protein design. In particular, we show that our models are able to learn distributions over the space of d x d symmetric positive definite (SPD) matrices S4 under trace constraints in Section 5.2-a a setting that is essential to describing and controlling the motions and exerted forces of robotic platforms (Jaquier al., 2021). In Section 5.3, we use the parametrisation introduced in Han & Rudolph (2006) to map the problem of modelling the conformational ensombles of proteins under positional constraints on their endpoints to the product manifold of a convex polytope and a torus. The code is available here. We refer the reader to Appendix J for more details on the experimental setup.

---

Published in Transactions Machine Learning Research (07/2023)
Table 1: MMD metrics between samples from syuthetic distributions and trained constrained and unconstrained (Euclidean) diffuusion models. Means and confidence intervals are computed over different. runs.
Log-barrier
Reflected
Euclidean
Space
MMD
% in M
MMD
% in M
MMD
% in .M
0066+.006
100.0
.055..015
100.0
.0622.011
98.8
[-,,1]
207
100.0
080:+.001
100.0
.076+.000 001
98.5
10
+30
100.0
.383..048
100.0
.081
96.4
.050+.012
100.0
.043+.002
100.0
.055+
96.4
A
2.38
100.0
+181
100.0
.068+.014
96.3
.275+.015 015
100.0
.290 .009
100.0
.060#.00 008
92.6
Data
Log-bane Log-barrier
Reflected
Dala
Log-barrier
Reflected
(a) 21) square data.
(b) 21) Dirichlet data.
Figure Histograms of samples from the data distribution and from trained constrained diffusion models.
5.1 Method characterisation convex polytopes
First, we aim to assess the empirical performance of our methods on constrained manifolds of increasing dimensionality. For this, we focus on polytopes and construct synthetic datasets that represent bimodal distributions. In particular, Appendix J.1 we also present results on the Birkhoff polytope. We quantify the performance of each model via the Maximum Mean Discrepancy (MMD) (Gretton et al., 2012), which is a kernel- based metric between distributions. We present qualitative comparison of the logarithmic barrier and reflected Brownian motion models in Figure 8, and observe that both hods recover the data distribution on the two-dimensional hypercube and simplex, although the reflected method seems to produce better fit. In Table 1, we report the MMD between the data distribution and the learnt diffusion models. Here, similarly observe that the reflected method consistently yields better results than the log-barrier one.
investigate two specific instances of polytopes: hypercubes and simplicos. In
We additionally compare both of these models to set of unconstrained Euclidean diffusion models, noting that they are outperformed by the constrained models in lower dimensions, but generate better results in higher dimensions. There are number of potential explanations for this: First, we note that the mixture synthetic dala-generating process places significantly less probability
of Normal distributions we use mass near the boundary as its dimensionality increases, more closely resembling an unconstrained mixture distribution that casier for the Euclidean diffusion models to learn, while posing challenge to the log- lorrbarrier and reflected diffusion models that initialise at the uniform distribution within the constraints. Additionally, we note that the design space and hyperparameters used all experiments were informed by best. practices for Euclidean models that may be suboptimal for the more complex dynamics of constrained diffusion models.
5.2 Modelling robotic arms under force and velocity constraints
Accurately determining and controlling the movement and exerted forces of robotic platforms fundamental problem in many real-world robotics applications. A kinctostatic descriptor that is commonly used to quantify the ability of a robotic arm to move and apply forces along cortain dimensions is the so-called manipulability ellipsoid E e ıd (Yoshikawa, 1985), which is naturally described symmetric positive definite (SPD)

---

Published in Transactions on Machine Learning Research (07/2023)
(a) Samples from the data distribution.
(b) Samples from our log- barrier diffusion model.
(c) Samples from our re- flected diffusion model.
(d) Samples from the uni- form distribution.
Figure 9: Samples in + x R from (a) the data distribution, (b) our log-barrier diffusion model, (c) our reflected diffusion model and (d) the uniform distribution. Each sample visualised as the manipulability ellipsoid encoded by the SPD matrix M E 87 placed al the corresponding location in g². Additional results and full correlation plots are postponed to Appendix J.2.
matrix M E IRdxd (Jaquier et al., 2021). The manifold of such d x d SPD matrices, denoted as 54 is defined as the of matrices {x Mr 2 O, z R M Rdxd), In many practical settings, 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 M, i., as the inequality constraint >i= My < with C O. 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 docomposition. Each SPD matrix has a unique decomposition of the form = L.I.T where I, 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 Tr(M)- = Tr(LI.T) = Eis 12 and results in ball of radius C. We additionally model the two-dimensional
the constraint. on the entries of L to live position of the arm. In summary, the space over which we parameterise the diffusion models is defined as [L E Rd(d+1)/2, Lis > 0, E, L2, < <C} IR.. Under the Euclidean metric, we can apply both our log- barrier and reflected approaches. The positive diagonal (linear) constraint is handled similarly to the polytope sctting. The reflection on the ball boundary is defined and illustrated on Figures 6 and 7.
Using this framework, we model the datasets presented in Jaquier et al. (2021) (see Appendix H for full experimental details). The joint distribution over the SPD matrices (represented as ellipsoids) and their positions is presented in Figure 9. We qualitatively observe that the reflected method is able to model the joint data distribution better than the barrier one. This is reflected by an MMD of 0.161 and 0.247, respoctively.
5.3 Modelling protein loops 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 discovery. In many data-scarce practical settings such as antibody or enzyme design, it 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.
This problem can be reformulated as modelling a spatial chain with spherical joints and fixed end points. Following the framework outlined in Han & Rudolph (2006), we parametriss the conformations of such chains with d fixed-length links and arbitrary end-point constraints as the product of convex polytope P and torus T: p(d-3) x T(d-2), The essential ial idea of 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 adjacent joints in the chain. A point in the polytope corresponds to the lengths of the diagonals of these triangles, and a point in the torus to the angles between each pair of subsoquent triangles. See Figure 10 for an illustration and Appendix I.1 for full description.

---

Published in Transactions Machine Learning Research (07/2023)
(a) An illustrative diagram of the parameeterisation used for the conformational modelling of the Co trace of eyclic peptide, introduced in Han Rudolph (2006).
(b) The convex polytope coustraining the diagonals
of the triangles for the given bond lengths in the illustrated molecule. The total design space the
product of this polytope with the 4D flat torus.
Figure 10: Illustrations of the carametrisation used to model distribution of polypeptide backbone conforma- tions under anchor point distance constraints as the product of a convex polytope P and torus T: p x T
(a) Samples from the data distributivn.
(b) Samples from our log- barrier difffuson model. flected dilfusion model. form distribution.
(c) Samples from our re-
(d) Samples from the uni-
Figure 11: Planar projection of the modelled C, chains from (a) the training dataset, (b) our log- barrier diffusion model, (c) our reflected diffusion model and (d) the uniform distribution. Additional results and full correlation plots are postponed to Appendix J.3.
Using this framework, we model the conformational landscape of the cyclic peptide C-AAGAGG, consisting of a circular polypeptide chain with coinciding endpoints. We generale 10th backbone conformations using tools from molccular dynamics (Eastman et al., 2017; Hornak et al., 2006) and divide them into training and evaluation datasets, (see Appendix 1.2 for full experimental details). Drawing on the definition above, the space on which we effectively paramcterise our constrainod diffusion models for circular polypeptide chain of length df = 6 is given by the product manifold gpl x TT.
To learn a distribution over this space, we leverage the methodology introduced Section 3 the polytope component Oo and in De Bortoli et al. (2022) for the torus component T A qualitative comparison of samples from the data distribution, our log-barrier and reflected difflusion models, and the uniform distribution presented Figures 1la 11d. For enhanced visual clarity, we project the modelled spatial chain onto the 2D plane by removing the (unconstrained) torus component of the product manifold and only plotting the planar chains encoded by the (constrained) polytope component (a correlation plot the full product manifood is presented in Figure 25).
is apparent that the data distribution is highly multimodal, encompassing a large number of locally optimal conformational clusters. Nevertheless, our reflected and log- far-barrier diffusion models are able to robustly approximate this energetic landscape, producing samples that reflect key conformational states and producing comparable MMD metrics of 0.03210.021 and 0.032 0001, respectively. As a point of comparison, the uniform distribution the polytope-torus product has an MMD of 0.1 12+0.001.

---

Published in Transactions Machine Learning Research (07 (07/2023)
Discussion
Learning complex distributions whose support is confined to constrained spaces is a crucial task with many applications in the natural and engineering sciences, including computational statistics (Morris, 2002), robotics (Han & Rudolph, 2006), quantum physics (Lukens et al., 2020) and computational biology (Thiele et al., 2013). In this work we extend continuous diffussin models to this setting, proposing two complementary approaches one based on log-barrier methods and the other on the reflected Brownian motion. For both methods, we derive the time-reversal formula, proposc discrctisation schemes and extend the score-matching toolbox. We demonstrate the utility of our methods on a range of syuthetic and real-world tasks, including the constr conformational modelling of proteins and robotic arms, and find that reflected methods, while enjoying fewer theoretical guarantees than their log- barrier counterparts, often yield preferable results.
We conclude by highlighting important directions of future research. First, the computational cost of performing the reflection when discretising the reflected Brownian motion is high. Finding numerically efficient approximations of the reflected process therefore necessary extend this methodology to very high dimensional scttings. Second, the retraction used in place of the exponential map for the barrier method leads to high number of discretisation steps to ensure good approximation. Designing a faster forward process for the log- barrier method key largeting more complex distributions.

---

Published in Transactions Machine Learning Research (07 (07/2023)
References
P-A Absil and Jérôme Malick. Projection-like retractions matrix manifolds. SIAM Journal on Optimizzation, 22(1):135-158, 2012.
RN Bhattacharya. Criteria for recurence and existence of invariant measures for multidimensional diffusions. The Annals of Probubilily, pp. 541-553, 1978.
Mircille Bossy, Emmanuel Gobet, and Denis Talay. A symmetrized euler scheme for an efficient approximation of reflected diffusions. Journal of opplied probability, 41(3):877 889, 2004.
Nicolas Boumal. An introduction to optimization on smooth manifolds. Cambridge University Press, 2023.
Stephen Boyd, Stephen P Boyd, and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2004.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-MMlnee, and Qiao Zhang. JAX: composable transformations of Python+ programs, 2018. URL http:// http:://gitte htt///gtttbt///gg/ www/ jaa/
Krzysztof Burdzy, Zhen-Qing Chen, and John Sylvester. The heat equation and reflected brownian motion in time-dependent domains. The Annals of Probability, 32(1B):775 804, 2004.
Krzysatof Burdzy, Zhen-Qing Chen, and Donald E Marshall. Traps for reflected brownian motion. Mathema- tische Zeitschrift, 252:103-1388 132, 2006.
Patrick Cattiaux, Giovanni Conforti, Ivan Gentil, and Christian Léonard. Time reversal of diffusion processes under a finite entropy condition. arXio preprint arXim:2104.0000 2021.
Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. DiflDook: Diffusion Steps, Twists, and Turns for Molecular Docking, October 2022. URL http:///axxxvv.orabbbaeeee 01776.
Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thomton, Yee Whye Teh, and Arnaud Doucet. Riemannian score-based generative modcling, 2022.
Patrick Dougherty, Ashweta Sahni, and Dohua Pci. Understanding coll penctration of cyclic poptides. Chemical Reviews, 1024 10287, 2019.
Peter Eastman, Jason Swails, John Chodera, Robert 'T' McGibbon, Yutong Zhao, Kyle A Beauchamp, Lee-Ping Wang, Andrew C Simmonett, Matthew P Harrigan, Chaya D Stern, et al. Openmm 7: Rapid development. high performance algorithms for molecular dynamics. PLoS compulational biology. 13(7): c1005659, 2017.
Lawrence Craig Evans and Ronald Gariepy. Measure theory and fine properties of functions. CRC press, 2015.
Khashayar Gatmiry and Santosh S Vempala. Convergence of the riemannian langevin algorithm. orXiv preprint orXi::200,.10118 2022.
Gene H Golub and Charles F Van Loan. Matriz computations. JHU press, 2013.
Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel Iwo-sample test. Journal of Machine Learning Research, 13(mull):723-777 March 2012. ISSN 1532-4435.
Li Han and Lee Rudolph. Inverse kinematics for serial chain with joints under distance constraints. In Robotics: Science and systems, 2006.
J Michael Harrison and Ruth J Williams. Brownian models of open qucucing networks with homogencous customer populations. Stochastics: An International Journal of Probability and Stochastic Processes, 22(2): 77 115, 1987.

---

Published in Transactions Machine Learning Research (07 (07/2023)
Ulrich G Haussmann and Etienne Pardoux. Time reversal of diffusions. The Annals of Probobility, pp-
1188 1205, 1986.
Laurent Heirendt, Sylvain Arreckx, Thomas Plau, Scbastián N Mendoza, Anne Richelle, Almut Heinken, Hulda S Haraldsdóttir, Jacek Wachowiak, Sarah M Keating, Vanja Vlasov, et al. Creation and analysis biochemical constraint-based models using the cobra toolbox 3.0. Nature protocols, 14(3):639-702, 2019.
Jonathan Ho, Ajay Jain, and Picter Abbecl. Denoising diffusion probabilistic models. Advances in Newral Information Processing Systems, 2020.
Viktor Hornak, Robert Abel, Asim Okur, Bentley Strockbine, Adrian Roitberg, and Carlos Simmerling. Comparison of multiple amber force fields and development of improved protein backbone parameters. Proteins: Structure, Punction, and Bioinformaties, 65(3):712-725, 2006.
Elton P Hsu. Stochastic analysis on manifolds. Number 38. American Mathematical Soc., 2002.
Chin-Wer Huang, Jae Hyun Lim, and Aaron C Courville. A variational perspective on diffusion-basse generative models and score matching. Advances Neural Information Processing Systems, 34:22863- 22876, 2021.
Chin-Wei Huang, Milad Aghajohari, Avishek Joey Bose, Prakash Panangaden, and Aaron Courville. Rieman- nian Diffusion Models, August 2022.
Aapo Hyvärinen. Estimation of non-nonn normaooidledd statistical models by score matching. Journal of Machine Learning Reseorch, 6(4), 2005.
Nobuyuki Ikeda and Shinzo Watanabe. Stochastic differential equations and diffusion processes. Elsevier, 2014.
Noémie Jaquier, Leonel Rozo, Darwin G Caldwell, and Sylvain Calinon. Geomelry-aware manipulability learning, tracking, and transfer. The International Journol of Robotics Research, 40(2-3):624 650, 2021.
Bowen Jing, Gabriele Corso, Jelfrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional Diffusion for Molecular Conformer Generation, June 2022. URL attp://arxx
Erik Jørgensen. The central limit problem for gendesic random walks. Zeitschrift für Wutrachheiniccheettttteeee und verwandte Gebiele, 32(1-2):1 64, 1975.
Weining Kang and Kavita Ramanan. On the submartingale problem for reflected diffusions in domains with piecewise smooth boundaries. 2017.
Ravi Kannan and Hariharan Narayanan. 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 570, New York, NY, USA, 2009. Association for Computing Machinery. ISBN 9781605585062. doi: 10. 15555/56111..16666
URL https://doi: htttp://ddi..org/10..oo 1.0.1145//136411 111.///666111.11499.... 1536491.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2014. URL http:
//arxiv. perrivv.org/aai abssa412..989
Yunbum Kook, Yin Tat. Lee, Ruoqi Shen, and Santosh S Vempala. Sampling with riemannian hamiltonian monte carlo in constrained space. arXiv preprint 2XXi:2202.0190: 2022.
Thomas J Lane. Protcin strueture prediction has reached the single-structure fronticr. Nature Methods, pp. 1 4, 2023.
Adam Leach, Sebastian M Schmon, Mattco T Degiacomi, and Chris G Willcocks. Denoising Diffusion Probabilistic Models on SO(3) for Rotational Alignment. ICLR 2022 Workshop Geometrical and Topological Representation Learning, pp. 8, 2022.
John M Lee. Riemannian manifolds: an introduction to carvature, volume 176. Springer Science & Businoss Media, 2006.

---

Published in Transactions Machine Learning Research (07 (07/2023)
John M Lec. Smooth manifolds. In Introduction to Smooth Manifolds, pp. 1-31. Springer, 2013.
Yin Tat Lee and Santosh Vempala. Geodesic walks in polytopes. In Proceedings of the 49th Annual A ACM SIGACT Symposium Theory of Computing, pp. 927-940, Montrcal Canada, June 2017. ACM. ISBN
10011-450-45228--4 doi: 10.1145/ 35555//35559996/ 3055446..
Yin Tat Lee and Santosh S Vempala. 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.
Pierre-Louis Lions and Alain-Sol Sznilanan. Stochastic differential equations with reflecting boundary conditions. Communications on pure and applied Mathematics, 37(4):511-58) 537, 1984.
Song Liu, Takafumi Kanamori, and Daniel J Williams. Estimating density models with truncation boundaries using score matching. Journal of Machine Learning Research, 23(186):1-38, 2022.
Jackson Loper. Uniform ergodicity for brownian motion in bounded convex set. Journal of Theoretical Probabilityy, 33(1):22 35, 2020.
Aaron Lou and Stefano Ermon. Reflected diffusion models. arXiv preprint (xXi::8344.047400 2023.
László Lovász and Santosh Vempala. Hit-and-rum from corner. SIAM Journal Computing, 35(4): 085-1005, 2006. doi: 10. 10. 11113//5000733775547727
Joseph M Lukens, Kody JH Law, Ajay Jasra, and Pavel Lougovski. A practical and efficient approach for bayesian quantum state estimation. New Journal of Physics, 06600888303, 2020.
Sean P Meyn and Richard L Tweedie. Stability of markovian processes ii: Contimuous-time processes and sampled chains. Advances Applied Probubility, 25(3):487-617, 1993.
Nina Miolane, Nicolas Guigui, Alice Le Brigant, Johan Mathe, Benjamin Hou, Yann Thanwerdas, Stefan Heyder, Olivier Peltre, Niklas Koep, Hadi Zaatiti, Hatem Hajri, Yann Cabanes, Thomas Gerald, Paul Chauchat, Christian Shewmake, Daniel Brooks, Bernhard Kainz, Claire Donnat, Susan Holmes, and Xavier Pennec. Geomstals: A python package for riemannian geomelry in machine learning. Journal of Machine Learning Research, 21(223):1-9, 2020. URL http: /mmrr. http:appers/v21/111/0
.tml.
Ben J. Morris. Improved bounds for sampling contingency tables. Random Structures & Algorithms, 21(2): 135-146, 2002. doi: //s:///orirrg/////// (ooiioor/11.1000 10.1000/1 10049. URL https: www.bbrrrrrrrllll...........eeee /oii Abs/10.1002/ raaa.10033
Yurii Nesterov et al. Lectures on conver optimization, volume 137. Springer, 2018.
Maxence Noble, Valentin De Bortoli, and Alain Durmus. Barricr Hamiltonian Monte Carlo, October 2022.
David Nualart. The Molliavin calculus and related topics, volume 1995. Springer, 2006.
Barbara Pacchiarotti, Cristina Costantini, and Flavio Sartoretto. Numerical approximation for functionals of reflecting diffusion processes. SIAM M Journal on Applied Mathematics,, 58(1):73-100 102, 1998.
Frédérique Petit. Time reversal and reflected diffusions. Stochastic Processes and their Applications, 69(1): 25-56 53, July 1997. ISSN 03044149. doi:
Daniel Revuz and Mare Yor. Continuous martingales and Browniun motion, volume 293. Springer Science & Business Media, 2013.
Piotr Rotkiewicz and Jeffrey Skolnick. Fast procedure for reconstruction of full-atom protein models from reduced roprosentations. Journal of computational chemistry. 2969))146011465 2008.
Laurent Saloff-Coste. Preciso estimates on the rate at which cortain diffusions tend to equilibrium. Mathema- Bische Zeitschrift, 217(1):641 677, 1994.

---

Published in Transactions Machine Learning Research (07 (07/2023)
Hirohiko Shima and Katsumi Yagi. Goometry of hessian manifolds. Differential geometry and applications, 7(3):277 290, 1997.
Mykhaylo Shkolnikov and Toannis Karatzas. Time-reversal of reflected Brownian motions in the orthant, July 2013. URL http://arxxv..rg// http///aaa///...... 1t////337..44444 .4222.
Anatoliy V Skorokhod. Stochastic equations for diffusion processes in bounded region. Theory of Probability & Hss Applications, 6(3):264-217 1961.
Robert L. Smith. Efficient Monte Carlo Procedures for Generating Points Uniformly Distributed over Bounded Regions. Operations Research, 32(0):1296-1308, December 1984. doi: 10.1287/opre.32...1 6........... html.
URL https: www.:deasrreee..e. http:///////rrppe////////89/2210
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrimm thermodynamics. In International Conference Machine Learning, pp.
PMLR, 2015.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems, 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishck Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
Ines Thiele, Neil Swainston, Ronan M T Fleming, Andreas Hoppe, Swagatika Sahoo, Maike K Aurich, Hulda Haraldsdottir, Monica L Mo, Ottar Rolfsson, Miranda D Stobbe, Stefan G Thorleffsson, Rasmus Agren, Christian Bölling, Sergio Bordel, Arvind K Chavali, Paul Dobson, Warwick B Dumn, Lukas Endler, David Hala, Michael Hucka, Duncan Hull, Daniel Jameson, Neema Jamshidi, Jon J Jonsson, Nick Juty, Sarah Keating, Intawat. Nookaew, Nicolas Le Novère, Naglis Malys, Alexander Mazcin, Jason A Papin, Nathan D Price, Evgeni Selkov, Martin I Sigurdsson, Evangelos Simconidis, Nikolaus Sonnenschein, Kieran Smallbone, Anatoly Sorokin, Johannes H G M van Beek, Dieter Weichart, Igor Goryanin, Jens Nielsen, Hans V Westerhoff, Douglas B Kell, Pedro Mendes, and Bernhard Ø Palsson. A community-drrvenr global reconstruction of human metabolism. Nature Biotechnology. March 2013. ISSN 1087-0156. doi: 10.1038 1..1008/nb..20 n22.8888..
Brian L. Trippe, Jason Yim, Doug Tischer, Tamara Broderick, David Baker, Regina Barzilay, and Tommi Jaakkola. Diffusion probabilistic modeling of protein backbones in 3D for the motif-scaffolding problem, June 2022.
Julen Urain, Niklas Funk, Georgia Chalvatzaki, and Jan Peters. Se (i)-difuusionfield:: Learning cost functions for joint grasp and motion optimization through diffusion. ar preprint ar KK::220900383 2022.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Computation, 23 (7):1661-1674, 2011.
Joseph L. Watson, David Juergens, Nathaniel R. Bennett, Brian L. Trippe, Jason Yim, Helen E. Eisenach, Woody Ahern, Andrew J. Borst, Robert J. Ragotte, Lukas F. Milles, Basile I. M. Wicky, Nikita Hanikel, Samuel J. Pellock, Alexis Courbet, William Sheffler, Jue Wang, Prectham Venkatesh, Isaac Sappington, Susana Vázquez Torres, Anna Lauko, Valentin De Bortoli, Emile Mathieu, Regina Barzilay, Tommi S. Jaakkola, Frank DiMaio, Minkyung Baek, and David Baker. Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models, December 2022.
R. J. Williams. Reflected Brownian motion with skew symmetric data in polyhedral domain. Probability Theory and Related Fields, 75(4):450-48) 485, August 1987. ISSN 0178-8051, 1432-2064. doi: 10.1007/ BF00320328.
Kevin E Wu, Kevin K Yang, Rianne van den Berg, James Y Zou, Alex X Lu, and Ava P Amini. Protein structure generation via folding diffusion. Xrx preprint ar 2209: 15611, 2022.

---

Published in Transactions Machine Learning Research (07 (77/2023)
Jason Yim, Brian Trippe, Valentin De Bortoli, Emile Mathicu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. Se (3) diffusion model with application to protein backbone generation. orn preprint arXiv:2302.088 02277, 2023.
Tsunco Yoshikawa. Manipulability of robotic mochanisms. The international journal of Robotics Research, 4 (2):3-9, 9, 1985.


Reviews:
[{"review_id": "0Hh44TMTqY", "reviewer": "Reviewer_926N", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n\n- The authors propose two methods for diffusion modeling relevant for problems defined on Riemannian manifolds with boundary conditions.\n- Extensive numerical experiements.\n- The paper is mathematically sound, with detailed derivations in the supplement.\n\n\n\n\n**Weaknesses:**\n\n- The log barrier method does not seem to work so well for constraining the diffusion in the numerical experiments (especially the ones in the supplement). Can the authors comment on that? What is the reason behind this and how could it potentially mitigated?\n- The scaling function used for estimating the score on the truncated domain has been actually introduced before, see the comments below.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper considers the training of diffusion models on Riemanninian manifolds with inequality constraints. Diffusion models implicitly assume that the forward and time-reversed diffuison process used to perturb the data distribution are well-defined for all times. However, when the considered space upon which the diffusion evolves is a Riemannian manifold defined through inequaity constraints this is not necessarily the case. The authors tackle this problem by proposing to consider diffusion processes defined on the logarithmic barrier metric and with reflected Brownian motion induced by boundary constraints.\n\nThe first approach uses intuitions from barrier methods employed in optimisation and sampling, and considers a 'potential' that blows up at the boundary positions, which they introduce into the geometry of the considered space, that prevents the sampled process to hit the boundary. The resulting constrained diffusion is defined with respect to a metric that comes as the Hessian of the considered potential function.\n\nThe second approach introduces a reflected Brownian in the dynamics of the diffusion process that prevents the process to trespass the boundary of the domain.\n\nThe authors present numerical results of the two methods on applications which require sampling on constrained domains (like distributions symmetric positive definite matrices with trace constraints and conformational ensembles of proteins under position constraints). While the reflected method seems to work better on the experiments, and is claimed to scale better with system dimension, the evaluation of the likelihood for constrained models with the latter approach is not straightforward. \n\nOverall, the paper is well written and provides a useful extension of diffusion modeling to bounded non-Eucledian domains. The proposed methods are supported with mathematical proofs in the supplement and numerical experiments, while the topic is relevant for the audience of the TMLR.\n\n\n \n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "\n- I think the authors should also cite [1] for the generalised score matching estimator on a generic domain $V$ that uses shortest distance to the boundary, which is effectively similar to the scaling function the authors considered here.\n\n- For estimating the truncated loss the authors use a scaling function similar to the one used in [1] and a weighting function indicated as $\\lambda_t$ in the text (function of time). Can the authors comment on the purpose of including the weighting function $\\lambda_t$ that cannot be accommodated by the scaling function in Eq.(3) ? Also I think the expression of the loss should include also the scaling function to accurately reflect the implemented solutions.\n\n- In page 3, after Eq. 2, the authors state that the $\\nabla log p_t$ is unavailable in practice, but this is not true for the equation they mention above. For Ornstein-Uhlenbeck processes with additive noise the marginal density $p_t$ and thus also the Stein score have an analytically tractable solution. The process considered in [2] has time dependent diffusion and is therefore its marginal density is not analytically tractable. I think the authors probably should rewrite this sentence to refer to a more general case, or update the SDE to a multiplicative noise SDE.\n\n- There are some recent approaches on constraining diffusions that employ score matching and the Doob-h transform to constrain diffusion processes [3,4] based on linearly solvable stochastic control [5]. Would such an approach defined for stochastic processes on manifolds provide an alternative solution for constraining the diffusion processes within the domain? I imagine this would result in altering the diffusion process to include an extra drift term that would be similar to a potential that blows up at the boundary. But as I understand this would not result in a similar formulation as the logarithmic barrier approach proposed here.\n\n- Figure 9 needs a better explanation in the caption on what is depicted. Also for completeness I would consider necessary to provide in the supplement the same plot with the omitted samples.\n\n\n**Minor:**\n\n- page 3, middle: the sentence “(iv) A score matching loss, in this paper we will focus on the implicit score matching loss.” reads strange. I think it requires a full stop or a  dash after the first loss.\n\n- In Algorithm 1: for completenes explain the parameters T,N X, b that the algorithm requires as input. As far as I understand $N$ is the number of discretisation steps, and $T$ the total simulation time, b the drift of the forward process, and $X$ the initial condition (?). Also please consider to use a different letter to indicate the drift $b$ since in the previous page $b$ was a parameter used to define the inequality constraints.\n- Page 5 last paragraph: I think the standard metric with the previous notation on the unconstrained domain $\\mathcal{N}$ is denoted by $\\mathfrak{h}$ not $\\mathfrak{g}$.\n\n- The hyperlinks of the pdf files do not seem to work and given the length of the paper and the supplement it is difficult to read.\n------\n**References:**\n\n[1] Liu, S., T. Kanamori, and D. J. Williams. \"Estimating density models with truncation boundaries.\" Journal of Machine Learning Research (2022)\n\n[2] Song, Yang, et al. \"Maximum likelihood training of score-based diffusion models.\" Advances in Neural Information Processing Systems 34 (2021)\n\n[3] Zhang, Benjamin, Tuhin Sahai, and Youssef Marzouk. \"Sampling via Controlled Stochastic Dynamical Systems.\" I (Still) Can't Believe It's Not Better! NeurIPS 2021 Workshop (2021)\n\n[4] Maoutsa, D., and Opper, M. \"Deterministic particle flows for constraining stochastic nonlinear systems.\" arXiv preprint arXiv:2112.05735 (2021)\n\n[5] Kappen, Hilbert J. \"Linear theory for control of nonlinear stochastic systems.\" Physical review letters 95.20 (2005).", "broader_impact_concerns": "There are no ethical implication directly related to the proposed methods.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "286zMsjlUj", "reviewer": "Reviewer_KexQ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The paper addresses an important gap in the literature on designing diffusion models for constrained spaces. This problem is motivated by both theoretical considerations---i.e. De Bortoli et. al's approach cannot handle manifolds with constraints, and practical problems as domains such as robotics have rigid constraints that must be met.\n- Conceptually, the main ideas in the paper are very clear. Both the log-barrier method and reflected diffusion model and their associated SDEs are natural fits for designing diffusion models in constrained spaces.\n- The overall presentation of the material is clear and the writing is of high quality.\n\nWeaknesses:\nOverall the paper does not have a single major weakness but a series of minor ones that I list below.\n- First, it seems both the main ideas---i.e. the log barrier method, and the reflected diffusion model have already been known to some degree in the literature. The technical hard part of the log-barrier method (designing the SDEs) can be found in Lee and Vampala 2017. While the reflected diffusion model appeared concurrently in Lou and Ermon 2023. So one must wonder about the amount of novelty here. I do admit that the authors use the ISM loss parametrization rather than DSM one---but its unclear if this is a meaningful enough differentiator. Can the authors comment on where the main differences between their contributions lie and the ones that already exist in the literature?\n- Moreover, the reflected diffusion model in Lou and Ermon seems to have a much higher degree of execution. They scale it up to large-scale image datasets and show improved metrics over previous baselines. Since pixels within images that are represented in computers fall in the range of $[0, 255]$ and are often normalized to $[-1,1]^d$ we can think of them as a hypercube which is one the motivating examples used in this papers introduction. Consequently, I was surprised to see that there were no image experiments conducted in this paper. Also, one of the main claims in this paper is the scalability and flexibility of the proposed approach but the current experiments do not show this as the largest manifolds in the paper have $d=10$ (please correct me if I'm wrong here). \n- I also suspect the scalability and flexibility of the log-barrier method. Specifically, the log-barrier method considered in this work requires the Hessian which results in a Hessian-manifold and this is only a specific type of barrier function that has a specific Riemannian metric $g(x) = A^TS^{-2}(x)A$. Thus the current paper does not seem that flexible in this regard. Furthermore, requiring the computation of a Hessian means that such an approach will be inherently harder to scale to manifolds of higher dimensionality. In fact, it is clear that this method is not scalable as evidenced by section 5.1 which requires the inversion of a $d \\times d$ matrix which as the authors correctly remark requires $O(d^3)$ cost. Scaling this to the size of large-scale image datasets seems challenging.\n- Lastly, for the log-barrier method warping the geometry of the problem seems like an undesirable tradeoff for generative modeling. The manifold itself is a rich object of study and by changing its structure we lose part of the problem definition so it seems not as elegant as the reflected diffusion model.\n\nQuestion:\n- The motivation at the beginning mentions that De Bortoli et. al and Huang et. al's approach cannot work with constrained manifolds. While I agree this is definitely the case for De Bortoli et. al who use an intrinsic perspective to parameterize the diffusion process which can escape the boundary I'm not entirely sure this problem exists with RDM as proposed in Huang et. al. Concretely, they use an extrinsic formulation and require a tangential projection operator so the diffusion automatically moves in the tangent space of the given manifold.  Perhaps I am wrong here but this would suggest that the process would not exit the manifold? I'm happy to be corrected here through a worked example.\n\nMinor Comments:\n- The formal approach to ’warping the geometry’ of\" - Notice the first quote is formatted incorrectly in latex for this sentence found on page 4\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the design of diffusion models in constrained spaces. In particular, the diffusion models are built for Riemannian manifolds with a set of inequality constraints. Examples include polytopes and convex sets of Euclidean spaces. To facilitate a proper stochastic process in these constrained manifolds the authors propose two strategies: 1.) a logarithmic barrier where the inequality constrained is formulated using the geodesic distance to the boundary and 2.) a reflected diffusion model where the stochastic process is reflected back in the inward normal direction at the boundary. Using these two methods the paper proposes the construction of both a diffusion process and the time-reversed denoising process as well as sampling algorithms for each method. On the empirical side, the authors conduct experiments on toy synthetic domains, robotic arms under force and velocity constraints, and modeling protein loops with anchored endpoints. Finally, the experiments demonstrate that the reflected diffusion model slightly outperforms the log-barrier method.\n\n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- A clearer distinction between contributions in this paper and ones borrowed from Lee and Vampala 2017.\n- A more thorough discussion between the reflected diffusion model in this paper and the one in Lou and Ermon 2023. This I believe will be useful to the community at large.\n- Inclusion of larger scale experiments on image data and comparisons to baseline diffusion models in this case.\n- All results in Tables should have quantification of statistical uncertainty over multiple trials and not just the raw number. The current presentation is difficult to interpret because we don't know what the standard error is.\n\n", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "xeXg6hmsPY", "reviewer": "Reviewer_x71F", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The article has the following strengths. \n\n- The article is generally well written, although the topic is very technical. However one needs a solid understanding in stochastic analysis and differential geometry to grasp it. I think this is due to the nature of the topic and inevitable.\n\n-  Related work is covered well. \n\n - The plenty visualizations are very helpful for understanding the concepts and are well executed. \n\n - The article is building a rigorous theory of both approaches. All the necessities of defining a diffusion model (establishing a forward process, stationary distribution, reverse SDE) are proved in a lot of detail. One has to refer to the appendix at several points for a solid understanding. \n\n- The two experiments are interesting applications where the data is in a constrained domain and there is a need for adapted algorithms to fit these constraints.\n\nI will list the major weaknesses here: \n\n - In the concurrent work [Lou, Ermon] the ism loss approach is not taken. They argue that the ism loss will scale poorly to the high dimensional setting since almost all mass is concentrated close to the boundary (in the case of $[0,1]^d$) and this loss is downweighted due to their slightly differing derivation. How do you think this loss behaves in high dimensions? Have you tried the classical image generation task?\n\n - It is hard to interpret the experimental results, as they require some domain knowledge. While this is fine, I would like to see a standard diffusion model (in Euclidean space) trained on some of those to get a feeling of how the two approaches perform. In particular if a metric like MMD is chosen, one can even compare the result since standard diffusion should approximate the data manifold as well. \n\n- The protein loop experiment is missing quantitative evaluations. I think it is not too high dimensional in order to take MMD (or even NLL) to compare the different methods. This is important because visually the log barrier method looks closer to the ground truth to me, however this would contradict the conclusions of the preceding experiments. \n\nMinor Weaknesses/Questions: \n\n- The DSM loss is missing a minus. Furthermore for both the DSM and the ISM loss standard references should be cited in section 3 as well not only in the intro. \n\n - In section 3.3 the clipped network parameterization does not make sense. Do you mean minimum here?\n\n - In section 5.1 you write that one advantage of the log barrier method is the equivalent ODE. However in appendix F you derive this equivalent ODE in the reflected case. This seems to be a relict of an earlier version?\n\n - Proposition F.1 does not contain a mathematical statement, it is only assumptions (as far as I see). Also the proof of it needs some work. I dont understand how smoothness of the score implies that $Y_t$ is on the boundary of M. \n\n  - In Lemma G.2 you use the surface measure on $\\partial M$ without having introduced it. \n\n - For the MMD evaluations you should write out which kernels you use. Also how many samples are used to calculuate the MMD? How much do the results depend on the used kernel?\n\n- The fact that the reflected Brownian motion converges to the uniform distribution on the manifold from [Loper,2020] only seems to hold on convex sets which is more specific than your definition of M. Please clarify if this only holds on convex sets. \n\n - The typesetting in E.1 does not look too nice. \n\n - [Ablin, 2022] is not the standard reference for retractions in differential geometry.\n\n- Can you comment on the regularity of the score in the reflected case? Here, the density at time t is not given anymore by the convolution of the data density and a gaussian transition, but something more complicated (reflected Gaussian kernel)? (This is just a curiosity of me, not a weakness)\n\n- The description of Algorithm 2 is too long. I would suggest putting a more detailed description of it into the appendix. Also the abstract notions of parallel transport,... could be described in a more computational way.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces diffusion models on constrained domains. While diffusion models have already been explored in the manifold setting, the case of bounded domains remains open, as one can show that the forward process a.s. escapes the domain. Therefore, two approaches are introduced: the log barrier and the reflected diffusion model. The log barrier defines a potential which penalizes if the process gets close to the boundary. The reflected diffusion model establishes another (local) BM which stays inside the domain and admits a decomposition into a standard BM and a bounded variation process. This allows to define a reverse process.\n\nThe two approaches are underlined by rigorous theorems that prove that the forward processes admit a stationary distribution, the corresponding statements for the reverse SDEs and some experiments. The first experiment is on toy data and the last two are from robotic modelling and molecular chemistry. ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please clarify the weaknesses. The most pressing issue in my opinion is the lack of comparison with a \"standard\" diffusion model (trained with a DSM loss). ", "broader_impact_concerns": "As this is mainly a theoretical paper, I do not see any negative implications and no need for a broader impact statement.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The conclusion of the rebuttal process is that the paper is well written and contains sufficient evidence to back up their main claims. The authors toned down claims of scalability of the proposed method because of the lack of evidence that supports this claim. There are several domains of applications for diffusion models on constrained domains, and the topic is also of interest from a theory point of view. Therefore we expect this manuscript to be of interest to the audience of TMLR, and create opportunities for interesting follow-up work. Although the novelty of the proposed method is perhaps limited compared to prior/concurrent work, the manuscript nicely exposes the similarities and differences with this related work. Taking all this into consideration, the recommendation is to accept the manuscript as is.
Decision: Accept (poster)