SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\mathcal{M}$ ,  $\forall \varepsilon > 0$ ,  $\exists \bar{\gamma} > 0$  such that for any  $\gamma \in (0, \bar{\gamma})$  and for any  $x \in \mathcal{M}$ ,  $\gamma \in (0, \bar{\gamma})$  and  $Z \sim \mathrm{N}(0, \mathrm{Id})$  we have  $\mathbb{P}(x + \sqrt{\gamma} Z \in \mathcal{M}) \geq 1/2 - \varepsilon$ , with  $Z \sim \mathrm{N}(0, \mathrm{Id})$ .
Proposition 5 tells us that locally the boundary looks like a half-space when integrating w.r.t. a Gaussian measure. A corollary is that, for  $\gamma >0$  small enough and for any  $k\in \mathbb{N}$ , the probability that  $X_{k + 1}^{\gamma} = X_{k}^{\gamma}$  is upper bounded uniformly w.r.t.  $X_{k}^{\gamma}\in \mathcal{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 as  $\gamma \rightarrow 0$ . This is enough to apply the same results as before, giving the desired convergence.
Assumptions on  $\mathcal{M}$ . Before concluding this section, we detail the assumptions we make on  $\mathcal{M}$ . For Theorem 2 to hold, we assume that  $\mathcal{M} = \{x\in \mathbb{R}^d:\Phi (x) > 0\}$  is bounded, with  $\Phi \in \mathrm{C}^2 (\mathbb{R}^d,\mathbb{R})$  concave. We have that  $\partial \mathcal{M} = \{x\in \mathbb{R}^d:\Phi (x) = 0\}$ . In addition, we assume that for any  $x\in \partial \mathcal{M}$ ,  $\| \nabla \Phi (x)\| = 1$ . 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  $\Phi$  to a 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 neighbourhoods. 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.
4 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 Euler-Maruyama schemes used to discretise SDEs without boundary. One of the most common approaches is to use the Euler-Maruyama discretisation and project the solution onto the boundary if it escapes from the domain  $\mathcal{M}$ .
Table 2: Log-likelihood  $(\uparrow)$  of a held-out test set from a synthetic bimodal distribution over convex subsets of  $\mathbb{R}^d$  bounded by the hypercube  $[-1,1]^d$  and unit simplex  $\Delta^d$ . Means and standard deviations are computed over 3 different runs. Average training time is provided in hours.
MANIFOLD	DIMENSION	REFLECTED	METROPOLIS
log-likelihood	runtime	log-likelihood	runtime
[-1,1]d	2	2.25±.01	8.95	2.32±.05	0.72
3	3.77±.13	8.97	4.15±.15	0.71
10	7.42±.77	10.1	10.80±.34	0.90
Δd	2	1.01±.01	9.17	1.06±.02	0.82
3	2.64±.01	9.43	3.23±.17	0.78
10	7.00±.13	10.5	7.81±.20	0.97
Figure 4: Convergence time of the Reflected (green) and Metropolis (red) forward noising processes to the uniform distribution on the hypercube  $[-1,1]^d$  and unit simplex  $\Delta^d$ . The lines indicate functions fit with the PYSR symbolic regression package (Cranmer, 2023) and correspond to empirical runtime complexities of  $\mathcal{O}(d^2)$  and  $\mathcal{O}(d)$ , 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; Słominski, 1994). Concretely this means that  $\mathbb{E}[\| \bar{\mathbf{B}}_t - X_n^{t/n} \|^2] = O(n^{-1 + \varepsilon})$  with  $\varepsilon > 0$  arbitrary small and where  $(X_k^\gamma)_{k \in \mathbb{N}}$  is the projection scheme. The rate  $1/2$  is tight (Pacchiarotti et al., 1998). It is possible to use the Euler-Peano method to get slight improvements for a mean-square error rate of order of  $1/2$ , but this is impractical as it assumes that one can 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  $1$  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 to a half-space. We refer to 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 of 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 et 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 best of our knowledge, its limiting behaviour when the step size goes to zero has not been investigated.
(a) Data distribution.
(b) Metropolis samples.
(c) Reflected samples.
Figure 5: A qualitative visual comparison of  $10^{6}$  samples from the data distribution, our Metropolis diffusion model, a reflected diffusion model and the uniform distribution for the constrained configurational modelling of robotic arms on  $\mathcal{S}_{++}^2 \times \mathbb{R}^2$ .
(d) Uniform distribution.
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 on 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 our method extends to constrained manifolds with highly non-convex boundaries—a setting that is intractable with existing approaches.
As we found—in line with Fishman et al. (2023)—that log-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 a 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 et al. (2022) and Fishman et al. (2023). We set  $T = 1$ ,  $\beta_0 = 1 \times 10^{-3}$  and tune  $\beta_{1}$  to ensure that the forward process reaches the invariant distribution with a linear  $\beta$ -schedule. We use a learning rate of  $2 \times 10^{-4}$  with a cosine learning rate schedule and an ism loss with a modified loss weighting function of  $(1 + t)$ , a 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.
All source code that is needed to reproduce the results presented below is made available under https://github.com/oxcsml/score-sde/tree/metropolis, which requires a supporting package to handle the different geometries that is available under https://github.com/oxcsml/geomstats/tree/polytope.
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 hypercube  $[-1, 1]^d$  and unit simplex  $\Delta^d$ . 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 Metropolis 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 distribution on hypercubes  $[-1, 1]^d$  and simplices  $\Delta^d$  of increasing dimensionality. The results are presented in Section 4 and show that the convergence time of the Metropolis process scales linearly in the dimension, while the reflected process scales quadratically.
Figure 6: A qualitative comparison of  $10^{5}$  samples from the data distribution, our Metropolis diffusion model, a reflected diffusion model and the uniform distribution for the constrained conformational modelling of cyclic peptide backbones. For visual clarity, the figures only show the constrained planar projections encoded by  $\mathbb{P} \subset \mathbb{R}^3$ .
(c) Reflected samples.
(d) Uniform distribution.
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 et al. (2023) to 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 configurations and trajectories of a robotic arm can be formulated as learning a distribution over the locations and manipulability ellipsoids of its joints, parameterised on  $\mathbb{R}^d\times S_{++}^d$ , where  $S_{++}^{d}$  is the manifold of symmetric positive-definite (SPD)  $d\times d$  matrices (Yoshikawa, 1985; Jaquier et al., 2021). For practical robotics applications, it may be desirable to restrict the maximal velocity with which a robotic arm can move or the maximum force it can exert. This manifests in a trace constraint  $C > 0$  on  $S_{++}^{d}$ , resulting in a constrained manifold  $\{M\in S_{++}^{d}:\sum_{i = 1}^{d}M_{ii} < C\}$ . Following Fishman et al. (2023), we parametrise this constraint via the Cholesky decomposition (Lin, 2019) and use the resulting setup to model the dataset presented in Jaquier et al. (2021).
Conformational modelling of protein backbones. Modelling the conformational ensembles of proteins is a data-scarce problem with a range of important applications in biotechnology and drug discovery (Lane, 2023). In many practical settings, it may often be unnecessary to model the structural ensembles of an entire protein, as 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 on 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 the backbone conformations of a cyclic peptide of length  $N = 6$  as learning a distribution over the product of a polytope  $\mathbb{P} \subset \mathbb{R}^3$  and the hypertorus  $\mathbb{T}^4$ .
Table 3: Log-likelihood  $(\uparrow)$  of a held-out test set for the robotics and protein applications. Means and standard deviations are computed over 3 different runs. Average training time is provided in hours.
DATASET	DOMAIN	REFLECTED	METROPOLIS
log-likelihood	runtime	log-likelihood	runtime
Robotics	S2++×R2	8.39±.06	9.52	9.13±.03	1.36
Proteins	P⊂R3×T4	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.
(c) Uniform distribution.
Figure 7: Orthographic projection of  $10^{5}$  samples from (a) the data distribution, (b) our Metropolis diffusion model, and (c) the uniform distribution, for geospatial data (wildfire incidence rates) within a 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 on tasks with challenging convex constraints, we investigated its ability to model distributions whose support is restricted to manifolds with highly non-convex boundaries—a setting that is intractable with existing approaches. To this end, we derived a geospatial dataset based on wildfire incidence rates within the continental United States (see Appendix E for full details) and trained a Metropolis diffusion model constrained by the corresponding country borders on the sphere  $S^2$ . 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^2$  (De Bortoli et al., 2022) is given in Table 4. Both demonstrate that our approach is able to successfully capture challenging multimodal and sparse distributions on constrained manifolds with highly non-convex boundaries.
Table 4: MMD  $(\downarrow)$  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	MMD	RUNTIME	% IN BOUNDARY
Unconstrained	S2	0.1567±0.013	0.81	63.3
Metropolis	P ⊂ S2	0.1388±0.015	3.86	100.0
Accurately modelling distributions on constrained Riemannian manifolds is a challenging problem with a range of impactful practical applications. In this work, we have proposed a mathematically principled and computationally tractable extension of existing diffusion model methodology to this setting. Based on a Metropolisation of random walks in Euclidean spaces and on 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 our method, we 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. a Dikin ellipsoid, while maintaining the simplicity and efficiency of the Metropolis approach.
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.
B. D. Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313-326, 1982. Cited on 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 on page 7.
J. A. Bednar, J. Crail, J. Crist-Harif, P. Rudiger, G. Brener, I. Thomas, C. B, J. Mease, J. Signell, M. Liqueit, J.-L. Stevens, B. Collins, A. Thorve, thuydotm, S. H. Hansen, esc, kbowen, N. Abdennur, O. Smirnov, maihde, A. Hawley, A. Oirekhov, A. Ahmadia, B. A. B. Jr, C. H. Brandt, C. Tolboom, E. G., E. Welch, J. Bourbeau, and J. J. Schmidt. Holoviz/datashader: version 0.14.4, version v0.14.4, Feb. 2023. DOI: 10.5281/zenodo .7599872. URL: https://doi.org/10 .5281/zenodo . 7599872. Cited on page 16.
M. Bevis and J.-L. Chatelain. Locating a point on a spherical surface relative to a spherical polygon of arbitrary shape. Mathematical geology, 21: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, 41(3):877-889, 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 a finite entropy condition. arXiv preprint arXiv:2104.07708, 2021. Cited on page 3.
R. Chitashvili and N. Lazrieva. Strong solutions of stochastic differential equations with boundary conditions. Stochastics: an international journal of probability and stochastic processes, 5(4):255-309, 1981. Cited on page 7.
G. Corso, H. Stärk, B. Jing, R. Barzilay, and T. Jaakkola. DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking, Oct. 2022. URL: http://arxiv.org/abs/2210.01776. Cited on page 1.
M. Cranmer. Interpretable machine learning for science with pysr and symbolicregression. jl.arXiv preprint arXiv:2305.01582, 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 [cs.LG]. Cited on pages 1-4, 10, 18.
N. Fishman, L. Klarner, V. De Bortoli, E. Mathieu, and M. Hutchinson. Diffusion models for constrained domains. arXiv preprint arXiv:2304.05364, 2023. Cited on pages 2, 4, 5, 8, 9, 1, 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 on page 7.
A. Gretton, K. M. Borgwardt, M. J. Rasch, B. Scholkopf, and A. Smola. A kernel two-sample test. Journal of Machine Learning Research, 13(null):723-773, Mar. 2012. ISSN: 1532-4435. Cited on page 16.
L. Han and L. Rudolph. Inverse kinematics for a serial chain with joints under distance constraints. In Robotics: Science and systems, 2006. Cited on pages 1, 20.
In Appendix B, we recall some basic concepts of Riemannian geometry which are key to 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 of 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 is presented in Appendix E. All supplementary experimental details and empirical results are gathered in Appendix F.
B Manifold concepts
In the following, we aim to 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 a more thorough treatment with reference to reflected diffusion models, we refer to (Fishman et al., 2023). For a detailed presentation of smooth manifolds, see Lee (2013).
A Riemannian manifold is a tuple  $(\mathcal{M},\mathfrak{g})$  with  $\mathcal{M}$  a smooth manifold and  $\mathfrak{g}$  a metric that imbues the manifold with a notion of distance and curvature and is defined as a smooth positive-definite inner product on each of the tangent spaces of the manifold:
$$
\mathfrak {g} (p): \mathrm {T} _ {p} \mathcal {M} \times \mathrm {T} _ {p} \mathcal {M} \to \mathbb {R}.
$$
The tangent space  $\mathrm{T}_p$  of a point  $p$  on a manifold is an extension of the notion of tangent planes and can be thought of as the space of derivatives of scalar functions on the manifold at that point.
To establish how different tangent spaces relate to one another, we need to additionally introduce the concept of a connection. This is a map that takes two vector fields and produces a derivative of the first with respect to the second, typically written as  $\nabla(X,Y) = \nabla_X 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\cdot (\mathfrak{g}(Y,Z)) = \mathfrak{g}(\nabla_XY,Z) + \mathfrak{g}(Y,\nabla_XZ),$
(ii)  $[X,Y] = \nabla_X Y - \nabla_Y X,$
where  $[\cdot, \cdot]$  is the Lie bracket. These conditions ensure that the connection is (i) metric-preserving and (ii) torsion-free, with the latter guaranteeing a unique connection and integrability on the manifold.
Using the metric and Levi-Cevita connection, we can define a number of key concepts:
Geodesic. Geodesics extend the Euclidean notion of 'straight lines' to manifolds. They are defined as the unique path  $\gamma : (0,1) \to \mathcal{M}$  such that  $\nabla_{\gamma'} \gamma' = 0$  and are the shortest path between two points on a manifold, in the sense that  $L(\gamma) = \int_0^1 \sqrt{\mathfrak{g}(\gamma(t)) (\gamma'(t), \gamma(t))} \mathrm{d}t$  is minimal.
Exponential map. The exponential map on a manifold is given by the mapping between an element  $\pmb{v} \in T_p\mathcal{M}$  of the tangent space at point  $p$  and the endpoint of the unique geodesic  $\gamma$  with  $\gamma(0) = p$  and  $\gamma'(0) = \pmb{v}$ .
Intersection. The intersection along a geodesic is the first point at which the geodesic intersects the boundary. We recall that the boundary is defined by  $f = 0$ . We can define this via an optimisation problem: compute the minimum  $t > 0$  such that we have that  $\exp_{\mathfrak{g}}(x, tz)$  is a root of  $f$ :  $f(\exp_{\mathfrak{g}}(x, tz)) = 0$ . We will say that  $\exp_{\mathfrak{g}}(x, tz) = \operatorname{intersect}_{\mathfrak{g}}(x, z; f)$  and that  $t = \arg \operatorname{intersect}_{\mathfrak{g}}(x, z; f)$ .