Metropolis Sampling for Constrained Diffusion Models
Authors: Nic Fishman
Leo Klarner
Emile Mathieu
Michael John Hutchinson
Valentin De Bortoli
Venue: NeurIPS 2023 poster (2023)
Keywords: diffusion model
generative modelling
manifold
constraints
proteins
robotics
TL;DR: We present a tractable extension of Riemannian diffusion models to arbitrarily constrained manifolds.

Abstract: Denoising diffusion models have recently emerged as the predominant paradigm for generative modelling on image domains. In addition, their extension to Riemannian manifolds has facilitated a range of applications across the natural sciences. While many of these problems stand to benefit from the ability to specify arbitrary, domain-informed constraints, this setting is not covered by the existing (Riemannian) diffusion model methodology. Recent work has attempted to address this issue by constructing novel noising processes based on the reflected Brownian motion and logarithmic barrier methods. However, the associated samplers are either computationally burdensome or only apply to convex subsets of Euclidean space. In this paper, we introduce an alternative, simple noising scheme based on Metropolis sampling that affords substantial gains in computational efficiency and empirical performance compared to the earlier samplers. Of independent interest, we prove that this new process corresponds to a valid discretisation of the reflected Brownian motion. We demonstrate the scalability and flexibility of our approach on a range of problem settings with convex and non-convex constraints, including applications from geospatial modelling, robotics and protein design.

Metropolis Sampling for Constrained Diffusion Models
Nic Fishman University of Oxford
Leo Klarner University of Oxford
Emile Mathieu University Cambridge
Michael Hutchinson University of Oxford
Valentin De Bortoli
ENS Ulm
Abstract
Denoising diffusion models have recently emerged the predominant paradigm for generative modelling on image domains. In addition, their extension to Riemannian manifolds has facilitated a range of applications across the natural sciences. While many of these problems stand to benefit from the ability to specify arbitrary, domain- informed constraints, this setting not covered by the existing (Riemannian) diffusion model methodology. Recent work has attempted to address this issue by constructing novel noising processes based on the reflected Brownian motion and logarithmic barrier methods. However, the associated samplers are either computationally burdensome only apply convex subsets of Euclidean space. this paper, we introduce alternative, simple noising scheme based on Metropolis sampling that affords substantial gains in computational efficiency and empirical performance compared to the earlier samplers. Of independent interest, we prove that this new process corresponds to a valid discretisation of the reflected Brownian motion. We demonstrate the scalability and flexibility of our approach on a range problem settings with convex and non-convex constraints, including applications from geospatial modelling, robotics and protein design.
Introduction
In recent years, denoising diffusion models (Sohl-Dickstein al., 2015; Song et al., 2019; Song et al., 2021; Ho et al., 2020) have emerged as a powerful paradigm for generative modelling, achieving state-of-the-art performance across a range of domains. They work by progressively adding noise to data following a Stochastic Differential Equation the forward noising process-unti it is close to the invariant distribution of the SDE. The generative model is then given by an approximation of the associated time-reversed denoising process, which is also SDE whose drift depends on the gradient of the logarithmic densities of the forward process, referred to as the Stein score, Building on the success of diffusion models for image generation tasks. De Bortoli et al. (2022) and Huang et al. (2022) have recently extended this framework to a wide range of Riemannian manifolds, enabling the specification of inherent structural properties of the modelled domain. This has broadened the applicability of diffusion models to problems in the natural and engineering sciences, including the conformational modelling of small molecules (Jing et al., 2022; Corso et al., 2022), proteins (Trippe et al., 2022; Watson et al., 2022; Yim et al., 2023) and robotic platforms (Urain et al., 2022).
However, in many data-scarce safety-critical settings, researchers may want to restrict the modelled domain even further by specifying moblem--nformen constraints to make maximal use of limited experimental data prevent unwanted behaviour (Morris, 2002; Han et al., 2006; Thiele et al., 2013; Lukens et al., 2020). As illustrated in Figure 1, such domain- inniiin-ffrmeed constraints can be naturally represented as Riemannian manifold with boundary. Training diffusion models on such constrained manifolds is thus an important problem that requires principled noising processes and corresponding discretisationss that stay within the constrained set.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

(a) R
(b) S C TR
(c) M C 8 CR.
Figure 1: When operating in data-scarce settings, it may often be beneficial to specify as much prior knowledge of the modelled domain as possible. Consider a distribution over a subset M of the unit sphere 5 C 3R3 While Euclidean diffusion model can approximate the distribution in R (a), directly modelling it on S can make learning significantly easier (b). Restricting the problem space even further by only constructing a diffusion model on . can lead to even better performance (c).
Recent work by Fishman et al. (2023) has attempted to derive such processes and extend the applicability of diffusion models to modelling applications of classic sampling schemes based on log-barrier methods (Kannan et al., 2009; Lee et al., 2017; Noble et al., 2022; Kook et al., 2022; Gatmiry et al., 2022; Lee et al., 2018) and the reflected Brownian motion (Williams, 1987; Petit, 1997; Shkolnikov al., 2013). While empirically promising, the proposed algorithms can be computationally and numerically burdensome, and require bespoke implementations for different manifolds and constraints. Concurrently, Lou et al. (2023) have investigated the use of reflected diffusion models for image applications. They focus on the high-dimensional hypercube, as this setting admits theoretically grounded treatment of the static thresholding method which is widely used in image models such as Saharia et al. (2022). More recently, Liu et al. (2023) have investigated the use of log-barriet beggbrrrrr-baed based mirror maps transform constrained domain into an unconstrained dual space for applications in image watermarking. While both methods exhibit robust scaling properties and impressive results, they only consider convex subsets of Euclidean space and do not extend to more general manifolds.
manifolds by investigating the generative
Here, we propose a new method for generative modelling on constrained manifolds based on Metropolis-based discretisation of the reflected Brownian motion. The Metropolised process' chief advantage is that it is lightweight: the only additional requirement over those outlined in De Bortoli et al. (2022) that is needed to implement constrained diffusion model is an efficient binary function that indicates whether any given point is within the constrained set. This Metropolised approximation of the reflected Brownian motion is substantially easier to implement, faster to compute and more numerically stable than the previously considered sampling schemes. Our core theoretical contribution is to show that this new discretisation converges to the reflected SDE by using the invariance principle for SDEs with boundary (Stroock et al., 1971). To the best of our knowledge, this is the first time that such a process has been investigated. We demonstrate that our method attains improved empirical results on diverse manifolds with convex and non-convex constraints by applying it to a range of problems from geospatial modelling, robotics and protein design.
2 Background
Riemannian manifolds. A Riemannian manifold is defined as a tuple (M, g) with M a smooth manifold and a metric defining inner product on tangent spaces. In this work, will use the exponential map exp, T + M, as well as the extension of the gradient V, divergence div and Laplace A operators to M. All of these quantities can be defined in local coordinates in terms Beltrami operator, M,
of the metric. The extension of the Laplace operator to M is called the Laplace-Beltrami
also denoted A when there is ambiguity. Using A, we can define a Brownian motion denoted (B,) and with density w.r.t. the volume form of M denoted pt for any t > O. We refer to Appendix B for a more detailed exposition, to Lee (2013) for a thorough treatment of Riemannian manifolds and to Hsu (2002) for details on stochastic analysis on manifolds. In the following, we consider a constrained manifold M defined by
M {I N f((z) < 0, i I),
(I)

---

Unconstrained
Reflected
Metzopolls
Time
Figure 2: Visual comparison of discretisation of the unconstrained Brownian motion (blue) and two discretisations of the reflected Brownian motion: one based on reflection scheme (green) and the other based our Metropolis sampler (red). The Metropolised trajectory is very close to that of the reflected one while being significantly easier to implement and cheaper to compute.
where (N, g) is Riemannian manifold, I is arbitrary finite indexing family and for any E I, fi E C(N, R). Since is finite and fi continuous for any E I, .M is an open set of N and inherits its metric g. This captures simple Euclidean polytopes and complex constrained spaces like Figure 1.
Denoising diffusion models. Denoising diffusion models (Song et al., 2019; Ho et al., 2020; Song al., 2021) work as follows: let (Xe)re(o.T) be noising process that corrupts the original data distribution Po- We assume that (X))r20 converges to N(0, o² Id), with 0. Several such processes exist, but in practice we consider the Omstein-Uhlenbeck (OU) process, also referred to as VP-SDE, which is defined by the following Stochastic Differential Equation (SDE)
dx,- - + adB,,
X ~ po-
(2)
Under conditions on po. for any T > 0. (E.) = (XT- is also the (weak) solution to a SDE (Anderson, 1982; Haussmann et al., 1986; Cattiaux et al., 2021)
dY: - ( Y, V pr (Y.))dd + odB,, Yo ~ pr:
(3)
where Pt denotes the density of X,, In practice, V log Pt is approximated with score network (t,z) x)) + st(t, æ)) trained by minimising either a denoising score matching (dsm) loss or an implicit score matching (ism) loss (Vincent, 2011)
{))- (x x.),e..))AA(( (l, X.) (1,
where A 0. For flexible score network, the global minimiser 8* argmin, satisfies 88. (t. = V logg: pt. De Bortoli et al. (2022) and Huang et al. (2022) have extended denoising diffusion models the Riemannian setting. The time-reversal formula (3) remains the same, replacing the Euclidean gradient with its Riemannian equivalent. The ism loss can still be computed in that setting. However, the samplers used in the Riemannian setting differ from the classical Euler- Maruyama discretisation used in the Euclidean framework. De Bortoli et al. (2022) use Geodesic Random Walks (Jorgensen, 1975), which ensure that the samples remain on the manifold at every step. In this paper, we propose sampler with similar properties in the case of constrained manifolds.
Reflected SDE. We conclude this section by recalling the framework for studying reflected SDEs, which is introduced via the notion of the Skorokhod problem. For simplicity, we focus on Euclidean space Rd here, but note that reflected processes can be defined on arbitrary smooth manifolds N. In the case of the Brownian motion, a solution to the Skorokhod problem is process of the form (B,, k)) Locally, (B,)ı20 can be seen as a regular Brownian motion (B,)220 while (k))20 forces (B,) to remain in M. Under mild additional regularity conditions on M and (B,, k, k)) Skorokhod (1961), (B,, k2) is solution to the Skorokhod problem if for any t 0
B. Bo B: K; M,
(5)
[k|, - S/ da k), and k, f n (B.)d where (k)))) is the total variation of (kr (k)) Let us provide some intuition on this definition. When (Bt) hits the boundary i,,, -k, pushes the process back into M along the inward normal n (B.), according to k, = J n (B.)d[k)-
In this case (kz) not regular enough, but if were the class C², its total variation would be given by 56 (r |ds in the one- me--imensioonl case.
3

---

(a) Reflected Brownian motion
(b) Metropolised approximation of Brownian motion
Figure 3: Evolution of the density of the reflected Brownian motion and its Metropolis sampling- based approximation on the unit interval starting from delta mass.
The condition Ikl; = L B.ea.d is more technical and can be seen as imposing that k, remains constant so long as (Be) does not hit a.M. We refer to Fishman et al. (2023) and Lou et al. (2023) for a more thorough introduction of these concepts in the context of diffusion models.
3 Diffusion models for constrained manifolds via Metropolis sampling
In Section 3.1, highlight the practical limitations of existing constrained diffusion models and propose an alternative Metropolis sampling- approach. In Section 3.2, we outline our proof that this process corresponds to valid discretisation of the reflected Brownian motion, justifying its use in diffusion models. An overview of the samplers we cover in this section is presented in Table 1.
3.1 Practical limitations of existing constrained diffusion models
Barrier methods. In the barrier approach, a constrained manifold is transformed into an uncon-
strained space via a barrier metric. This metric is defined by V²6(z) with ((x) >iez & (d(xr, f() where d(x, f.) is the minimum distance from the point x to the set defined by f((x) = and Φ; is monotonically decreasing function such that lim,-o (2)
Under additional regularity assumptions, & is called a barrier function (see Nesterov et al. (1994)). This definition ensures that the barrier function induces well-defined exponential map on the manifold, making the Riemannian diffusion model frameworks of De Bortoli et al. (2022) and Huang et al. (2022) applicable. In the log-barrier method of Fishman al. (2023), evaluating é requires computing d(x, 8.)) (and its derivatives), which can be prohibitively expensive. Furthermore, since the exponential map under the induced manifold difficult compute, it is approximated by projecting the exponential map on the original manifold back onto the constraint set, incurring an additional bias. Liu et al. (2023) propose more tractable method by constructing a mirror map that transforms a constrained domain into an unconstrained dual space. in which one can train a standard Euclidean diffusion model. However, this approach is only applicable convex subsets of Rd and does not extend arbitrary Riemannian manifolds. More generally, warping the geometry of the modelled domain can adversely impact the interpolative performance of log-barrier-batt based diffusion models, as the space between data points expands rapidly when approaching the boundary.
Reflected stochastic processes. Fishman et al. (2023) and Lou et al. (2023) introduce diffusion models based on the reflected Brownian motion (RBM). In Fishman et al. (2023), the reflected SDE is discretised by (i) considering a classical step of the Buler-Maruyama discretization (or the Geodesic Random Walk in the Riemannian setting) and (ii) reflecting this step according to the boundary defined by a.M. To compute the reflection, one must check whether the step crosses the boundary. If it does, the point of intersection needs to be calculated in order to reflect the ray and continue the step in the reflected direction. This can require an arbitrarily large number of reflections depending on the step size. the geodesic on the manifold, and the geometry of the bounded region within the manifold. We refer to Appendix C for the pseudocode of the reflection step and additional comments. An alternative approach to discretising reflected SDE is to replace the reflection with a projection (Stomiński, 1994). However, the projection requires the most expensive part of the reflection algorithm: computing the intersection of the geodesic with the boundary. Lou et al. (2023) propose a more tractable approach that exploits the product structure of the unit hypercube to afford simulation-fris sampling but does not extend to arbitrary Riemannian manifolds. Additionally, specifying convex constraints their framework requires a bijection onto the hypercube, distorting the modelled geometry and incurring the same issues as outlined above.

---

Table Comparison of the advantages and disadvantages of the different constrained (Riemannian) diffusion models covered in Section 3.1.
Both required for fast DSM loss
TRACTABLE
SIMULATION- FORWARD SAMPLING
MODELLING DOMAIN
DIFFUSION MODEL
CONDITIONAL SCORE
METRIC OF M
Reflected Diffusions
LOU ET AL. (2023)
COnveX C R
X
FISHMAN ET AL. (2023)
Old))
Any M Any M
V
METROPOLIS (OURS)
((d)
Barrier Diffusions
FISHMAN ET AL. (2023)
X
COnveX C any M
LIU ET AL. (2023)
/ Convex C R"
X
Metropolis approximation. Existing approaches to constrained (Riemannian) diffusion mod- els either only apply to convex subsets of IRd or require manifold- and constraint-specife implementations that become computationally intractable as the complexity of the modelled geometry increases. This limits their practicality even for relatively simple manifolds with well-defined exponential maps and linear inequality constraints such as for example polytopes. In the following, we introduce method that aims to solve both of these problems. The sampler we pro- classical
Algorithm Metropolis approx. of RBM
pose is similar sation of the Brownian motion, except that, whenever step would carry the Brownian motion of the con- strained region, we reject (see Algorithm 1). This is Metropolised version of the usual discretisation and is trivial to implement compared to the existing barrier, reflection and projection methods. Hence, this method enables the principled extension of diffusion models to arbitrarily constrained manifolds virtually no added implementational complexity or computational expense.
discreti-
Require: p E M, ffihez
Sample p' f- exp,(o)
N(0, Id) TpM
if f((2) < <0Vi then
p Dt-p' end if return p
3.2 Relating the Metropolis sampler to the reflected Brownian motion
In this section, prove that the proposed Metropolis sampling- baasling-based process (Algorithm 1) corre- sponds to valid discretisation the reflected process, justifying its use in diffusion models. Here focus concise presentation of the concepts and the main results. A full proof can be found in Appendix D. For simplicity, we restrict ourselves to the Euclidean setting. All of our results require particular assumptions M,, which we discuss at the end of this section. We begin with definition of the Metropolis approximation of RBM.
Definition 1. For any 7i > 0 and l E N, let Xj E M and X = +X V M and X otherwise. The sequence (XE) is called the Metropolis approximation of RBM.
For any >> >> O, we consider (X)) the linear interpolation of (Xī) )EEN such that for any k E N. X X7. The following result is the main theoretical contribution of our paper.
Theorem 2. Under assumptions on M. for any T (X)))etory weakly converges the RBM (Be)cejorr as O.
The rest of the section is devoted to a high level presentation of the proof of Theorem 2. It is theoretically impractical work directly with the Metropolis approximation of RBM. Instead, introduce an auxiliary process, show this converges to the RBM, and finally prove that the convergence of the auxiliary process implies the convergence of our Metropolis discretisation.
M and X X2 + vRLi with Z3 Gaussian random variable conditioned on X + RZ3 E M. The sequence (XP) is called the Rejection approximation of RBM.
Definition For any 7 and k E N, let X3
5

---

We call this process Rejection approximation of RBM since in practice, 27 sampled using rejection sam- pling, see Algorithm 2. For any 7 > O, we also consider (X))(20. the linear interpolation of (XE))EN such that for any N, X2, = XX. In Appendix D, we prove the following result.
Algorithm 2 Rejection apprex. of RBM
Require: p E M, (ffhee
Sample v ~ (0, E TTM
p' f exp,(v)
while f:(p') 2 0 for any i do Sample e ~ [d, 1) E T.M
Theorem 4. Under assumptions on M, for any T O. IE(o.T weakly converges to the Reflected Brownian Motion (Be)se(o.T] as 7 - O.
p' f- (xp,(r)
end while return p'
Proof. Here we give some elements of the proof. Details and full derivations are postponed to Appendix D. Our approach is based on the invariance principle of Stroock et al. (1971). More precisely, we show that we can compute an equivalent 'drift' and 'diffusion matrix' for the discretised process and that, the Euclidean setting, this result, accompanied by mild regularity and growth assumptions, ensures that the discretization weakly converges to the original SDE. However, the case with boundary is much more complicated, primarily because the approximate drift might explode near the boundary, thus we need to verify exactly how the drift behaves as 7 0 and as the process approaches the boundary. We show that the normalised drift converges to the inward normal near the boundary. This ensures that (a) in the interior of M the drift converges to zero, i.e. locally in the interior of M the Brownian motion and the Reflected Brownian Motion coincide, (b) on the boundary, the drift pushes the samples inside the manifold according the inward normal, mimicking (k, (2)) in (5). Finally, with results from Stroock et al. (1971) and Kang et al. (2017), we show the convergence to the RBM.
-> O, the drift converges to zero and the diffusion matrix converges to Id. In
Our next step is to show that the approximate drift and diffusion matrix of the Metropolised process are upper and lower bounded by their counterparts in the rejection process. While the upper-bound is easy to derive, the lower-bound requires the following result.
Proposition 5. Under assumptions on M, V E M, E (0,5) and Z ~ N(0, Id) we have P(r + VJZ E M) 2 1/2 E, with Z ~ N(0, Id).
O, 35 such that for any 7 E (0,5) and for any
Proposition tells us that locally the boundary looks like a half-space when integrating w.r.t. a Gaussian measure. A corollary is that, for > 0 small enough and for any & E N, the probability that X = X7 is upper bounded uniformly w.r.t. X E .M. The proof of Proposition 5 uses Theorem 7 in Appendix D, whose proof relies on the concept of tubular neighborhoods (Lee et al., 2012).
Having established the lower and upper bound, we can conclude the proof by noting that the approximate drift and the diffusion matrix in the rejection and Metropolis case coincide 7 - 0. This is enough to apply the same results as before, giving the desired convergence.
Assumptions on M. Before concluding this section, we detail the assumptions we make on M. For Theorem 2 to hold, we assume that M {z E gd 4(x) > 0} is bounded, with & E C² R) concave. We have that aM = fx E Rd ((x) = 0}. In addition, we assume that for any E aM, ||V((=)]|
These assumptions match those Stroock et al. (1971) use for their study of the existence of solutions to the RBM. While it seems possible to relax the global existence of & to local one, the regularity assumption of the domain is key. This regularity is essential to establish Proposition 5 and the associated geometrical result on tubular neighhourhoods. We also emphasize that the smoothness of the domain is central in the results of Kang et al. (2017) on the equivalence of two definitions of RBMs which we rely on.
Related work on approximations of reflected SDEs
Several schemes have been introduced to approximately sample from reflected Stochastic Differential Equations. They can be interpreted as modifications of classical Fuler-Maruyama schemes used to discretise SDEs without boundary. One of the most common approaches to use the Fuler- Maruyama discretisation and project the solution onto the boundary if it escapes from the domain M.

---

Table Log-likelihood (1) of a held-out set from synthetic bimodal distribution over convex subsets of IR bounded by the hypercube [-1, and unit simplex A". Mcans and standard deviations are computed over 3 different runs. Average training time provided in hours.
REFLECTED
METROPOLIS
MANIFOLD
DIMENSION
log- lggliikelhood
runtime
log-likelihood
runtime
2
2.25+.01
8.95
2.32#.05
0.72
1
3
3.77
8.97
4.15#.15
0.71
10
7.421.77
10.1
10.80
0.90
2
1.01#.01
9.17
1.06#.02
0.82
A
3
2.641.0
9.43
3.28
0.78
10
7.001.13
10.5
7.811.20
0.97
Reflected Hypercube
10
Reflected Simplex
Rejection Hypercube
10
Rejection Simplex O(d)
Convergence time rel. to
O(d")
10a
10'
10
*
Dimension of manifold
Figure 4: Convergence time of the Reflected (green) and Metropolis (red) forward noising processes to the uniform distribution on the hypercube |-1, 1] and unit simplex A The lines indicate functions fit with the PYSR symbolic regression package (Cranmer, 2023) and correspond to empirical runtime complexities of (dd²) and (dd, respectively, matching the superimposed scaling law isocontours.
In this case, mean-square error rates of order almost 1/2 have been proven under various conditions (Liu, 1995; Chitashvili et al., 1981; Pettersson, 1995; Slomiński, 1994). Concretely this means that || B, - |2 O(n 0 with e > O arbitrary small and where (XT) is the projection the Euler-Peano method
scheme. The rate 1/2 is tight (Pacchiarotti al., 1998). It is possible to get slight improvements for a meennssqarr еттоr rate of order of 1/2, but this is impractical as assumes that one solve a (simplified) Skorokhod problem, which is usually intractable. Liu (1993) introduced a penalised method which pushes the solution away from the boundary and shows a mean-square error of order 1/4, see also Pettersson (1997). Weak errors of order have been obtained in Bossy et al. (2004) and Gobet (2001) by introducing a reflection component in the discretisation or using some local approximation of the domain half-space. We refer Pilipenko (2014) for an introduction to the discretisation of reflected SDEs. Closer to our work, Burdzy et al. (2008) consider three different methods to approximate reflected Brownian motions on general domains (two based on discrete methods and one based on killed diffusions). Only qualitative results are provided. To the best our knowledge, no previous work in the probability literature has investigated the Metropolised scheme we propose. Our Metropolis scheme is also related to the ball walk (Applegate al., 1991), which replaces the Gaussian random variable with a uniform over the ball (or the Dikin ellipsoid). Applegate et al. (1991) and Lovász et al. (2007) have studied the asymptotic convergence rate of the ball walk, but, to the hest of our knowledge, its limiting behaviour when the step size goes to zero has not been investigated.
7

---

(a) Data distribution.
(b) Metropolis samples.
Reflected samples.
(d) Uniform distribution.
Figure 5: A qualitative visual comparison of 10 samples from the data distribution, Metropolis diffusion model, a reflected diffusion model and the uniform distribution for the constrained configu- rational modelling of robotic arms on S R
5 Experimental results
To demonstrate the practical utility and empirical performance of the proposed Metropolis diffusion models, we conduct a comprehensive evaluation on a range of synthetic and real-world tasks. In Section 5.1, we assess the scalability of our method by applying it to synthetic distributions on hypercubes and simplices of increasing dimensionality. In Section 5.2, we extend the evaluation to real-world tasks manifolds with convex constraints by applying our method to the robotics and protein design datasets presented in Fishman et al. (2023). In Section 5.3, we additionally demonstrate that method extends to constrained manifolds with highly non-convex boundaries setting that intractable with existing approaches.
As we found-in in line with Fishman et al. (2023)-thee log- bog-barrier diffusion models perform strictly worse than reflected approaches across all experimental settings, we focus on a more detailed comparison with the latter here and postpone additional empirical results to Appendix F.1. These include additional performance metrics and a comparison to an unconstrained Euclidean diffusion model on the synthetic datasets presented in Section 5.1.
For all experiments, we use simple 6-layer MLP with sine activations and a score rescaling function to ensure that the score reaches zero at the boundary, scaling linearly into the interior of the constrained set as in Liu al. (2022) and Fishman et (2023). We 7 1, .00 ensure that the forward process reaches the invariant distribution with a linear 3-schedule. We use learning rate of 2 x 10 10-4 with cosine learning rate schedule and an ism loss with modified loss weighting function of (1 + t), batch size of 1024 and 8 repeats per batch. All models were trained on a single NVIDIA GeForce GTX 1080 GPU. Additional details are provided in Appendix F.2.
10 10-3 -3 and tune 31 to
All source code that is needed to reproduce the results presented below is made available under which requires supporting package handle the different geometries that is available under
https//github..co/occ//////////
5.1 Synthetic distributions on simple polytopes
In this section, we investigate the scalability of the proposed Metropolis diffusion models by applying them to synthetic bimodal distributions over the d-dimensional hypereube [-1, 1]d and unit simplex A A quantitative comparison of the log-likelihood of a held-out test set is presented in Table 2, while a visual comparison is postponed to Appendix F.3. We find that our Mctropolis models outperform reflected approaches across all dimensions and constraint geometries by a substantial and statistically significant margin while training in one tenth of the time. The degree of improvement seems to scale with the dimensionality of the problem: the larger the dimension of the experiment, the larger the gain in performance from using our proposed Metropolis scheme.
We observe a similar difference in the scaling properties of reflected and Metropolis models when measuring the convergence times of the respective forward noising processes to the uniform distribu- tion on hypercubes |-1, 1] and simplices A of increasing dimensionality. The results are presented in Section and show that the convergence time of the Metropolis process scales linearly in the dimension, while the reflected process scales quadratically.
8

---

(a) Data distribution.
(b) Metropolis samples.
Reflected samples.
(d) Uniform distribution.
Figure 6: A qualitative comparison of 10 samples from the data distribution, our Metropolis diffusion model, a reflected diffusion model and the uniform distribution for the constrained conformational modelling of eyclic peptide backbones. For visual clarity, the figures only show the constrained planar projections encoded by P C RR.
5.2 Modelling proteins and robotic arms under convex constraints
In addition to illustrating our method's scalability on high-dimensional synthetic tasks, we follow the experimental setup from Fishman al. (2023) additionally demonstrate its practical utility and favourable empirical performance on two real-world problems from robotics and protein design.
Constrained configurational modelling of robotic arms. The problem of modelling the configu- rations and trajectories of a robotic arm can be formulated as learning distribution over the locations and manipulability ellipsoids of its joints, parameterised on E2 x S where S + is the manifold of symmetric positive-definite (SPD) d x matrices (Yoshikawa, 1985; Jaquier et al., 2021). For practical robotics applications, it may be desirable to restrict the maximal velocity with which robotic arm can move or the maximum force can exert. This manifests in a trace constraint C' > O on St resulting in a constrained manifold {M E St Ei Mi < C). Following Fishman et al. (2023), we parametrise this constraint via the Cholesky decomposition (Lin, 2019) and use the resulting setup model the dataset presented in Jaquier al. (2021).
Conformational modelling protein backbones. Modelling the conformational ensembles of proteins is data-scarce problem with range of important applications in biotechnology and drug discovery (Lane, 2023). many practical settings, may often be unnecessary to model the structural ensembles of entire protein, researchers are primarily interested in specific functional sites that are embedded in a structurally conserved scaffold (Huang et al., 2016). Modelling the conformational ensembles of such substructural elements requires positional constraints their endpoints to ensure that they can be accommodated by the remaining protein. Using the parametrisation and dataset presented in Fishman et al. (2023), we formulate the problem of modelling cyclic peptide of length N = as learning a distribution over the product of polytope P C 3 and the hypertorus T
the backbone conformations
Table 3: Log- lig-liield (f) of a held-out test set for the robotics and protein applications. Means and standard deviations are computed over different runs. Average training time is provided in hours.
REFLECTED
METROPOLIS
DATASET
DOMAIN
log-likelihood
runtime
log-likelihood
runtime
Robotics
S X R
8.39
9.52
9.13=.03
1.36
Proteins
PCR R3 x T
15.20+.06
24.80
15.33+.02
3.12
We quantify the empirical performance of different methods by evaluating the log- likelihood of a held-out test set and present the resulting performance metrics in Table 3. Again, we find that our Metropolis model outperforms the reflected approach by a statistically significant margin while training 7-8 times as fast. Qualitative visual comparisons of samples from the true distribution, the trained diffusion models and the uniform distribution are presented in Figures 5 and 6, with full univariate marginal and pairwise bivariate correlation plots postponed to Appendices F.4 and F.5.

---

(a) Data distribution.
(b) Metropolis samples.
(c) Uniform distribution.
Figure 7: Orthographic projection of 10 samples from (a) the data distribution, (b) our Metropolis diffusion model, and (c) the uniform distribution, for geospatial data (wildfire incidence rates) within non-convex boundary (the continental United States). The projections are aligned with the geometric centre of the boundary and zoomed in ten-fold for visual clarity.
5.3 Modelling geospatial data within non-convex country borders
Motivated by the strong empirical performance of our approach tasks with challenging convex constraints, we investigated its ability to model distributions whose support is restricted to manifolds with highly non-convex boundaries- setting that intractable with existing approaches. To this end, we derived a geospatial dataset based on wildfire incidence rates within the continental United States (see Appendix for full details) and trained a Metropolis diffusion model constrained by the corresponding country borders on the sphere S A qualitative visual comparison of samples from the true distribution, our model, and the uniform distribution is presented in Figures 7a to 7c and a quantitative comparison to a Riemannian diffusion model on S (De Bortoli et al., 2022) is given in Table Both demonstrate that our approach able to successfully capture challenging multimodal and sparse distributions on constrained manifolds with highly non-convex boundaries.
Table 4: MMD (L) of a held-out test set for the geospatial modelling dataset. Means and standard deviations are computed over 3 different runs. Average training time is provided in hours.
MODEL
DOMAIN S2
MMD
RUNTIME
% IN BOUNDARY
Unconstrained
0. 10.001
0.81
63.3
Metropolis
P PCS²
0.138810.015
3.86
100.0
6 Conclusion
Accurately modelling distributions on constrained Riemannian manifolds a challenging problem with range of impactful practical applications. In this work, we have proposed mathematically principled and computationally tractable extension of existing diffusion model methodology to this setting. Based on Metropolisation of random walks in Euclidean spaces and Riemannian manifolds, we have shown that our approach corresponds to a valid discretisation of the reflected Brownian motion, justifying its use in diffusion models. To demonstrate the practical utility of our method, we have performed an extensive empirical evaluation, showing that it outperforms existing constrained diffusion models on a range of synthetic and real-world tasks defined on manifolds with convex boundaries, including applications from robotics and protein design. Leveraging the flexibility and simplicity of method, have also demonstrated that it extends beyond convex constraints and is able to successfully model distributions on manifolds with highly non-convex boundaries. While we found our method to perform well across the synthetic and real-world applications we considered, we expect it to perform poorly on certain constraint geometries. For instance, the current implementation relies on an isotropic noise distribution which could impede its performance on exceedingly narrow constraint geometries, even with correspondingly small step sizes. In this context, an important direction of future research would be to investigate whether we can instead sample from more suitable distributions, e.g. Dikin ellipsoid, while maintaining the simplicity and efficiency of the Metropolis approach.
10

---

Acknowledgements
NF thanks the Rhodes Trust for supporting their studies at Oxford and this work. LK acknowledges support from the University of Oxford's Clarendon Fund.
References
B. D. Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):11-332, 1982. Cited page 3.
D. Applegate and R. Kannan. Sampling and integration of near log-concave functions. In Proceedings of the twenty-third annual ACM symposium on Theory of computing, pages 156-163, 1991. Cited page 7.
J. A. Bednar, J. Crail, J. Crist-Harif, P. Rudiger, G. Brener, I. Thomas, C. J. Mease, J. Signell, M. Liquet, J.-L. Stevens, Collins, A. Thorve, thuydotm, H. Hansen, esc, kbowen, Abdennur, O. Smirnov, maihde, A. Hawley, A. Oriekhov, A. Ahmadia, B. A. B. Jr, C. H. Brandt, C. Tolboom, E. G., E. Welch, Bourbeau, and J. Schmidt. Holoviz/datashader, version 0. 14.4, version v0. v0.14.4, Feb. 2023. DOI: 10 zenodoo 7599872. URL: https: /oi. org/ 10. 5281 2081/zenodo 7599872. Cited on page 16.
M. Bevis and J.-L. Chatelain. Locating a point on spherical surface relative to a spherical polygon of arbitrary shape. Mathematical geology, 21:81 31:811-828, 1989. Cited on pages 15, 16.
M. Bossy, E. Gobet, and D. Talay. A symmetrized euler scheme for an efficient approximation of reflected diffusions. Journal of applied probability, (1(3)8777-88, 2004. Cited on page 7.
K. Burdzy and Z.-Q. Chen. Discrete approximations to reflected brownian motion, 2008. Cited on page 7.
P. Cattiaux, G. Conforti, I. Gentil, and C. Léonard. Time reversal of diffusion processes under finite entropy condition. arXiv preprint arXiv:2104.0770
2021. Cited on page 3.
R. Chitashvili and N. Lazrieva. Strong solutions of stochastic differential equations with boundary conditions. Stochasties: an international journal of probability and stochastic processes, 5(4):255
309, 1981. Cited on page 7.
G. Corso, H. Stärk, B. Jing, Barzilay, and Jaakkola. DiffDDok: Diffusson Steps, Twists, and Turns for Molecular Docking, Oct. 2022. URL: http://aarxiv.orr//bb/22111 01776. Cited on page 1.
M. Cranmer. Interpretable machine learning for science with pysr and symbolicregression,, jl. arXiv preprint arXiv:2305.01 5822, 2023. Cited on page 7.
V. De Bortoli, E. Mathieu, M. Hutchinson, J. Thornton, Y. W. Teh, and A. Doucet. Riemannian score-based generative modeling, 2022. arXiv: 2202.02763 [G]. Cited pages 1-4, 10, 18.
N. Fishman, L. Klamnrr, V. De Bortoli, E. Mathieu, and M. Hutchinson. Diffusion models for constrained domains. arXiv preprint arXivv2304.05364 2023. Cited on pages 2, 4, 5, 8, 9, 16-18, 20.
K. Gatmiry and S. S. Vempala. Convergence of the riemannian langevin algorithm. arXiv preprint arXiv:2204. 10818, 2022. Cited on page 2.
E. Gobet. Euler schemes and half-space approximation for the simulation of diffusion in a domain. ESAIM: Probability and Statistics, 5:261-297 2001. Cited page 7.
A. Gretton, K. M. Borgwardt, M. J. Rasch, B. Schölkopf, and A. Smola. A kernel two-sample test. Journal of Machine Learning Research, 13(mull):723-77 -73,,73, Mar. 2012. ISSN: 1532-4435. Cited on page 16.
L. Han and L. Rudolph. Inverse kinematics for serial chain with joints under distance constraints.
In Robotics: Science and systems, 2006. Cited on pages 1, 20.

---

U. G. Haussmann and E. Pardoux. Time reversal of diffusions, The Annals of Probability: 188- 1986. Cited on page
J. Ho, A. Jain, and Abbeel. Denoising diffusion probabilistic models, Advances in Neural Informa- tion Processing Systems, 2020. Cited on pages 1, 3.
E. P. Hsu. Stochastic analysis on manifolds, number 38. American Mathematical Soc., 2002. Cited on page 2.
C.-W. Huang, M. Aghajohari, A. Bose, P. Panangaden, and A. Courville. Riemannian Diffusion Models, Aug. 2022. DOT: 10. 48550/ 48550/arXiv. arxiv. .7949. Cited on pages 1, 3,4.
P.-S. Huang, S. E. Boyken, and D. Baker. The coming of age de novo protein design. Nature, 337076200:320-327, 2016. Cited on page 9.
N. Jaquier, I., Rozo, D. G. Caldwell, and S. Calinon. Geometry-aware manipulability learning, tracking, and transfer. The International Journal of Robotics Research, 40(2-3)::24-6550 2021. Cited on page
B. Jing, G. Corso, J. Chang, R. Barzilay, and T. Jaakkola. Torsional Diffusion for Molecular Con- former Generation, June 2022. URL: http http//aawww..rrraa///aaa//////// 01729. Cited on page 1.
E. Jørgensen. The central limit problem for geodesic random walks. Zeitschrift für Wahrscheinlichkeii- stheorie und verwandte Gebiete, 32(1-2):1-644 1975. Cited page 3.
W. Kang and K. Ramanan. On the submartingale problem for reflected diffusions in domains with piecewise smooth boundaries, 2017. Cited on pages 6, 4, 14.
R. Kannan and H. Narayanan. Random walks on polytopes and an affine interior point method for linear programming. In Proceedings of the Forty-First Annual ACM Symposiun Theory of Computing. STOC '09, pages 561-570, Bethesda, MD, USA. Association for Computing Machinery, 2009. ISBN: 978 100585006.. DO1: 10 1145 1536414 1536491. //doi. oro/0...rg11... 1145/ 11344411 1536491. Cited on page 2.
URL: https:
R. Ketzner, Ravindra, and M. Bramble. A robust, fast, and accurate algorithm for point in spherical polygon classification with applications in geoscience and remote sensing. Computers & Geosciences, 167:105185, 2022. Cited on pages 15, 16.
Y. Kook, Y. T. Lee, R. Shen, and S. Vempala. Sampling with riemannian hamiltonian monte carlo in constrained space. arXiv preprint arXiv:2202.019000 2022. Cited on page 2.
T. J. Lane. Protein structure prediction has reached the single-structure frontier. Nature Methods: 1 1-4, 2023. Cited on page
B. Laurent and Massart. Adaptive estimation of quadratic functional by model selection. Annals of Statistics: 1. 1302-1338, 2000. Cited on page 6.
J. M. Lee. Smooth manifolds. In Introduction to Smooth Manifolds, pages 1-31. Springer, 2013. Cited on pages 1.
M. Lee and J. M. Lee. Smooth manifolds. Springer, 2012. Cited on pages 6, 5.
Y. T. Lee and S. 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, pages 1 1115-1121, 2018. Cited on page 2.
Y. T. Lee and S. Vempala. Geodesic walks in polytopes. In Proceedings of the 49th Annual ACM SIGACT Symposium on Theory of Computing, pages 927-940, Montreal Canada. ACM, June 2017. ISBN; 978-1-4503-4522-6. DOI: 10.1140/3331 00.1145/3055399 3055416. Cited on page 2.
Z. Lin. Riemannian geometry of symmetric positive definite matrices via cholesky decomposition. SIAM Journal on Matrix Analysis and Applications, 40(4):1353-137) 2019. Cited on pages 9, 2.
G.-H. Liu, T. Chen, E. A. Theodorou, and M. Tao. Mirror diffusion models for constrained and watermarked generation, arXiv preprint 2rXiv:2310.0/236 2023. Cited on pages 2, 4, 5.
12

---

S. Liu, T. Kanamori, and D. J. Williams. Estimating density models with truncation boundaries using score matching. Journal of Machine Learning Research, 23(186)::-38, 2022. Cited on page 8.
Y. Liu. Discretization of class of reflected diffusion processes. Mathematics and computers in simulation, 38(1-3):103-108, 1995. Cited on page 7.
Y. Liu. Numerical approaches to stochastic differential equations with boundary conditions. Purdue University, 1993. Cited on page 7.
A. Lou and S. Ermon. Reflected diffusion models. arXiv preprint arXiv:2304.047740 2023. Cited on pages 2, 4,
L. Lovász and Vempala. The geometry of logconcave functions and sampling algorithms. Random Structures & Algorithms, 30(3):307-355 2007. Cited on page 7.
J. M. Lukens, K. Law, A. Jasra, and P. Lougovski. A practical and efficient approach for bayesian quantum state estimation. New Journal of Physics, 22(6):063038, 2020. Cited on page 1.
Met Office. Cartopy: cartographic python library with Matplotlib interface. Exeter, Devon, 2010
2015. URL: https::/t/t http:::/ssccttoll:: org. uk/ cartopy. Cited on page 15.
B. J. Morris. Improved bounds for sampling contingency tables. Random Structures & Algo- rithms, 21(2):135-146, 2002. DOT: https / doi org / 10 1002 / rsa 10049. eprint: https: : // onlinelbbrrrrr wiley com / doi / påf / 10 1002 rsa 10049. URL: https //onlimelibrary wiley. com/doi/abee
1002 10049. Cited on page 1.
Natural Earth. Natural earth, 2023, URL: https // www naturalearthdata, com/. [Online;
accessed 31 31-Jan-200 -023]. Cited on page 15.
Y. Nesterov and A. Nemirovskii. Interior-point polynomial algorithms in convex programenog. SIAM, 1994. Cited on page 4.
M. Noble, V. De Bortoli, and A. Durmus. Barrier Hamiltonian Monte Carlo, Oct. 2022. DO1: 10. 48550/ arXiv 2210.2211000
10...19999 11925. Cited page 2.
B. Pacchiarotti, C. Costantini, and F. Sartoretto. Numerical approximation for functionals of reflecting diffusion processes. SIAM Journal on Applied Mathematics, 58(1):73-102, 1998. Cited on page
Petit. Time reversal and reflected diffusions. Stochastic Processes and their Applications, 69())225- 53, July 1997. ISSN: 03044149, DOI: 10. 10016//00000-4444917700003333330 5. Cited on page 2.
R. Pettersson. Approximations for stochastic differential equations with reflecting convex boundaries. Stochastic processex and their applications, 59(2):295-308, 1995. Cited on page
R. Pettersson. Penalization schemes for reflecting stochastic differential equations. Bernoulli:403- 414, 1997. Cited on page 7.
A. Pilipenko. An introduction to stochastic differential equations with reflection, volume 1. Univer- sitätsverlag Potsdam, 2014. Cited on page 7.
K. Ramanan. Reflected diffusions defined via the extended skorokhod map, 2006. Cited on pages 4, 14.
P. Rudiger, J. A. Bednar, J.-L. Stevens, M. Liquet, C. B, S. H. Hansen, B. Little, Andrew, J. Signell, M. Hedley, C. Bosley, R. Hattersley, G. Brener, ea42gh, kbowen, A. Hilboll, I. Lustig, J. Werdd N. Hand, R. Signell, J. Bampton, pmav99, scainel, and zassa. Holoviz/geoviews: version 1.9.6, version v1.9.6, Jan. 2023. DOI: 10. 5281 / zenodo. 7543863. URL: stttps//dde org 10. 5281/zenodo 7543863. Cited on page 16.
C. Saharia W. Chan, S. Saxena, L. Li. J. Whang, E. L. Denton, K. Ghasemipour, R. Gontijo Lopes, B. Karagol Ayan, T. Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems, 35:36479 36494, 2022. Cited on page
M. Shkolnikov and I. Karatzas. Time-reversal of reflected Brownian motions in the orthant, July 2013.
URL: http: sttp://arxits arxiv. org/ abs//307.... 1107 4422. Cited
page 2.
13

---

A. V. Skorokhod. Stochastic equations for diffusion processes in bounded region. Theory of Probability & Its Applications, 6(3):264-274, 1961. Cited on page 3.
L. Slomiński. On approximation of solutions of multidimensional sde's with reflecting boundary conditions. Stochastic processes and their Applications, 50(2):197-219, 1994. Cited on pages 4, 7.
J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamies. In International Conference on Machine Learning, pages 2256- 2265. PMLR, 2015. Cited on page 1.
Y. Song and S. Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems, 2019. Cited on pages 1, 3.
Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021. Cited on pages 1, 3.
D. W. Stroock and S. S. Varadhan. Diffusion processes with boundary conditions. Communications on Pure and Applied Mathematics, 24(2):147-225, 1971. Cited on pages 2, 6, 4, 14, 15,
I. Thiele, N. Swainston, R. Fleming, A. Hoppe, S. Sahoo, M. Aurich, H. Haraldsdottir, M. Mo, O. Rolfsson, M. Stobbe, S. Thorleifsson, R. Agren, C. Bölling, S. Bordel, A. Chavali, P. Dobson, W. Dunn, [. Endler, D. Hala, M. Hucka, D. Hull, D. Jameson, N. Jamshidi, Jonsson, N. Juty, S. Keating, I. Nookaew, N. Le Novère, N. Malys, A. Mazein, J. Papin, N. Price, E. Selkov, M. Sigurdsson, E. Simeonidis, N. Sonnenschein, K. Smallbone, A. Sorokin, J. van Beek, D. Weichart, Goryanin, J. Nielsen, H. Westerhoff, D. Kell, Mendes, and B. Palsson. A community-drive global reconstruction of human metabolism. English. Nature Biotechnology, Mar. 2013. ISSN: 1087-0156. DOI: 10. 1038/nbt..... 2488. Cited on page
B. L. Trippe, Yim, D. Tischer, T. Broderick, D. Baker, R. Barzilay, and Jaakkola. Diffusion probabilistic modeling of protein backbones in 3D for the motif-scaffolding problem, June 2022.
DOr: 10. 48550/ arriv..iiiivv.iii/ 2206 04119. Cited on page 1.
J. Urain, N. Funk, G. Chalvatzaki, and J. Peters. Se (3)-diffusionfiells:: learning cost functions for joint grasp and motion optimization through diffusion. arXiv preprint arXiv:2209.33855 2022. Cited on page
Vincent. A connection between score matching and denoising autoencoders. Neural Computation, 23(7):6661-1674 2011. Cited on page 3.
J. L. Watson, D. Juergens, N. R. Bennett, B. L. Trippe, J. Yim, H. E. Eisenach, W. Ahern, A. J. Borst, R. J. Ragotte, L. F. Milles, B. I. M. Wicky, N. Hanikel, S. J. Pellock, A. Courbet, W. Sheffler, J. Wang, P. Venkatesh, I. Sappington, S. V. Torres, A. Lauko, V. D. Bortoli, E. Mathieu, R. Barzilay, T. S. Jaakkola, F. DiMaio, M. Back, and D. Baker. Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models, Dec. 2022. DOI: 10. 1001111222.2111 12. 09. 519842. Cited on page
Welty and M. Jeffries. Combined wildfire datasets for the united states and certain territories, 1878-2019: us geological survey data release, 2020. Cited on page 15.
R. J. Williams. Reflected Brownian motion with skew symmetric data in a polyhedral domain. Probability Theory and Related Fields, 75(4):459 485, Aug. 1987. ISSN: 0178 0178-8051, 1432-2064.
DOr: 10. 1007 /007/BF00320322 Cited on page 2.
J. Yim, B. L. Trippe, V. De Bortoli, F. Mathieu, A. Doucet, R. Barzilay, and T. Jaakkola, Se (3) diffusion model with application protein backhone generation. arXiv preprint arX1v:2302.02277, 2023. Cited on page 1.
Yoshikawa. Manipulability of robotic mechanisms. The international journal of Robotics Research, 4(2):3-9, 1985. Cited on page 9.
14

---

Supplementary to:
Metropolis Sampling for Constrained Diffusion Models
A Overview
In Appendix B, we recall some basic concepts of Riemannian geometry which are key defining discretisations of the reflected Brownian motion. In Appendix C, we give some details on the reflection step in reflected discretizations. In Appendix D, we prove the convergence the rejection and Metropolis discretizations to the true reflected Brownian Motion. The geospatial dataset with non-convex constraints based on wildfire incidence rates in the continental United States presented Appendix E. All supplementary experimental details and empirical results are gathered in Appendix F.
B Manifold concepts
In the following, aim introduce key concepts that underpin diffusion models on Riemannian manifolds, with a particular focus on notions relevant to the reflected Brownian motion that we build on in Appendix C. For more thorough treatment with reference to reflected diffusion models, refer to (Fishman et al., 2023). For a detailed presentation of smooth manifolds, see Lee (2013).
A Riemannian manifold is tuple (M, g) with M smooth manifold and metric that imbues the manifold with a notion distance and curvature and defined as a smooth positive-defintte inner product each of the tangent spaces of the manifold:
((p) T T.M T.M +
The tangent space Tp of point p manifold is extension of the notion of tangent planes and can be thought as the space of derivatives of scalar functions on the manifold at that point.
To establish how different tangent spaces relate to another, need to additionally introduce the concept connection. This is map that takes two vector fields and produces a derivative of the first with respect to the second, typically written V(X, Y) V Y. While there are infinitely many connections on any given manifold, the Levi-Cevita emerges as a natural choice if we impose the following two conditions:
(i) X (g(Y, 2) = = VxZ), (ii) [X VrX.
where [ is the Lie bracket. These conditions ensure that the connection is (i) metric-preserving and (ii) corsion-free. with the latter guaranteeing unique connection and integrability on the manifold.
Using the metric and Levi-Cevita connection, we can define a number of key concepts:
Geodesic. Geodesies extend the Euclidean notion of "straight lines" manifolds. They are defined as the unique path y (0, 1) + M such that V.rr = 0 and are the shortest path between two points on manifold, in the sense that L(y) = fo (t) (t))dt is minimal.
Exponential map. The exponential map manifold is given by the mapping between element E T,. of the tangent space at point and the endpoint of the unique geodesic y with 7()) and (0) V.
The intersection along geodesic is the first point which the geodesic inter- sects the boundary. We recall that the boundary is defined by f = O. We can define this via an optimisation problem: compute the minimum > 0 such that have that exp, (2. t2) is root of f: ff (exp, t2)) 0. We will say that exp (2) intersect, , f) and that t arg intersect, (2,2;)
Intersection.

---

Parallel transport. We say that a vector field X is parallel to curve (0,1) MM 0. where y/: (0,1) + MM For two points on the manifold p. E M that are connected by a curve 7. and an initial vector Xo E M. there is a unique vector field X that is parallel to y such that ( = Xo. This induces a map between the tangent spaces p and q 74 T.M - M, which is referred to as the parallel transport of tangent vectors between and and satisfies the condition that for v, u E TpM (()) u) = (T,) T,), (u)).
Reflection. For an element E M.M in the tangent space of the manifold point p and constraint characterised by its unit normal vector n E Tp.M, the reflection of v in the tangent space is given by r 2g(2 n.
C Full Reflected Discretisation
Here, we reproduce the central algorithm for the full discretisation of the reflected Brownian motion (Algorithm 4) derived for Euclidean models in (Lou et al., 2023) and for Reimannian models in (Fishman et al., 2023). Its central component is the Reflected Step Algorithm (Algorithm 3), which gives generic computation for the reflection in any manifold. Due the need to balance speed and numerical instability issues around the boundary, an efficient practical implementation of the reflected step highly non-trivial, even for simple polytopes Euclidean space. More complex geometries and boundaries make this problem significantly worse: a constraint on the trace of SPD matrices under the log-Cholesky metric of (Lin, 2019) requires solving complex non-convex optimisation problems for each sample at each discretised sampling step in both the forward and reverse process. This motivates our work in this paper.
These problems motivated the development of our Metropolis approximation, which significantly simplifies the random walk. Instead of requiring the intersection, parallel transport and reflection, we simply need to be able to evaluate the constraint functions ft. We highlight this simplicity in Algorithm
Reflected Step Algorithm. The algorithm operates by repeatedly taking geodesic steps until one of the constraints is violated, or until the step is fully taken. Upon hitting the boundary, we taralleltrannportt the tangent vector to the boundary and then reflect against it. We then start new geodesic from this point in the new direction. The arg intersect, function computes the distance one must travel along geodesic in direction until constraint fi violated. For discussion of paralleltannport, exp, and reflect see appendix B.
Algorithm 3
Require: x E M, E T.M, {fsheez
{
[
+ v/|
while 2 do
d; arg intersect, (2, Ji)
f- arg min; d; min (d., E)
x' + exp, a (x, as)
8 + paralleltransport. 8,
8 -- reflect(s, fi)
g - 8-a æ x'
end while return æ

---

Algorithm 4 Reflected Random Walk. Discretisation of the SDE dx, = b(1, d.) + dB. -dk,
Require: T, N,N,X) (fther
y T/N
for E {0. .0...... N 13 do
Zk+1 ~ N(0, Id)
X ReflectedStepl 2441 (fibez)
end for return {x])==
Algorithm 5 Metropolis Random Walk. Discretisation of the SDE dX, b(t, X;)df dB. dk,
Require: N,N,X ifihez 7 T/N
for E {O..... N - 1} do
Zx+1 -d)
X' expa (X) k: vT
if maxieI fi()') S then
X = X"
else
X X
end if
end for return
{
3

---

D Convergence to the reflected process
In this note, we assume that M {ir TR ((x) > 0} is compact, with & E ( R). We have that 0. f- E - (xx) 0}. In addition, we assume that for any E A.M, [[)(())] = 1 and that & is concave. The closure of M is denoted M. The assumption that & is concave is only used in Theorem 7-(d) and can be dropped. We consider it for simplicity.
Let (XE) JEN given for any Y > and i E by X3 = E M and for X2+1 X7 + vizg
with Z Gaussian random variable conditioned on X v E M. In practice, Z2 is sampled using rejection sampling. We define X I + M given for any E N by X2, = XX and for any E [Kr,, (k + 1)7), X = Xx. Note that (o)) a D([0, T) .)) valued random variable, where D( .M) is the space of right-continuous with left- processes which take values in M. We denote IPT the distribution of ()),), on D[ M)..
Our goal is to show the following theorem.
Theorem 6. For any T 2 0.
weakly converges to (X (X))xe(0,7) such that any t [0,
= B. k2
So 1x.eandlkk.
k, So mX,))dkk)
(6)
Proof. In order to prove the result, we prove that the distribution of the Markov chain seen as an element of D( T) M)) converges to solution of the Skorokhod problem (6). particular, we first show that the limiting distribution satisfies submartingale problem following (Stroock et al., 1971, Theorem 6.3). The transition from a solution of a submartingale problem to a weak solution of the Skorokhod problem is given by (Kang et al., 2017, Theorem 1, Proposition 2.1 and (Ramanan, 2006, Corollary 10). In order to apply (Stroock et al., 1971, Theorem 6.3), we define an intermediate drift and diffusion matrix, see (55) and (51). To prove the theorem one needs to control the drift and diffusion matrix inside M and show that converges to 0 and Id respectively. The technical part of the proof comes from the control of the drift coefficient near the boundary. In particular, we show that if the intermediate drift is large then we are close to the boundary and the intermediate drift is pointing inward. To investigate the local properties of the drift near the boundary we rely on the notion of tubular neighborhood, see (Lee et al., 2012, Theorem 6.24).
Some key properties of the tubular neighborhood are stated in Appendix D.I. We then establish few technical lemmas about the tail probability of some distributions in Appendix D.2. Controls on the diffusion matrix and lower bounds on the probability of belonging in M are given in Appendix D.3. Properties of large drift terms are given in Appendix D.4. The convergence of the drift and diffusion matrix on compact sets is given in Appendix D.5. The convergence of the boundary terms is investigated in Appendix D.6. Finally, we conclude the proof in Appendix D.7.
D.1 Properties of the tubular neighborhood
Using the results of (Lee et al., 2012), we establish the existence of an open set of M (for the induced topology of Rd) satisfying several important properties.
Theorem There exist U M open and C Z 1, > 0 such that for any y E (0,5) with 5 = 1 the
following properties hold:
(a) For any E U, there exist unique ā E 8. and ā > 0 such that 2 (V).
(b) For any ā E [0, r) and & E aM such that E + (V E M. let x 7 (V and C(x,7) such that x: + v/z E C(x, if
- 1/2 S a < F. 112
||ll S (27 + ()/(CT),
(7)
with a 4Z) with 1 V Then C(a,n) M.
(c) V o +
3 aM, a E [0.F)] is open in M.
(d) For any x E U, x E ( 7) then a 2 Fy 1/2 or ||e|2 2 (an 1/2 +a)/(Cn)
and 1/2 + 2 0, with av)(T) + Fu, with ã given in (a) and v 1 V4(T).

---

(e) There exists R > O such that fx E M: d(z, 0M) S 2R) C V.
Proof. Let y E (0,7) with 7 1. First, note that for any 2 E i.M, the normal space is given by
(α) a E RR. Using this result and (Lee et al., 2012, Theorem 6.24) there exists ño > 0
such that Ug- = {i + (T) 2 ZEOM, a E Fo)) C R is open". We have that for any
[-ro. 0) and E am
D V( = (T) + a [|] + fo V2 taV 2f dt
(8)
Ea+Cho' <0,
with ro = mminffn./(((C))) where we have used that (()) = 0, [ = 1 and defined
Co = sup + (z)(l) IE a E [ofo, Reciprocally, we have for any
E (0) and 8M
o ( = x + Vo/
a - Goa
(10)
where we have used that = E 8M, a [-To. rol). Let r min 1/(2)))). Then, U, ( + av
0, [[][ = 1 and defined Co supp + aV)())
OM, a a (-ri,r)) Rd is open and
Un (=+aV((3) 3 OM, a E (0,r1))-
(11)
In what follows, we define == -J-11 M. Note that U is open for the induced topology and that
RM c U. In particular, AM is compact, U is closed and aM n U = 25. Hence, there exists > 0
such that (()M.UU) 2 4r. Without loss of generality we can assume that r S 1/2. We also have
{a E M d(x, A.M) S 2r) C U. The proof of (a) follows from the definition of Up. In the rest of the proof, we define
C²/2 = 2
a E 0., ||/| S r(r+1))), f= min(11((2 1/2..
(12)
Let us prove (b). Consider T2 C(x, 7) with C(z, 7) given by (7) and x (V and
2 = a (T) + u with v 1 (T).). In particular, we recall that we have
-āy 1/2 < < my 1/2
|ul| < 1/2 + (O))
(13)
This implies that
ā V/ S 27,
T|lla1 s 2F/C.
(14)
First, using that C 2 1,)) = 1, (14) and (12), we have
|| + -1== = (ā + V/a)² + 2|lla| < 12 + +r/C S <r(r+ (1). 1).
(15)
Then, we have that
Φ(x + RT2) = (()) + + + + fo (2 + -)== + v/2 - -r)
(16)
2 ã Vya /2 (2) v + +
(17)
where we recall that
C1/22 = 2 max ( 3 OM, ||/||2 S sr(r+1))), f= m
)//2).
(18)
First, using that S 1/2 and (14), we have + via S 2r S 1. Since, ||u||2 S (a / and we have that ||0:| < 1/(C)). Let P(X)= X - P(x) 0 if and only if [Tmin, Tmax] with
- (C4/2/2)nllulp We have that
Emin = (1 - - (11-CCClll)))))/)//CCC
Fmax = (1 (1 - (1 1//9///)//////1111
(19)
Using that for any & E (0, 1). (1 - 1)2/2 2 1-t we have that
Emin 5 ~011101///2
Tmax 2 1/C1/2
(20)
TTii is the tubular neighborhood theorem which is key to the rest of the proof.
5

---

Since || < (v7a + a)/(7C), we have that ā - VTa 2 Fmin- In addition, using that ā + v/a - 2F < 1/C < maax we get that P(a v/a)
and therefore v12 E M since
$(x v(z+ 2 This concludes the proof of (b). Note that the condition a 2 -7 a is implied by the condition ||u||@ I (via + a))(CC). Using that V C fz E M. d(x, a.) I 2r) C U, (c) is direct consequence of (Lee al., 2012, Theorem 6.24)]. Next, we prove (d). Let + v
M n C(a, 7). If < -ay 1/2 then since @ is concave, we have
(( + v12) (z) + Via+ fo V2 + t(x + vYz +)) v/2 - 3) dt < 0,
(21)
where we have used that $(3) O. This is absurd, hence either ry-112 1/2 ||0|2 > (a> 1/2 ())(Cr) and ary1/2 + â 2 0, which concludes the proof. The proof of (e) is similar to the proof that fo E M: d(x, AM) s 2r) C U.
The main message of Theorem 7 is that using Theorem 7-(d), if we move in the direction of V (the inward normal) with magnitude then are allowed to move in the orthonormal direction with magnitude a l/22, In the next paragraph, we discuss this fact in details and shows it is necessary for the rest of our study.
The necessity of Theorem 7-(b). At first sight one can wonder if the statement of Theorem 7- (b) could be simplify. In particular, would be simpler to replace this statement with: for any F) and OM such that E + aV M, let x = I + (V)( and C(2,,7 such that x V2z E C(x, if
-ay 1/2 S a < 1/2
||a|| S (a72/3 (/)///+))///SS) /C)))
(22)
with a (z) tu, with v 1 (T).). Then C(x, 7) C M. Note that ||0||2 S (a72/2+ + (CC) is replaced by || S (a-1/22 + a) (C)) see Figure for an illustration. However, in that
case Theorem 7-(d) becomes: in addition, if x + viz M n C(a,,) 7) then a 2 r. 17/1 or ||al|2 2 (a-,1/2 + a) /CT)) and 1/2 + ā 2 0.
In what follows, when controlling the properties of large drift, see the proof of Proposition 18 and the proof of Proposition 21, we need to control quantities of the form P(x + v E C(x,,) n// Using the original Theorem 7-(d) it is possible to show that this quantity is bounded. However, if one uses the updated version of Theorem 7-(d) then one needs to show that there exists M 2 0 and 5 > 0
such that for any > (0,5) (here we have assumed that ā 0, i.e. x E AM for simplicity)
So So 1/2
S M V7,
(23)
which is absurd
D.2 Technical lemmas
We start with few technical lemmas which will allow us to control some Gaussian probabilities outside of compact set. We denote V R x N - [0, 1] such that for any k E N. V (-,
(.,k) k) is the tail
probability of x-squared random variable with parameter k. i.e. for any & E N and t 2 0 we have
(t k) = IP( 22 2 t),
(24)
with 7 a Gaussian random variable in + with zero mean and identity covariance matrix. We will make extensive use of the following lemma which is direct consequence of (Laurent et al., 2000,
Section 4, Lemma 1).
Lemma 8. For any k E N and t e R4 with t 2 5k, (t, k) < expl-t/5)
Proof. Let k E N. First, note that for any x 2 k; we have that A: + 2(kxr) 2x < 5x. Combining this result and (Laurent al., 2000, Section 4, Lemma 1, Equation (4.3)), have that for any 2 k
P[| 2 5z) S exp|-2],
(25)
with X a IR -vauud Gaussian random variable with zero mean and identity covariance matrix. This 5x.
concludes the proof upon letting
The division by v/ comes from the definition of the intermediate drift (55).

---

{ (xx)= = 0}
{ D(x) 0}
v
a
T
The grey shaded area represents M while the blue shaded represents C(rr, 7) for an arbitrary value of y and x = E O.M
Figure 8:
Let o - R+ given for any u e by (2)) (2T) real Gaussian random variable with zero mean and unit variance. While Lemma 9 appears technical, will be central to provide quantitative upper bounds on the rejection probability, see Lemma for instance.
(0m))pp222 p/2 exp[-||u|] i.e. the density of
Lemma 9, For any E N, o0 > 0, 80 E (0,, 1] and > 0 we have
xb)) supiJo +00 (ant/88,k) p to/8)dt to 2 03 S Coo,
(26)
with Co = (/
Proof. Let k ag > 0, fo E (0, 1] and > Let ts 5k8/a0. Note that if 2 to then, ant/8 2 5k. In addition, we have
f V(ant/8, 1a00/88,)00000 pl - to/d)dt s (2n1-11/ 1/2 1/2
go (ant/88 Ro
(27)
s 12/2-11/2 1/2 for (ant/6 k) dt + (2Tt) 1/2 1: (ast)) 1 18, k) 80
1/2
(28)
Using that for any 0, J cxpl-wtldt (1/w), that for any 0, Vu, k) and that if 5k, (u, k) I expl-u/5) 5). get for any to 2 0
fo 1ot00 (ant//6 / 88, k)) - to/8) < (2m) 1/2 15kō au + 56/ (an (5(2m) +
1/2
(29)
which concludes the proof.
Finally, have the following lemma, which is similar to Lemma 8 but will be used to control quantities related to the norm.
Lemma 10. For any ks E N, an > 0, 3. E (0, 1] and ā > 0 we have
(8) 0 (aaot (oopp(t < Coo2,
(30)
with C. 25((2)) + 10)
Proof. Let k E N. 0o > 0, Bo E (0, 1] and 8 > 0. Let is = 5k00/00 Note that if ( > ly then, 0ol/5 2 5k. In addition, we have
oo
< (2m) fa ((ool/5, bdl + -2 +o0 (ool k3 tdtt
tdt.
(31)
'In the rest of the supplementary, we never precise the dimension E N which can be deduced from the variable.
7

---

In addition, using that if u 2 5k then ((u, k) - exp[-u/5], we get (2T))² Jex (agt/6, k) 2 (2rt) So expl-a0 (0))))too S (2m) (a),)))
(32)
Finally, using that for any u 2 0, v (2, k) S 1, we have
(2T) 2 V(oot/5, & idt s (27)--156428//o6
(33)
which concludes the proof.
D.3 Lower bound on the inside probability and control of moments of order two and higher
Lower bound on the inside probability. We begin with the following lemma which controls the expectation of + I| outside of C(zr, 7). We recall that V is defined in Theorem 7-(c).
Lemma 11. Ler 7 Ler x E V, E~ N(O, Id) and y (0,7) then we have
max ( (Btlaaaatteeeeeaaaeaa
[Z, ( S (77)
(34)
with vb R2+ - RR,, such that lim sup, (1//1/2 < +oc.
Proof. Let F > O given by Theorem 7. First, we have that
JK J2--1 (1 + |a| + [eell//aa////pppaaaaddddd
(35)
S {d (al)
+ +opp-//01))
(36)
Second, using Lemma 9, we have that
J& 1
(37)
(38)
S 1 -d (1)
(39)
Second, using Lemma 10, we have that
Sk aac-ialle/fie
= wtwp:///hhteeeeeocccccaaaa////ae
(40)
(41)
< f (0ap://"(..d)))ppp/doo
-
(42)
Note that we have limsup, 2/() + w (7 < +00. We conclude upon combining (36),
(39) and (42) with Theorem 7-(d) and the fact that | ( 11
The following lemma allow give lower bound the quantity E[ E[1 (V uniformly w.r.t
MM.
Lemma 12. There exists 7 such that for any 7 E (0,7) for any x E M. (0,7) and
2 ~ N(O,Id) Id) we have
2 1/4.
(43)
Proof. Let E (0,7). If & V then B(x, 2R) C M using Theorem 7-(e) and therefore
E S+VEZEM 2 1/4 for 7 > small enough. Now, assume that E V. Using Lemma 11,
we have that [[ AZzMMn I (r)). In addition, using Theorem 7-(b), we have that for any
7>> 7>0 0
E[1 A
2 E ElL=ttyzec(an)
(44)
> | - - /iipr//aaa////////cc///0/cdd/ddd
(45)
a // - 1aalaall-lii--alll - V((o==//aa/(Ch)ddd1))dl)
(46)
2 (1/2) (r /r/ - V((ā
(47)
Hence, using Lemma and Lemma 9, there exists 7 > such that for any E (0, 7), W (1 /7,
1)
fot to
V(a/(C)"/2)) 1/2 - 1/4, which concludes the proof.

---

Note that the result of Lemma 12 can be improved to 1/2 - = for any E > 0. In particular this result tells us that for > 0 small enough, M looks like the hyperplane from the point of view of the Gaussian with variance ry centered on ö.M.
Bound on moments of order two and higher. In what follows, we define for any > 0, A M + R given for any x: E M by
(² - (1/7) J 1z (T2l (z2)d2 J2 (et
(48)
Proposition 13. We have lim (1) E M 0.
Proof. Let 7 > 0 given by Lemma 12. Let E M and y E (0,5). We have using Lemma 12
Sa L=tvisemp())d 1/4.
(49)
We also have that
(1/7) Sal b S 3ydf. Therefore, we get that for any E 7). (z) - 12yd2, which concludes the proof.
(50)
In what follows, we define for any 0, Lr M St (R) given for any E M by
{x) = Sax V @ zp(=)
Ja 1=+)sxeMp(()dzz
(51)
Proposition 14. There exists & O such that for any æ M and y E (0,5) we have
(x)]l - 4d.
(52)
Proof. Let x E M and y > 0 given by Lemma 12. For any E (0, we have using Lemma 12
J L=tv 1/4.
(53)
We also have that
Se bet/ireelll///(()== < d.
(54)
which concludes the proof.
D.4 Properties of large drift terms
Finally, we define for any 7 > 80 M Rd given for any x E M by
- Sge 1=t/)rem=((())d/ Saa I (=)MM((())))
(55)
First, we show away from the boundary the drift iy converges to zero.
Proposition 15. There exists 7y > O such that for any E (0, 7), d(æ, aM) 2 we have [16((x)) < (d/(r//,,)) d 1/2
> 0 and x E M such that
Proof. Let E M and 7 > 0 given by Lemma 12. For any y (0,5) 5) we have using Lemma 12
S L=t/roemp(2)ds 2 1/4.
(56)
We also have that Sga L (semM(())dd)
< [1.e ( Inslssr//u//pp(()dd + Ser 1 1nellr//==lllll(((d) (57)
< 2 Sga 11lleee/vvllllllleee < 2///
(58)
which concludes the proof.
We have the following corollary.
Corollary 16. There exists 0 such that for any 8 > 0 there exists Ms > 0 such that for any 7' (0,7) and x E M, (16² (x)] 2 Ms. then (xx) 8.

---

Proof. Let 7 > 0 given by Lemma 12. Let f: R+ + R+ given for any r by f(r) suply > 0 (r/y,, 1) //2/2////.. We have that f is non-increasing and lim, f()) = too. Let 5 > 0 and M2 = (6) with C' sup(llV((z)) : M). Let E (0,5) and E M such that 16 (a)]l 2 Ms then using Proposition 15 we have that d(zz, 0M) s 8/C. Let E e 0. such that
x ø d(x, a.M. We have
( (v t(x x)), z)dt -
(59)
which concludes the proof.
For case of notation, for any > O, we define b = >/1/277 the renormalized version of the drift. First, we have the following result which will ensure that the drift projected on the normal component does not vanish.
Lemma 17. There exists 7 > such that for any E (0, 7) and a E V we have
(z), (T)) 2 [x)) *b7), with vb. : R4 + R4 such that sup, (n)/ V +00.
(60)
Proof. Let E M and y > 0 given by Lemma 12. For any y (0,5) we have using Lemma 12 Sa Lxt Letvise.p())d 1/4.
(61)
In addition, we have
Jou Aetell (+(3)) 2 fax (eedaan)) (z())) (
(62)
- - deuteetveeeeenann VZ(E))p(2).
(63)
Using Lemma 11, get that
See Aet (Z())) 2 Sar Aet V=(=)) v17)).
(64)
Let (ei)d=1 basis of ( Using Theorem 7-(b), we have that for any i E d 13
Sgr l=t/sseceen) ( fva(s) (0
(65)
Hence, combining this result and the Kauchy-Schware incquality we have for any i E (E((11...1 {L....,
13
(Szu 1xt (x (x)(()))))))) = 1/2////3 (va(e)= (tt
(66)
s $ (i (t 1/////2 V((ā az "3 (67)
s / Vi(ā + an
(68)
Hence, using Lemma 9, we get that
E Lx+ (e (e),),,)))))))))))))) ()d-)))())
(69)
with ib given by Lemma 9 with Bn 1/2. Therefore, get that
Ug (atvrrec(an) (=)))))(2)d2)
(70)
= (fee Ixt ( - 1 (ee
(71)
(Sgu (3)))
(72)
We conclude the proof upon using that for any a, 6 0, (a + b)2/ s a/2 6//2 and (61).
We are now ready state the following lower bound on the drift.
10

---

Proposition 18. There exist 7 > 0, M 2 0 and c > 0 such that for any E M and y E (0, if [16(x))] 2 M then E V and
(73)
Proof. Let > given by Lemma 12 and M, = s / (1) YE E (0, 7].. addition, let
C 1/4. Using Proposition 15 and Theorem 7-(e), there exists M1 2 0 such that for any any æ E M,
if []6)(z)] 2 M1 then z E V and = I + (V) with S 1/(4C) and C' = V M). We denote M max(MMo M11.. Let E (0,5) and z E M such that []7(x)] 2 M. Using Lemma 17, we have that
(o))))) 2 -())))
(74)
Using that //()/////3 M/2 < (b) 1/2, we have
(z), vo))) (1/2)
(75)
< 1/(4C) have (x), ( 2 (1/2-Ca)ll ( concludes the proof.
Since -|
(z))11) which
D.5 Convergence on compact sets
In this section, we show the convergence of the drift and diffusion matrix on compact sets. We recall that M does not include its boundary θ.M
Proposition 19. For any compact set K M and > O, there exists 7 > 0 such that for any y E (0,7) we have for any K
(0) S
[$) Id S
(76)
Proof. Let C .M be a compact set and y > O. Since n 8m Ø. there exists r > 0 such that for any K, d(x, a.M) Therefore, we have that for any x E K
|ه)) = 7 Sat zp(())dll/ Sa//Saaannn SatViEEM p(z)dz.
(77)
In addition, using the Cauchy-Sehhe Schharrz inequality we have
Satvreem =p(=)d=]| < Il =1(())=]1| Smr [=]]<(=)d=
(78)
s Jg 1 (112://1>=l=2llp(2)( s vau((2 (r,
(79)
Using Lemma 8 and Lemma 12, there exists 70 such that for any 7 E (0,50) we have that for any x E K
|66 (x)]l < 4r 1) SE,
(80)
which concludes the first part of the proof. Similarly, we have that for any E K
|| I1 @ z Id)x(=)dz]| < Il @= Id)p(z)d=|| + SMo I|=1b<(>)d=
(81)
fge 1 /0 0= Id [|))d)
(82)
S 2 + 3d3) (r² dr, 1/2
(83)
Using Lemma and Lemma 12, there exists 71 > such that for any 7 E (0, 72), we have that for any E K
(x) Id || < 4v 3d 3/2
1/2 1/, Se.
(84)
which concludes the proof upon letting min(7o, 71).

---

D.6 Convergence on the boundary
Finally, we investigate the behavior at the boundary of the diffusion matrix and the drift. First, we show that there is a lower bound to the diffusion matrix near the boundary. Second, we show that the renormalized drift converges the outward normal.
Proposition 20. There exist e > 0 and 7 > 0 such that for any Y E (0, E IR and x E V we have
(u, (2,)))))))) 2 c||ull?.
(85)
In particular, there exist r >> O such that for any (0, 5) and x E M with d(z, 0.M) r
(V x (T(x)) 2..
(86)
Proof. First, we show (85). Let E V. We have for any u E Rdd
(u, (z)m) = les lovisemm2z (z)dzz Jg 12-
(87)
2 S (at u)² (z))d.
(88)
For any u E Rd, let a (u, V Using Theorem 7-(b) we have for any u Rd
Jes ( (2, u) (z))dz
(89)
- (( ((u, 1) + aa.) (tap://aaaao/////aaaa//00pppppp / [oore] ((u, (),, 0)² + a hwwp////o/////a////tt/aaaaaadaad
(90)
(91)
a pr// d²pla)da + (u,) "aan
(92)
Using Cauchy-Schwarz/ SSwwwrh...wr.... inequality, we have
sg 12
d'pla)da (1/2) - - d'p(a)daa 2 (1/2) (1/2)-34(1))) 33p0)),.)))
(93)
In addition, using the Cauchy-Schwarz inequality, we have that
J-4/71/2 fro(a)= (u, w) " 11
(94)
= - (u. u/ (())d ( /" p(a)da
(95)
(96)
2 (1) 00,2/711))
(97)
- /3(dd1))llul fo000 d (e
(98)
Combining this result, (93), (92) and Lemma 9 there exists 7 > such that for any 7 E (0, 7] and Rdd
Jg 1at (reeC(e)))))) (2, (z)dd 2 (1/4)]]u/]?,
(99)
which concludes the proof of (85). Finally, using Theorem 7-(e), have that for any M if
d(x,, 8.) < R then x E V. Let r min R, 1/(2C)) with C = sup([ E M). We have that for any E M such that d(x. AM) <r
[|V(())|| >
V Cr > 1/2,
(100)
where To such that || - Foll < and To aM. Combining this result and (99) concludes the proof upon letting s = 1/16.
Finally, we investigate the behavior of the normalized drift near the boundary.
Proposition 21. For any Zo E êM and = > 0, there exist 7, r, M 0 such that for any x E M and Y (0,7) with ||æ Zoll S r and ||6 (z)]l 2 M
16 (x)/ (z), vo(x)) V
(101)
12

---

Proof. Let y be given by Proposition 18. Let & given by Lemma 9 and Mg = sup E } < +oo. Let M = /c= with e given in Proposition 18. Let R > 0 given by Theorem 7-(e) such that for any I E M with d(x, A.M) there exist ā E AM and E [0, c0,c=/(4C)) 2 E M and given in Proposition 18. L min ce/ R) and x E M with - Zoll S r. First, since d(z, iM) - R. there exist 0.M and a E [0, (/4C) such that r = (3). Therefore, we get that f Tol s e/CC)) and therefore [T [ - e/2. In addition, we have that
such that æ = Σ + (z) with C = sup[[|V²(<z))
(102)
(103)
Using Proposition 18, we get that
(1G00011)144))--070700000 (X(3) s s//4.
(104)
In what follows, we show that
(=) (1) v < 8/2.
(105)
In partieular, show that for any ( with | 1,
(b'()), (=( (c/16) (2)
(106)
Assuming (106), letting u= = (1110 ( and using that (zx) - u)) (6)) have
- o0tp (b()))
(107)
$ (
(108)
e/2.
(109)
which concludes the proof. Let u E - with |||||| = and { [ dd-1 an orthonormal basis of (T)+). There exist {a: ) such that Σ at = and Σ 04e;- Using Theorem 7-(b),
we have that for any i E (1,..... d-1
SR L (x (aanis=-ajpp 1a - (h
(110)
| = 1-a/11// 1 hwwp:///wwaaaaaaleeooooooooooooe
(111)
Hence, combining this result and the
inequality we have for any i E (1...., d-
(Jzr 1ut (es.seeeeeeeeeeeeceee)))) = (0-2//4/2 =
(112)
$ < (v, (w..)))(((e)ddd(((11
m/33// ((0++or/////yddd1)))44000
(113)
Hence, we get that
Dnn of (he be 11ull3
(114)
with de given by Lemma 9. Recalling that || (x)] 2 M we have
(x), u)² < ( < (c)1 (o)
(115)
which concludes the proof.
13

---

D.7 Submartingale problem and weak solution
We are now ready conclude the proof.
Theorem 22. There exists IP* a distribution on D( TT MM) such that lim P = P In addition, for any f E C ([o, TI x M, R) with (V (x)) 2 foror any t E [0, 7] and x E 0MM we have that the process (e)) given for any t [0,
- - (0 + (
(116)
a IP submartingale.
Proof. Condition (A) (Stroock al., 1971, p.197) is consequence of Proposition 13. Condition (B) (Stroock et al., 1971, p. 1.797) is a consequence of Proposition 18. Condition (C) (Stroock et al., 1971,
p.198) is a consequence of Corollary 16. Condition (D) (Stroock et al., 1971, p.198) is a consequence of Proposition 14. We fix p = 0 and condition (1) (Stroock et al., 1971, p.203) is a consequence of Proposition 19. Condition (2)-(iii) (Stroock et al., 1971, p.203) is consequence of Proposition 20. Condition (2)-(iv) (Stroock et al., 1971. p.203) is a consequence of Proposition 21. We conclude
upon using (Stroock al., 1971, Theorem 6.3) and (Stroock et al., 1971, Theorem 5.8).
O
We finally conclude the proof of Theorem 6 upon using the results of (Kang et al., 2017) which establish the link between weak solution to the reflected SDE and the solution to submartingale problem.
Theorem 23. For any T (X))e(oor weakly converges (Xt)cE(o.T) such that for any e [0,
X, = B. kt
x
k| = 1x
k, = [xx,n(X,)dlk
(117)
Proof. Using Theorem 22 and (Kang et al., 2017, Theorem 1, Proposition 2.12), we have that IP in Theorem 22is associated with a solution to the extended Skorokhod problem. We conclude that a solution to the extended Skorokhod problem is solution to the Skorokhod problem using (Ramanan, 2006, Corollary 2.10).
D.8 Extension to the Metropolis process
We recall that the Metropolis process is defined as follows. Let (XE)AEN given for any y >> 0 and A N by X²¹²
.M and for X =X vZz if X2 + v E M and X otherwise, Z% ~ N(O, Id). We recall that 87, Σd and À are given by (48), (51) and (55). In particular, denoting Ry the Markov kernel associated with (KE))EN, i.e. K : M x ((M) -> [0, 1] such that for any M,, (Z, -) probability measure, for any A B(M), RT (.,A) A) is a measurable function
and (X7) X7 = x² = (a: A). We have that for any 7 > and x e M
()(x) (1/7) San x)² (2, dy).
(118)
{)(x) (1/7) (1/7)/m(( 2) dy),
(119)
(² (1/7) JMn 11 x||[K1 (I,, dy).
(120)
In what follows, we denote a" (x) = E([ [[11ot/ZeeM Denote K" the kernel associated with (X2 KeNN We have that for any A E B(M).
and x E M
() A) = F1 X EA l [ (1 - (x)))(aa(()
(121)
= =x)(<) (x, A) + (1 (a(()))aa)())
(122)
We define for any 7 > 0 and x JM
(1)) = (1/7) (1 ( dy),
(123)
(x) = (K) dy),
(124)
(x) = (1/7) 11 - "I dg)
(125)
14

---

Using (122), we get that for any > 0 and E M
(x(x) a" (x)) (²)2) a² (x)
A² = ax))))
(126)
Using Lemma 12, we have that for any 7 E (0, 7) and M, 1/4.
In order to conclude for the convergence of the Metropolis process we adapt Theorem 22 and Theorem 23, We define X R+ + M given for any k E N by X = X and for any t E [No7, (K + 1)7), X2 = X2. Note that (X)) is a D(OO, M) valued random variable, where D[ MM) is the space of right-continuous with left-limit processes which take values in M.. We denote IP7 the distribution of (X))1e(orr) on D[ MM.
Theorem 24. There exists TP* distribution D( 7] .M) such that lim P = P.. In addition, for any f E C). T| M, with (V fz)) 2 0 for any E [0, T] and x E aM we have that the process (f(t, (e)) given.orr any t E [0,T]
f(t,w(t)) - (0.f + ( (
(127)
is IP submartingale.
Proof. Condition (A) (Stroock et al., 1971, p.797) is a consequence of Proposition 13 and (126).
Condition (B) (Stroock et al., 1971, p.79 is consequence of Proposition 18 and (126). Condition (C) (Stroock et al., 1971, p. 1.1 is consequence Corollary 16 and (126). Condition (D) (Stroock al., 1971, p. 198) is a consequence of Proposition 14 and (126). We fix p = 0 and condition (I) (Stroock et al., 1971. p.203) is consequence of Proposition 19 and that lim---00 a = 1 uniformly on compact subsets K M.. Condition (2)-(iii) (Stroock et al., 1971, p.203) is a consequence of Proposition 20 and (126). Condition (2)-(iv) (Stroock et al., 1971, p.203) is a consequence of Proposition 21 and (126). We conclude upon using (Stroock al., 1971, Theorem 6.3) and (Stroock
et al., 1971, Theorem 5.8).
Theorem 25. For any T 2 0, (X))s{(0 weakly converges to (Xt)cEjoTT) such that for any t t [0,
X, = x B. k,,
k|: = 1x,= = nXx,((X)
(128)
Proof. The proof is identical to Theorem 23.
E Modelling geospatial data within non-convex boundaries
To demonstrate the ability of the proposed method to model distributions whose support is restricted to manifolds with highly non-convex boundaries, we derived a geospatial dataset based on the historical wildfire incidence rate within the continental United States (described in in Appendix E. and, using the corresponding country borders, trained a constrained diffusion model by adapting the phintin-sphericaa--olyyooo conditions outlined in (Ketzner al., 2022) (described in Appendix E.2).
E.1 Derivation of bounded geospatial dataset
Specifically, we retrieved the rasterised version of the wildfire data provided by Welty et al. (2020), converted it to spherical geodetic coordinate system using the CARTOPY library (Met Office, 2010 2015), and drew a weighted subsample of size 1 x 10 We then retrieved the country borders of the continental United States from (Natural Earth, 2023) and mapped them to the same geodetic reference frame as the wildfire data. A visualization of the resulting dataset is presented in Figure 9,
E.2 loint-in-ppherical-plyyttt algorithms
The support of the data data-generatin generating distribution we aim to approximate is thus restricted to highly non-convex spherical polytope E S2 given by the country borders the continental United States. To determine whether query point g E S is within IP, we adapt an efficient reformulation of the paant-ii-sppericl-oolyyy algorithm (Bevis et al., 1989) presented in (Ketzner et al., 2022). The algorithm requires the provision of reference point 2 known to be located in P and determines
15

---

Figure 9: Orthographic projection of the wildfire dataset described in Appendix The projection is aligned with the centroid of the continental United States and zoomed in ten-fold for visual clarity. All visualisations of geospatial data were generated using the GROVIEWS (Rudiger et al., 2023) and DATASHADER (Bednar et al., 2023) libraries.
whether y is inside or outside the polygon by checking whether the geodesic between r and & crosses the polygon
odd number of times. Letting & R denote the Cartesian coordinates of point.>: S (Ketzner al., 2022) rely Cartesian reference coordinate system Q (with its z-axis given by f) and the corresponding spherical coordinate system Q to decompose the edge-crossing condition of Bevis al. (1989) into efficiently computable parts. That is, the geodesic between and r crosses an edge e; = (2%, Uz) of the polygon if:
the longitude of q in Q is bounded by the longitudes u; and by Q, i.e.
(9) E [min (vt), (v;)), max(eo (1:)) (Q(1)))
(ii) the plane specified by the normal vector p. 0g ûg represents equator that separates and r into two different hemispheres, i.e.
sign (p., f) (p.. 9)) 11
Especially when IP is fixed and the corresponding coordinate transformations and normal vectors can be precomputed for each edge, this algorithm affords an efficient and parallelisable approach to determining whether any given point on S is contained by a spherical polytope.
F
Supplementary Experimental Results
F. Evaluating log-barrier and Euclidean models
Following (Fishman et al., 2023), we approached the empirical evaluation of our Metropolis model by computing the maximum mean discrepancy (MMD) (Gretton et al., 2012) between samples from the true distribution and the trained diffusion models. The MMD is statistic that guantifies the similarity of two samples by computing the distance their respective mean embeddings in a reproducing kernel Hilbert space. For this, we use RBF kernel with the same length scales as the standard deviations of the normal distributions used to generate the synthetic distribution. We sum these RBF kernels by the weights of the corresponding components of the synthetic Gaussian mixture model.
This is essential to able to include the log-barrier the comparison since log-barrier methods suffer severe instabilities around the boundary, as the space is stretched to more and more. These instabilities cause the problems in fitting the log-barrier model and in computing the likelihood using the barrierri model.
16

---

Table 5: Maximum mean discrepancy (MMD) (4) of held-out from synthetic bimodal distribution over convex subsets of R bounded by the hypercube [-1,1] and unit simplex A Means and standard deviations are computed over different runs.
MMD
% in Manifold
Manifold
Dimension
Process
mcan
std
mean
Euclidean
0.027
0.011
0.969
Log-Barrier
0.050
0.012.
1.000
Reflected
0.041
0.008
1.000
Rejection
0.030
0.002
1.000
Euclidean
0.032
0.015
0.969
A
Log-Barrier
0.238
0.009
1.000
Reflected
0.179
0.013
1.000
Rejection
0.111
0.002
1.000
Euclidean
0.028
0.001
0.946
Log-Barrier
0.275
0.0015
1,000
10
Reflected
0,233
0.004
1.000
Rejection
0,226
0.005
1,000
Euclidean
0.069
0.004
0.992
Log-Barrier
0.66
0.006
1.000
Reflected
0.048
0.012
1.000
Rejection
0,025
0.005
1.000
Euclidean
0.074
0.004
0.991
[0, 1
Log-Barrier
0.209
0.0077
1.000
Reflected
0.085
0.006
1.000
Rejection
0.049
0.006
1,000
Euclidean
0.086
0.007
0.968
Log-Barrier
0.330
0,004
1,000
10
Reflected
0.314
0.049
1.000
Rejection 0.138
0.007
1.000
From the results in Table 5, it is clear that the log-barrier approach performs significantly than the Reflected model and the Metropolis models across all settings. This, in conjunction with numerical instabilities we encountered when attempting evaluate sample likelihoods with the log-barrier models as presented in (Fishman et al., 2023), motivated us to focus on the Reflected and Metropolis models in the main text.
Additionally, we note that the unconstrained Euclidean models outperform the constrained methods on both the simplex and the hypereube as the dimensionality of the problem space increases. Especially on the simplex, we attribute this performance primarily to the fact that the synthetic distribution is simply standard Normal with only small portion close the boundary. The amount of reflection needed to model the distribution decreases in higher dimensions, as the mass of the Normal distribution gets increasingly concentrated-wwhie Euclidean diffusion models will fit well. This same dynamic is partially responsible for the hypercube performance.
F.2
Implementational details
All source code that is needed to reproduce the results presented below is made available under https://iitub..co/oo hawcmm///oorredddee///mettttoooo
the different geometries that is available under which requires a supporting package to handle http:///ww/..oo.//////o///ttttta
We use the same architecture in of our experiments; 6-layer MLP with 512 hidden units and sine activation functions, except in the output layer, which uses linear activation function. Following (Fishman et al., 2023), we implement a simple linear function that scales the score by the distance to the boundary, approaching zero within 0.01 of the boundary. This ensures the score obeys the Neumann boundary conditions required by the reflected Brownian Motion. For the geospatial dataset within non-convex country borders, we do not use distance rescaling, Instead, substitute
17

---

it with series of step functions to rescale the score. This is proof-of-concept to show that even when computing the distance is hard, simple and efficient approximations suffice. When constructing Riemannian diffusion models on the torus and sphere for the protein and geospatial datasets, we follow (De Bortoli et al., 2022) and include an additional preconditioner for score on the manifold. We do not use the residual trick or the standard deviation trick, which are both common score-rescaling functions in image model architectures; in our setting, we find that they adversely affect model training.
For the forward/reverse process we always set = Su the forward process just reaches the invariant distribution with a linear 3-schedule. At sampling time we use N 100 steps of the discretised process. We discretise the training process by selecting random N between 0 and 100 for each example, rolling out to that time point. This lets us cheaply implement simple variance reduction technique: we take multiple samples from this trajectory by selecting multiple random N to save for each example. This technique was originally described in (Fishman al., 2023) and we find is also helpful for our Metropolis models. For all experiments, we use the ism loss with modified weighting function of (1 [1+13 t), which we found to be essential to model training. All experiments use batch size 256 with repeats per batch. For training, we use learning rate of 2 x 10-4 with cosine learning rate schedule. We trained for 100,000 batches on the synthetic examples and 300,000 batches on the real-world examples (robotics, proteins, wildfires).
10-h and then tune 31 ensure that
We selected these hyperparameters from a systematic search over learning rates (6 x 10 10-4, 2 x 10 10-4, x 10
10 2) learning rate schedules (cosine, log-linear), and batch sizes (128, 256, 512, 1024) on synthetic examples for the reflected and log-barrier models. Similar parameters worked well for both, and we used those for our Metropolis models allow straightforward comparison. We tried N = 100, 1000 for several synthetic examples but found that very large rollout times actually hurt performance for the Metropolis model, though the log-barrier performed a bit better with longer rollouts and the reflected was the same.
All models were trained on single NVIDIA Geforce GTX 1080 GPU. All of the Metropolis models presented here can easily be trained on this hardware in under 4 hours. The runtime for the log- log-barrier and reflected models is considerably longer.
F33 Synthetic Distributions Constrained Manifolds of Increasing Dimensionality
(a) Data distribution
(b) Metropolis model
(c) Reflected model
(d) Uniform distribution
Figure 10: Qualitiative comparison of samples from the data distribution, our Metropolis model, Reflected model and the uniform distribution for synthetic bimodal distribution on [-1, 1/2.
(a) Data distribution
(b) Metropolis model
(c) Reflected model
(d) Uniform distribution
Figure 11: Qualitiative comparison samples from the data distribution, our Metropolis model, Reflected model and the uniform distribution for synthetic bimodal distribution on A
18

---

F.4 Constrained Configurational Modelling of Robotic Arms
The following univariate marginal and pairwise bivariate plots visualise the distribution of different samples in
[
l2
(i) the three dimensions needed to describe ellipsoid M
E
S
and
13
(ii) the two dimensions needed to describe location in I
F,4.1 Visualisation of samples from the data distribution
100
(a) Plots of the univariate marginal and pairwise bivariate distributions of x 10" samples from the data distribution in S?_.
Plots of the univariate marginal and pairwise bivariate distributions 10 samples from data distribution in IR?.
Figure 12: Visualisation of the data distribution S x I22 using univariate marginal and pairwise bivariate plots.
F.4.2 Visualisation of samples from our Metropolis sampling- sampling-based diffusion model
L
-1
(a) Plots of the univariate marginal and pairwise bivariate distributions 10 samples from our Metropolis sampling-hasee diffusion model in S²..
(b) Plots of the univariate marginal and pairwise bivariate distributions 10 samples from our Metropolis sampling-bar baasdd diffusion model R².
Figure 13: Visualisation of the distribution learned by our Metropolis sampling- based diffusion model in S? x IR² using univariate marginal and pairwise bivariate plots.
19

---

F.4.3 Visualisation of samples from a reflected Brownian motion- based diffusion model
L
65
(a) Plots of the univariate marginal and pairwise bivariate distributions of x 10 samples from reflected Brownian motion-based diffusion model S7+.
(b) Plots of the univariate marginal and pairwise bivariate distributions of 1 10 samples from reflected Brownian motion-based diffusion model in RR..
Figure 14: Visualisation of the distribution learned by reflected Brownian motion-basee diffusion model S x R using univariate marginal and pairwise bivariate plots.
F.4.4
Visualisation of samples from the uniform distribution
1501
(a) Plots univariate marginal and pairwise bivariate distributions of x 10" samples from the uniform distribution in S.
(b) Plots of the univariate marginal and pairwise bivariate distributions of 1 x 10th samples from the uniform distribution in IR?.
Visualisation of the uniform distribution S2, x R22 using univariate marginal and pairwise bivariate plots.
Figure 15:
F.5
Conformational Modelling of Protein Backbones
The following univariate marginal and pairwise bivariate plots visualise distribution of different samples in (i) the polytope I c IR3 and (ii) the torus T used to parametrise the conformations of polypeptide chain of length N with coinciding endpoints. We refer to (Han al., 2006) for full detail on the reparametrisation and to (Fishman et al., 2023) for full description of the dataset.
20

---

F.5.1 Visualisation of samples from the data distribution
(a) Plots of the univariate marginal and pairwise bivariate distributions of x 10" samples from the data distribution in P C R
(b) Plots of the univariate marginal and pairwise bivariate distributions of 1 x 10 samples from the data distribution in T
Visualisation of the data distribution in P pairwise bivariate plots.
Figure 16:
R x T using univariate marginal and
F.5.2 Visualisation of samples from our Metropolis sampling- based diffusion model
(a) Plots of the univariate marginal and pairwise bivariate distributions of 10 samples from our Metropolis model in C 3
Plots of the univariate marginal and pairwise bivariate distributions 10 samples from our Metropolis model in T
Figure 17: Visualisation of the distribution learned by our Metropolis model in IP T x T using univariate marginal and pairwise bivariate plots.
21

---

F,5.3 Visualisation of samples from a reflected Brownian motion- based diffusion model
(a) Plots of the univariate marginal and pairwise bivariate distributions of x 10 samples from reflected Brownian motion-based diffusion model
(b) Plots of the univariate marginal and pairwise bivariate distributions of 1 10 samples from reflected Brownian motion-based diffusion model in T
in P C R
Figure 18: Visualisation of the distribution learned by reflected Brownian motion- baaied diffusion model in P c R x T using univariate marginal and pairwise bivariate plots.
E55.4
Visualisation of samples from the uniform distribution
(a) Plots of the univariate marginal and pairwise bivariate distributions of x 10" samples from the uniform distribution
(b) Plots of the univariate marginal and pairwise bivariate distributions of 1 x 10 samples from the uniform distribution in T
3
Figure 19: Visualisation of the uniform distribution in 123 T using univariate marginal and pairwise bivariate plots.
22


Reviews:
[{"review_id": "edVntCwz7H", "reviewer": "Reviewer_u5aU", "summary": "This paper introduces a new approach for generative modelling on constrained Riemannian manifolds, building upon diffusion models. The proposed method implements a Metropolis sampling scheme and offers computational efficiency and improved empirical performance over previous models, particularly with increased constraint complexity. It provides a valid discretisation of the reflected Brownian motion and demonstrates its scalability and flexibility across several application domains.", "questions": "MCMC types of approaches sometimes suffer from long 'burn-in', I am wondering whether the authors have consider this case and evaluate its impact to the performance of the proposed method\n\nIs it possible that Metropolis will give some samples that are not good enough (not representative of the latent space) for diffusion model, and mislead the training, - as samples accepted by Metropolis doesn't guarantee they are good samples.", "limitations": "Limitations are stated, e.g., potential poor performance on certain constraint geometries. The future work sounds sensible.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper has good originality of incorporating MCMC type of sampling with diffusion model on constrained Riemannian manifolds.\n\nThe overall quality of the paper is good, the proposed method should be of interests to a group of researchers working on diffusion models.", "weaknesses": "The paper is a bit hard to follow, it would be helpful to further refine the presentation of the methodology sections. Specifically, it would be nice to have a diagram illustrating the relationship between diffusion model, Metropolis sampler, and constrained Riemannian manifolds.\n\nIn my view, some of the mathematics in section 2 can be moved to supplementary, too much symbols in this section makes readers to spend time on each definition and formula, and hard to understand the whole picture and the relationship between different components.\n\nSome illustrations might need further refinement, for example, Figure 5 (b) and (c) it is very hard to tell the difference between them and it conveys little information.\n\nNumerical comparison with existing methods and evaluation metrics are relatively limited, the proposed method is mainly compared with the reflected method. Limited metric, does log likelihood comprehensively reflect the main performance improvement for such different datasets (robotics and proteins) Especially the proteins dataset, the changes is relative small, i.e., from 15.2 to 15.33.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our manuscript and for providing helpful and constructive feedback.\nWe were happy to see your review emphasising the improved computational efficiency and empirical performance of our approach, as well as acknowledging the quality of the paper and its relevance to other researchers in this space.\n\nThe concerns you raised in your review relate to the presentational clarity and empirical evaluation of the manuscript. We will address each of these points in turn below.\n\n## Empirical Evaluation\n\n> Numerical comparison with existing methods and evaluation metrics are relatively limited, the proposed method is mainly compared with the reflected method. Limited metric, does log likelihood comprehensively reflect the main performance improvement for such different datasets (robotics and proteins) Especially the proteins dataset, the changes is relative small, i.e., from 15.2 to 15.33.\n\nWe provide both additional comparisons to the Log-Barrier method proposed in [4] as well as comparisons using the MMD instead of the log-likelihood in Section F.1 of the Appendix, to which we refer at the beginning of Section 5. The agreement between the MMD and the log-likelihood signals the log-likelihood is a reasonable metric. It is true that for the proteins dataset the improvement is small; this is corroborated by both the visualizations and the MMD. The time to train (as well as the sampling time) is still a factor of 10 faster though, which is a crucial advantage for scaling this approach up to larger proteins.\n\n\n## Presentational Clarity\n\n\n> In my view, some of the mathematics in section 2 can be moved to supplementary, too much symbols in this section makes readers to spend time on each definition and formula, and hard to understand the whole picture and the relationship between different components.\n\nWhile we agree that shortening the outline of the proof would make the paper more accessible to applications-driven researchers, we believe that the theory presented in Section 3.2 is one of the central contributions of the manuscript is, which is why we chose to include it. We believe the proof is interesting and useful for machine learning theorists working on non-standard diffusion models.\n\n> Some illustrations might need further refinement, for example, Figure 5 (b) and (c) it is very hard to tell the difference between them and it conveys little information.\n\nThe difference between Figures 5b and 5c is hard to discern because the models fit the data distribution equally well - though, as mentioned in the previous section, our approach fits it slightly better and is 10 times faster. We believe that this makes any qualitative visual comparison difficult, and would like to mention that Figures 4 and 5 are mainly included to provide an intuitive visualisation of the data, learned and uniform distribution on constraint geometries that may not be well-known in the manuscript's intended target audience.\n\n> Specifically, it would be nice to have a diagram illustrating the relationship between diffusion model, Metropolis sampler, and constrained Riemannian manifolds.\n\nWe agree that such a diagram would make the connection between Figure 1 and the methodology outlined in Sections 2 and 3 more apparent. We have added a figure to the appendix to clarify how these components are related and will refer to it in the main text.\n\n## Questions\n\n> MCMC types of approaches sometimes suffer from long 'burn-in', I am wondering whether the authors have consider this case and evaluate its impact to the performance of the proposed method\n\n\nWe want to emphasize that our paper is not a MCMC paper. We are only marginally concerned with the mixing time of the algorithm as our target is always the uniform distribution for the forward noising process. We agree that it would be valuable to obtain 1) quantitative non-asymptotic result regarding the convergence of the proposed scheme (however, we again emphasize that we only target the uniform distribution) 2) bounds on the bias incurred by the choice of the stepsize.\n\nHowever, in practice the mixing time and the discretization stepsize are not the bottlenecks of the performance of the algorithm. As shown in [1,2], the main source of error in denoising diffusion models (in the unconstrained setting) is the approximation of the network. This is because the convergence to the invariant distribution is exponential (with respect to the total variation distance for instance) and the bias incurred by the discretization stepsize is largely dominated by the neural network approximation of the score.\n\n> Is it possible that Metropolis will give some samples that are not good enough (not representative of the latent space) for diffusion model, and mislead the training, - as samples accepted by Metropolis doesn't guarantee they are good samples.\n\nThis is a valid point that may warrant further exploration. However, we think that this is unlikely, as our central theoretical contribution shows that the Metropolis approach is a valid discretisation of the reflected Brownian motion, which was demonstrated to produce high-quality samples on the constrained geometries we consider [3,4].\n\n---\n\nWe hope that our additional clarifications and proposed changes to the layout of the paper address all of your questions and concerns.\n\nPlease let us know if you have any further questions!\n\n---\n\n## References\n\n[1] Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, Anru R. Zhang – Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions (2022)\n\n[2] De Bortoli – Convergence of denoising diffusion models under the manifold hypothesis (2022)\n\n[3] Lou and Ermon. \"Reflected diffusion models.\" ICML 2023\n\n[4] Fishman, et al. \"Diffusion Models for Constrained Domains.\" TMLR, 2023"}, {"review_id": "zHFpZDrmXM", "reviewer": "Reviewer_fvg9", "summary": "This paper proposes a Metropolis sampling for constraint diffusion in the context of generative modeling. The authors show that the proposed algorithm is a discretization of reflected Brownian motion on manifold (with rejection). The authors also apply the proposed algorithm to several synthetic and real data sets. ", "questions": "See weaknesses for my comments on the paper.", "limitations": "NA", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "This paper considers a Metropolis sampling algorithm on manifolds in light of the recent development of denoising diffusion models in generative modeling. The paper is well written, and I enjoyed reading it.", "weaknesses": "There are several weakness.\n\n(1) As the authors pointed out themselves, the idea of (Metropolis) sampling on manifolds is much related to ball walk. There have been extensive work in this direction, see e.g. Dwivedi, Chen, Wainwright, and Yu's paper Log-concave sampling: Metropolis-Hastings algorithms are fast, JMLR, 2018. The authors failed to make connections to recent developments (and literature).\n\n(2) In the MCMC literature, there are two general \"discretization' techniques: ULA (unadjusted Langevein Algo) and MALA (Metropolis adjusted Langevin Algo). The proposed algorithm is close to the idea of MALA, which the authors failed to make connections. \n\n(3) The main theorems (Theorem 2 & 4) are quite standard to experts in reflected Brownian motion. This generally from follows the techniques developed by Burdzy, Ramanan, Williams..., in the light of Stroock-Varadhan. Only Proposition 5 seems to be new (but not hard).\n\n(4) One of the most important aspects of MCMC sampling is the mixing time (or convergence rate). This paper does not seem to provide any related analysis. Though on the continuous level this can be read from existing results, the mixing time in the discrete algorithm is important (especially concerning the choice of the step size $\\gamma$).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our manuscript and for the helpful and constructive feedback. We were happy to see you emphasize that we extend methods for generative modeling on constrained Riemannian manifolds [1,2] and are glad you enjoyed the paper.\n\nThe concerns you raised relate to connections to existing prior art and an analysis of the mixing time of our algorithm. We will address each of these points in turn below.\n\n## Prior Literature\n\n### 1\n\nWe agree with the reviewer that the proposed forward process is linked with the ball walk (and related algorithms). We do not claim that our algorithm is superior for sampling from an unnormalized density defined on the constrained manifold. We are aware that there is a strong line of works dealing this problem (see e.g. [10] exhibiting strong theoretical guarantees). One of the contributions of our paper is to showcase that we can use these methods to derive efficient numerical schemes for generative modeling under constraints. To the best of our knowledge, the ball walk and derived algorithms have not been applied to define forward noising process for diffusion model algorithms. Nevertheless, we will extend our discussion of related work with [3,7,8,9] and the references therein. \n\n### 2\n\nWe will add a paragraph making the connection with MALA and the distinction with ULA. However, our point of view is somewhat different than the usual motivation behind MALA, which, in the unconstrained setting, acts as a way to remove the bias from ULA, see Besag's comment: \"if instead one uses s' merely as a Hastings proposal for the next state, then the usual acceptance probability ensures that p is an exact stationary distribution of the modified Markov chain\" in [4]. Our proposal can be recast in this context but our motivation is different since we aim at a computationally efficient discretisation of the forward noising  of the reflected diffusion.\n\n### 3\n\nThis is an important point. We are not aware of any existing proof of Theorem 2 and Theorem 4. We would be grateful if the reviewer could point us towards such a reference. We agree that the proof is not difficult and that our work is an application of the techniques of Stroock and Varadhan to prove the sub-martingale property. However, we would like to make the following comments:\n* If the proof does not exist in the literature, we believe it is valuable for the community (and especially newcomers to the field, as not everyone in ML is familiar with reflected processes). We strongly believe that making such results available to the community will help to foster the interaction between the ML community and the MCMC community. \n* While the proof is technical (and again we agree that our aim is to verify the conditions of [3], and not introduce any novel techniques to show that a numerical scheme is a valid discretisation), the main difficulty is to obtain the lower bounds in Proposition 5 and to control the behavior of the drift near the boundary (see for example Proposition 20). To do so, we needed Theorem 7 which ensures that, provided enough regularity on the boundary, one can obtain such controls. To the best of our knowledge, this approach is new. \n\nWe would appreciate if the reviewer could point us to a reference proving Theorems 2 and 4. Regarding Prop 5, we do not believe a result must be hard to serve as the theoretical foundation for a novel algorithm exhibiting state-of-the-art empirical performance.\n\n## 4 Mixing Time\n\nWe want to emphasize that our paper is not an MCMC paper. The only target we consider for the forward noising process is the uniform distribution on the constrained manifold under consideration. In practice, we set the parameter $T$ (running time of the forward noising process) so that, at time $T$, the distribution of the forward process is close to the uniform. While the parameter $T$ is important, it is easily tunable and in all of our experiments we observe fast mixing of the Markov chain. We agree that it would be valuable to obtain 1) quantitative non-asymptotic result regarding the convergence of the proposed scheme (note: we only target the uniform distribution) 2) bounds on the bias incur by the choice of the stepsize. \n\nHowever, in practice the mixing time and the discretization stepsize are not bottlenecks in performance of the algorithm. As shown in [5,6], the main source of error in denoising diffusion models (in the unconstrained setting) is the approximation of the score. This is because the convergence to the invariant distribution is exponential (with respect to the total variation distance for instance) and the bias incurred by the discretization stepsize is dominated by the neural approximation of the score.\n\n---\n\nWe hope that our additional clarifications and discussion address your questions and concerns.\n\nPlease let us know if you have any further questions!\n\n## References\n\n[1] Lou and Ermon. \"Reflected diffusion models.\" ICML 2023\n\n[2] Fishman, et al. \"Diffusion Models for Constrained Domains.\" TMLR, 2023\n\n[3] Dwivedi, Chen, Wainwright, and Yu's Log-concave sampling: Metropolis-Hastings algorithms are fast, JMLR, 2018\n\n[4] J. Besag (1994) --  \"Comments on \"Representations of knowledge in complex systems\" \n\n[5] Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, Anru R. Zhang -- Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions (2022)\n\n[6] De Bortoli -- Convergence of denoising diffusion models under the manifold hypothesis (2022)\n\n[7] Vempala -- Geometric random walks: a survey\n\n[8] Cousins, Vempala -- Bypassing KLS: Gaussian Cooling and an $O(n^3)$ Volume Algorithm (2015)\n\n[9] L. Lovasz and M. Simonovits. Random walks in a convex body and an improved volume algorithm. (1993)\n\n[10] Kook, Lee, Shen, Vempala -- Condition-number-independent convergence rate of Riemannian Hamiltonian Monte Carlo with numerical integrators (2022)"}, {"review_id": "0jFQimXEyr", "reviewer": "Reviewer_a8R9", "summary": "The paper proposes a new discretization of the forward and backward SDEs used for learning and sampling in the diffusion generative model framework in manifolds with inequality constraints. The method relies on \"metropolising\" a given discretisation of the \"unconstrained\" manifold (either Euler-Maruyama if the manifold is $\\mathbb{R}^d$ or Geodesic random walks if it's not). The discretisation is shown to converge weakly to the reflected SDE and is shown to be numerically simpler to implement and faster for diffusion models on manifolds.", "questions": "I feel that a lot of the implementation could be better explained, even though the algorithm is simple. The code should be made available and I would also like the authors to be more precise into how do they calculate the log likelihood of the test set, even though this can be considered somehow standard in the literature.\n\nHow the epochs were chosen for the log likelihood calculations? It'd be interesting to have a graph of computing time vs log-likelihood of the held out test set to understand how those behave in time for several dimensions for each algorithm (metropolised vs reflected).\n\nThe geospatial data example seems to show that the learned distribution is « coarser » than the target distribution and has some concentration around the boundaries. Is this something to expect? Maybe rejecting a lot of steps in the more non convex parts of boundaries lead to an augmentation of the likelihood? ", "limitations": "None", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper proposes a simple approach for sampling and training diffusion models on constrained manifolds. The approach is theoretically justified and is shown to be numerically more efficient than the state of the art. Thus, the proposed algorithm is an interesting tool for practitioners. The paper is mostly well written and accounts well for the state of the art. Several examples with increasing degree of complexity are shown.", "weaknesses": "The paper possesses some shortcomings in the presentation, namely some typos that I will refer to in the bottom part of this section.\nThere are also some fundamental things missing. For example, I could not find the code, since it's not given in the appendix and the link to the anonymous github was inactive when I tried to access it. This is a problem, specially when for example the paper do not explicitly tells how it calculates the log likelihood of the test samples, thus I'm admitting it uses the ODE formulation as presented for example in [1].\n\nI also feel that some other metric would also be of great value. The paper uses only the log likelihood of the held out test set as a metric, but some other distribution related metric such as the sliced wasserstein can be applied in the first example, specially in the lower dimensional examples.\n\nSome typos:\nSection 3.1 Practical imitations → Pratical limitations\nProposition 5 repeat the definition of Z\n\n\n[1] De Bortoli, V., Mathieu, E., Hutchinson, M., Thornton, J., Teh, Y. W., & Doucet, A. (2022). Riemannian score-based generative modelling. Advances in Neural Information Processing Systems, 35, 2406-2422.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our manuscript and for providing helpful and constructive feedback.\nWe were happy to see your review emphasising the improved computational efficiency and empirical performance of our approach and acknowledging its usefulness to a range of practitioners.\n\nThe concerns you raised in your review relate to the availability of our code, the log-likelihood evaluation and the use of an additional performance metric. We will address each of these points in turn below.\n\n---\n\n## Code Availability\n\n> For example, I could not find the code, since it's not given in the appendix and the link to the anonymous github was inactive when I tried to access it. This is a problem, specially when for example the paper do not explicitly tells how it calculates the log likelihood of the test samples, thus I'm admitting it uses the ODE formulation as presented for example in [1].\n> \n> The code should be made available and I would also like the authors to be more precise into how do they calculate the log likelihood of the test set, even though this can be considered somehow standard in the literature.\n\n\nWe have checked the anonymised links on multiple different devices and they seem to work correctly, as they did when we originally added them to the manuscript. If the code anonymisation site was down at any point during the review period, we apologise for the inconvenience. Hopefully you can access them now.\n\n## Additional Performance Metric\n\n> The paper uses only the log likelihood of the held out test set as a metric, but some other distribution related metric such as the sliced wasserstein can be applied in the first example, specially in the lower dimensional examples.\n\nWe would like to point out that we present further evaluations in Appendix F.1 using the MMD distance, as well as a comparison to an additional baseline method presented in [1]. If you would like us to additionally compare the models using the sliced Wasserstein distance, we are happy to do so during the author discussion period. The agreement between the MMD and the log-likelihood seems to signal the log-likelihood is a reasonable metric.\n\n## Log-likelihood Evaluation\n\n>  I would also like the authors to be more precise into how do they calculate the log likelihood of the test set, even though this can be considered somehow standard in the literature.\n\nWe will add details of how we compute the likelihood in the main paper, but it is indeed as an ODE in the reviewers reference. More precisely, we compute the log-likelihood using the equivalence between diffusion models and continuous normalizing flows (see [2,3] for instance). In the case of reflected diffusion models this equivalence was first used in [4]. As in continuous normalizing flow approaches, the log-likelihood computation is performed leveraging tools from Neural ODE [5].\n\n> How the epochs were chosen for the log likelihood calculations? It'd be interesting to have a graph of computing time vs log-likelihood of the held out test set to understand how those behave in time for several dimensions for each algorithm (metropolised vs reflected).\n\nWe fixed the number of optimization steps for the two methods at a point when the optimization had converged and used the same end point for both. This choice is somewhat arbitrary, and we can easily add a plot showing how the log-likelihood evolves over the course of training during the discussion period if the reviewer would like.\n\n> The geospatial data example seems to show that the learned distribution is « coarser » than the target distribution and has some concentration around the boundaries. Is this something to expect? Maybe rejecting a lot of steps in the more non convex parts of boundaries lead to an augmentation of the likelihood?\n\nWe agree that there may be some concentration around the non-convex parts of the boundaries. This is particularly interesting given the clear convergence of the forward process to the invariant distribution. We think it is likely that it is caused not by the Metropolis discretization but by the added complexity the score network must learn to \"jump\" out of the non-convex parts of the constrained set. It seems plausible that more careful architecture engineering could fix this problem, but we think the example in the manuscript serves as a proof of concept for the method.\n\n---\n\nWe hope that our additional clarifications and discussion address all of your questions and concerns.\n\nPlease let us know if you have any further questions!\n\n## References\n\n[1] Fishman et al. \"Diffusion Models for Constrained Domains.\" TMLR, 2023\n\n[2] Song et al. \"Score-Based Generative Modeling through Stochastic Differential Equations\", 2021\n\n[3] Huang et al. \"A Variational Perspective on Diffusion-Based Generative Models and Score Matching\", 2021\n\n[4] Lou and Ermon. \"Reflected Diffusion Models\", 2023\n\n[5] Chen et al. \"Neural Ordinary Differential Equations\", 2018"}, {"review_id": "feFK4cNH4f", "reviewer": "Reviewer_vCsC", "summary": "The paper presents a rejection-sampling technique applied to diffusion models when sampling defined on manifolds.  \n\nThe paper contains some theoretical results on the convergence of the proposed algorithm, in particular the assymptotic convergence to the reflected Brownian motion. The also paper contains a few comparisons to the other methods, in particular, methods using regularization (such as barrier functionals) and reflected stochastic processes. ", "questions": "See above. ", "limitations": "The authors mentions some of the limitations of their work, and there is not direct adverse impact of this work. ", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper applies the notion of rejection sampling to denoising diffusion models. \nThe proposed algorithm is very simple to modify to code, and it require much less intrucive modifications to the sampling algorithm compared to the alternatives the authors mention. ", "weaknesses": "- It needs to be clarified what the algorithmic innovation is in the paper. Rejection sampling is a standard technique, and the application to diffusion models directly applies an old idea to a trendy topic. Also, it is not clear how valuable is the theoretical treatment. The authors show that the methodology is asymptotically equivalent to RBM, but then in the experiments, the authors claim that their algorithm outperforms RBM. \n\n- The authors mention that scalability is one of the main advantages of its method. However, it is not clear what the scalability is with respect to. Also, for algorithms (which is the case here), scalability is often required to show how a bigger problem can be split into different pieces that can be run independently so that the overall runtime remains low. In this case, only one experiment showcasing scalability can be found in Section 5.1. The authors seem to claim their method is scalable with respect to the dimension, but still most of the wall-time are not that different. Along the same lines, I am not sure how interesting sampling from a cube (or a polytope) is in this case. I would assume that most of the \"mass\" of the distribution is along the boundary, which would make the problem easier, as seen in Table 2, where both methods start to do better as d increases. \n\n- Most of the theory seems to be around Manifolds of co-dimension 0, which is only mildly interesting. There is not much theory when the manifold has a positive co-dimension. Also, there is a mismatch between the definitions of $\\mathcal{M}$. In Line 72, it is defined as an open subset of another Riemannian manifold $\\mathcal{N}$, which is not clearly defined, using several inequalities, and there are the assumptions in Line 207, which greatly simplifies the setup. What are the assumptions that this paper is considering?.\n\n- The manifolds are smooth and of co-dimension 0; however, some of the experiments consider non-smooth manifolds, and others with manifolds of co-dimension greater than 0, thus, there seems to be a disconnect between the theory and experiments. Also, there is no quantitative metric for the sampling quality in section 5.3. They seem to \"look\" better than the very weak baseline of a uniform sampling, but it is not clear that the sampling does capture the correct statistics. \n\n- In general, the experiments are fairly weak, particularly the baselines. Using a uniform distribution is by no means a baseline, and it should be used as a sniff test. One baseline that the authors could have used is the sample on a sphere using their algorithm (seen the sphere embedded in a 3-dimensional space) and using a diffusion model defined on the sphere. Such a comparison would make the point that it can outperform a method specifically tailored for the geometry. Also, the reference to the figures is not really informative, particularly when compared to weak baselines. I would prefer to have some quantitative information. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our manuscript and for providing helpful and constructive feedback. We were happy to see your review emphasising the implementational simplicity of our approach, as well as the resuting ease with which it can be integrated into existing sampling algorithms.\n\nThe concerns you raised in your review relate to the originality and scalability of our method, our choice of baselines, and the type of manifolds covered by our theoretical justification. We will address each of these points in turn below.\n\n## Algorithmic Innovation\n\n> It needs to be clarified what the algorithmic innovation is in the paper. ... Also, it is not clear how valuable is the theoretical treatment.\n\nRejection sampling is indeed a standard technique; in this manuscipt we present several contributions. First is exactly as the reviewers notes: an application of this old idea to a trendy topic, and evidence that using this idea in this setting is more effective than existing approaches in the literature. Second is the theoretical treatment we give which is necessary to justify the use of rejection sampling in the diffusion models we present. The asymptotic equivalence establishes the RBM is the continuous time process corresponding to the metropolis discretization; this means we can swap the metropolis discretization in for a reflected walk in reflected diffusion models using the same losses and time-reversal. This is highly non-obvious as the dynamics of the two discretizations are superficially quite different. \n\n> The authors show that the methodology is asymptotically equivalent to RBM, but ... claim that their algorithm outperforms RBM.\n\nYes, this is because our approach affords much more tractable learning dynamics (see e.g. the discussion in Section 3.1), which enables it to unambiguously outperform the less tractable RBM-based diffusion models across a range of constraint geometries and applications, both in terms of empirical performance and computational efficiency.\n\n## Scalability\n\n> The authors mention that scalability is one of the main advantages of its method. ...\n\nThe proposed method is scalable in two senses: first it is much more efficient in general than the existing approaches in the literature, and hence is inherently _more_ scalable than existing approaches; second it is in fact scalable in exactly the way the reviewer notes when the underlying manifold is a product of constrained manifolds, e.g. a product of simplices or a product of SPD matrices with bounded trace. In this case we indeed have that the problem can be broken up into pieces and run independently, allowing straightforward scaling of this more efficient method in this specific regime. \n\n## Empirical Evaluation\n\n> In general, the experiments are fairly weak, ... I would prefer to have some quantitative information.\n\nWe apologize for the confusion and would like to emphasize that we benchmark our method against recently developed state-of-the-art constrained diffusion models, providing detailed empirical comparisons with models based on the reflected Brownian motion [1,2] and the log-barrier metric [2].\nAs neither of these methods extends to manifolds with non-convex bundaries, we currently do not include a baseline for the wildfire example on the sphere. We only provide the uniform distribution as it is a relatively standard practice in diffusion papers to highlight what the invariant distribution looks like, especially for non-standard invariant distributions like the uniform over the constrained set we investigate here. Motivated by your comment, we have run an additional experiment on the sphere to provide a quantitative comparison to an unconstrained model. We just present MMD results here: \n\n| Model    | MMD | % in boundary |\n| - | - | - |\n| Unconstrained | $0.1567 \\pm 0.013$ | $63.3%$ |\n| Constrained  | $0.1388 \\pm 0.015 $  | $100%$ |\n\nWe compute the RBF-MMD as described in the appendix, removing samples outside the boundary.\n\n## Manifolds covered by our proof\n\n>Most of the theory seems to be around Manifolds of co-dimension 0, which is only mildly interesting. ... What are the assumptions that this paper is considering?.\n>\n> The manifolds are smooth and of co-dimension 0 ... there seems to be a disconnect between the theory and experiments.\n\nWe agree with the reviewer that the theoretical results of Theorem 3 and Theorem 4 are stated for manifolds of co-dimension 0. We believe that the results can be extended to a more general setting. One approach to do so would be to replace the local parameterization of the manifold as $\\\\{x \\in \\mathbb{R}^d,\\Phi(x) < 0\\\\}$ by similar inequalities in local charts. Using the paracompact property of the manifolds under consideration it seems possible to extend our results. However, this approach seems to be highly technical and we postpone it to future work. \n\nThe definition in line 72 is to provide the broad class of manifolds which our method cover. Our experiments fall in the class of the manifolds defined in line 72. However, we greatly simplify the structure of the manifold to give the theoretical result in line 207. So to summarize:\n* The methodology is defined on the class defined in line 72. Experiments are defined w.r.t. manifolds defined in line 72. \n* In order to simplify the analysis, the validity of the discretization scheme is only assessed under stronger assumptions (as pointed out by the reviewer).\n\nWe thank the reviewer for pushing us to clarify this point and will update the paper accordingly. We leave the extension of our theoretical results to future work (as highlighted above this extension is likely to be highly technical).\n\n\n---\n\nWe hope that our additional clarifications and discussion address all of your questions and concerns.\n\nPlease let us know if you have any further questions!\n\n---\n\n\n## References\n\n[1] Lou and Ermon. \"Reflected diffusion models.\" ICML 2023\n\n[2] Fishman et al. \"Diffusion Models for Constrained Domains.\" TMLR, 2023"}]

Author rebuttal:
# General Response

We would like to thank all reviewers for the time and effort they have put into reviewing our manuscript and for the valuable and constructive feedback they have provided.

---

In our manuscript, we present a new method for generative modelling on constrained Riemannian manifolds that affords substantial gains in computational efficiency and empirical performance compared to the current state of the art [1, 2].

We are delighted to see that reviewers appreciated the practical significance of our method across a range of settings, highlighting that it is both "of interest to a group of researchers working on diffusion models" (**reviewer u5aU**) and an "interesting tool for practitioners" (**reviewer a8R9**).

We are also happy to see that reviewers appreciated the "good originality" (**reviewer u5aU**) and "theoretical justification" (**reviewer a8R9**) of our approach, noting that "the proposed algorithm is very simple to modify [and] code" (**reviewer vCsC**), "requires much less intrusive modifications to the sampling algorithm" (**reviewer vCsC**) and is "numerically more efficient than the state of the art" (**reviewer a8R9**). 

Additionally, we are pleased that reviewers found the paper to be "well-written" (**reviewers fvg9 and a8R9**) and that they "enjoyed reading it" (**reviewer fvg9**).

---

The reviewers have also raised a range of points regarding different parts of the paper. To address these, we have responded with additional discussions and clarifications of our work under the respective reviews. We have also added a baseline model to the paper, which further shows the strength of the proposed method (discussed in response to **reviewer vCsC**). We hope that these remarks convincingly address the reviewers' concerns and are happy to answer any further questions.


Sincerely,

The Authors

---


## References

[1] Lou and Ermon. "Reflected diffusion models." ICML 2023

[2] Fishman et al. "Diffusion Models for Constrained Domains." TMLR, 2023

Meta-review: The paper offers a novel approach to challenging issues in generative modeling by extending by now standard denoising diffusion probabilistic models (DDPM) to Riemannian manifolds. Efficient solutions for handling manifolds defined by numerous constraints— commonly faced in various applied sectors like natural sciences— remain outside the reach of most existing methodologies. Previous works attempted to address this problem through complex noising processes (eg. introducing complex project steps), leading to higher computation demands.

The proposed strategy in this paper introduces a simplified noising scheme based on rejection sampling (i.e. simply reject "increments" when stepping out of the boundaries), contributing significant improvements in computational efficiency and empirical performance. The independent analysis proving the validity of this new process as a discretization of the reflected Brownian motion is both technically sophisticated and interesting (while the proposed method is very simple). Despite its simplicity, the method's effectiveness is convincingly demonstrated by numerical experiments that show that it can substantially outperform the current state-of-the-art. The authors manage to convincingly apply the proposed method to a range of problem settings.

Some reviewers proposed a relatively low score. Nevertheless, after careful reading of these reviews, it was concluded that they were either not relevant to the paper (eg. discussion of other MCMC schemes or mixing times considerations), or convincingly answered by the authors during the rebuttal phase (eg. numerical stability, scalability, numerical experiments).

In conclusion, the paper's proposed solution, with both its theoretical analysis and its proven practical impact, makes a strong case for acceptance.
Decision: Accept (poster)