SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 5: Maximum mean discrepancy (MMD)  $(\downarrow)$  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.
Manifold	Dimension	Process	MMD	% in Manifold mean
mean	std
Δd	2	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
3	Euclidean	0.032	0.015	0.969
Log-Barrier	0.238	0.009	1.000
Reflected	0.179	0.013	1.000
Rejection	0.111	0.002	1.000
10	Euclidean	0.028	0.001	0.946
Log-Barrier	0.275	0.0015	1.000
Reflected	0.233	0.004	1.000
Rejection	0.226	0.005	1.000
[0,1]d	2	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
3	Euclidean	0.074	0.004	0.991
Log-Barrier	0.209	0.0077	1.000
Reflected	0.085	0.006	1.000
Rejection	0.049	0.006	1.000
10	Euclidean	0.086	0.007	0.968
Log-Barrier	0.330	0.004	1.000
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 worse than the Reflected model and the Metropolis models across all settings. This, in conjunction with numerical instabilities we encountered when attempting to 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 hypercube 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 a standard Normal with only a small portion close to 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—which Euclidean diffusion models will fit well. This same dynamic is partially responsible for the hypercube performance.
F.2 Implementational details
We use the same architecture in all of our experiments: a 6-layer MLP with 512 hidden units and sine activation functions, except in the output layer, which uses a 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  $\epsilon = 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, we substitute
it with a series of step functions to rescale the score. This is a 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 the 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  $T = 1$ ,  $\beta_0 = 1 \times 10^{-3}$  and then tune  $\beta_{1}$  to ensure that the forward process just reaches the invariant distribution with a linear  $\beta$ -schedule. At sampling time we use  $N = 100$  steps of the discretised process. We discretise the training process by selecting a random  $N$  between 0 and 100 for each example, rolling out to that time point. This lets us cheaply implement a 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 et al., 2023) and we find it is also helpful for our Metropolis models. For all experiments, we use the ism loss with a modified weighting function of  $(1 + t)$ , which we found to be essential to model training. All experiments use a batch size of 256 with 8 repeats per batch. For training, we use a learning rate of  $2 \times 10^{-4}$  with a 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).
We selected these hyperparameters from a systematic search over learning rates  $(6 \times 10^{-4}, 2 \times 10^{-4}, 6 \times 10^{-5}, 2 \times 10^{-5})$ , 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 to allow a 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 a 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-barrier and reflected models is considerably longer.
F.3 Synthetic Distributions on Constrained Manifolds of Increasing Dimensionality
(b) Metropolis model
(c) Reflected model
(d) Uniform distribution
Figure 10: Qualitative comparison of samples from the data distribution, our Metropolis model, a Reflected model and the uniform distribution for a synthetic bimodal distribution on  $[-1,1]^2$ .
(a) Data distribution
Figure 11: Qualitative comparison of samples from the data distribution, our Metropolis model, a Reflected model and the uniform distribution for a synthetic bimodal distribution on  $\Delta^2$ .
F.4 Constrained Configurational Modelling of Robotic Arms
The following univariate marginal and pairwise bivariate plots visualise the distribution of different samples in
(i) the three dimensions needed to describe an ellipsoid  $M = \begin{bmatrix} l_1 & l_2 \\ l_2 & l_3 \end{bmatrix} \in S_{++}^2$  and
(ii) the two dimensions needed to describe a location in  $\mathbb{R}^2$
F.4.1 Visualisation of samples from the data distribution
(a) Plots of the univariate marginal and pairwise bivariate distributions of  $1 \times 10^{5}$  samples from the data distribution in  $S_{++}^{2}$ .
(b) Plots of the univariate marginal and pairwise bivariate distributions of  $1 \times 10^{5}$  samples from the data distribution in  $\mathbb{R}^2$ .
F.4.2 Visualisation of samples from our Metropolis sampling-based diffusion model
Figure 12: Visualisation of the data distribution in  $\mathcal{S}_{++}^{2} \times \mathbb{R}^{2}$  using univariate marginal and pairwise bivariate plots.
(a) Plots of the univariate marginal and pairwise bivariate distributions of  $1 \times 10^{5}$  samples from our Metropolis sampling-based diffusion model in  $\mathcal{S}_{++}^2$ .
Figure 13: Visualisation of the distribution learned by our Metropolis sampling-based diffusion model in  $S_{++}^{2} \times \mathbb{R}^{2}$  using univariate marginal and pairwise bivariate plots.
(b) Plots of the univariate marginal and pairwise bivariate distributions of  $1 \times 10^{5}$  samples from our Metropolis sampling-based diffusion model in  $\mathbb{R}^2$ .
F.4.3 Visualisation of samples from a reflected Brownian motion-based diffusion model
(a) Plots of the univariate marginal and pairwise bivariate distributions of  $1 \times 10^{5}$  samples from a reflected Brownian motion-based diffusion model in  $\mathcal{S}_{++}^2$ .
(b) Plots of the univariate marginal and pairwise bivariate distributions of  $1 \times 10^{5}$  samples from a reflected Brownian motion-based diffusion model in  $\mathbb{R}^2$ .
F.4.4 Visualisation of samples from the uniform distribution
Figure 14: Visualisation of the distribution learned by a reflected Brownian motion-based diffusion model in  $S_{++}^2 \times \mathbb{R}^2$  using univariate marginal and pairwise bivariate plots.
(a) Plots of the univariate marginal and pairwise bivariate distributions of  $1 \times 10^{5}$  samples from the uniform distribution in  $S_{++}^{2}$ .
Figure 15: Visualisation of the uniform distribution in  $S_{++}^2 \times \mathbb{R}^2$  using univariate marginal and pairwise bivariate plots.
(b) Plots of the univariate marginal and pairwise bivariate distributions of  $1 \times 10^{5}$  samples from the uniform distribution in  $\mathbb{R}^2$ .
F.5 Conformational Modelling of Protein Backbones
The following univariate marginal and pairwise bivariate plots visualise the distribution of different samples in (i) the polytope  $\mathbb{P} \subset \mathbb{R}^3$  and (ii) the torus  $\mathbb{T}^4$  used to parametrisse the conformations of a polypeptide chain of length  $N = 6$  with coinciding endpoints. We refer to (Han et al., 2006) for full detail on the reparametrisation and to (Fishman et al., 2023) for a full description of the dataset.
F.5.1 Visualisation of samples from the data distribution
(a) Plots of the univariate marginal and pairwise bivariate distributions of  $1 \times 10^{5}$  samples from the data distribution in  $\mathbb{P} \subset \mathbb{R}^3$ .
(b) Plots of the univariate marginal and pairwise bivariate distributions of  $1 \times 10^{5}$  samples from the data distribution in  $\mathbb{T}^4$ .
F.5.2 Visualisation of samples from our Metropolis sampling-based diffusion model
Figure 16: Visualisation of the data distribution in  $\mathbb{P} \subset \mathbb{R}^3 \times \mathbb{T}^4$  using univariate marginal and pairwise bivariate plots.
(a) Plots of the univariate marginal and pairwise bivariate distributions of  $1 \times 10^{5}$  samples from our Metropolis model in  $\mathbb{P} \subset \mathbb{R}^3$ .
(b) Plots of the univariate marginal and pairwise bivariate distributions of  $1 \times 10^{5}$  samples from our Metropolis model in  $\mathbb{T}^4$ .
Figure 17: Visualisation of the distribution learned by our Metropolis model in  $\mathbb{P} \subset \mathbb{R}^3 \times \mathbb{T}^4$  using univariate marginal and pairwise bivariate plots.
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  $1 \times 10^{5}$  samples from a reflected Brownian motion-based diffusion model in  $\mathbb{P} \subset \mathbb{R}^3$ .
(b) Plots of the univariate marginal and pairwise bivariate distributions of  $1 \times 10^{5}$  samples from a reflected Brownian motion-based diffusion model in  $\mathbb{T}^4$ .
F.5.4 Visualisation of samples from the uniform distribution
Figure 18: Visualisation of the distribution learned by a reflected Brownian motion-based diffusion model in  $\mathbb{P} \subset \mathbb{R}^3 \times \mathbb{T}^4$  using univariate marginal and pairwise bivariate plots.
(a) Plots of the univariate marginal and pairwise bivariate distributions of  $1 \times 10^{5}$  samples from the uniform distribution in  $\mathbb{P} \subset \mathbb{R}^3$ .
(b) Plots of the univariate marginal and pairwise bivariate distributions of  $1 \times 10^{5}$  samples from the uniform distribution in  $\mathbb{T}^4$ .
Figure 19: Visualisation of the uniform distribution in  $\mathbb{P} \subset \mathbb{R}^3 \times \mathbb{T}^4$  using univariate marginal and pairwise bivariate plots.