SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\lim  _ {\gamma \rightarrow 0} \sup  \left\{\| b _ {\gamma} (x) - b (x) \| \mid x \in \mathbb {R} ^ {d}, \| x \| \leq R \right\} = 0, \tag {86}
$$
$$
\lim  _ {\gamma \rightarrow 0} (1 / \gamma) \sup  \left\{\mathbb {P} \left(\left\| \mathbf {X} _ {(k + 1) \gamma} ^ {\gamma} - \mathbf {X} _ {k \gamma} ^ {\gamma} \right\| \geq \varepsilon \mid \mathbf {X} _ {k \gamma} = x\right) | x \in \mathbb {R} ^ {d}, \| x \| \leq R \right\} = 0. \tag {87}
$$
Where we recall that for any  $x\in \mathbb{R}^d$
$$
b (x) = 2 \int_ {\mathbb {R} ^ {p}} \nabla_ {x _ {t _ {0}}} \log p _ {t _ {0}} ([ x _ {t _ {0}}, x _ {t _ {0}} ^ {c} ]) p _ {t | 0} ^ {x} \left(x _ {t _ {0}} ^ {c} \mid x _ {0} ^ {c}\right) \mathrm {d} x _ {t _ {0}} ^ {c}, \quad \Sigma (x) = 4 \operatorname {I d}. \tag {88}
$$
Proof. Let  $R, \varepsilon > 0$  and  $\gamma \in (0,1)$ . Using Assumption 1, there exists  $C > 0$  such that for any  $x_{t_0} \in \mathbb{R}^d$  with  $\| x_{t_0} \| \leq R$ , we have  $\| \nabla_{x_{t_0}} \log p_{t_0}([x_{t_0}, x_{t_0}^c]) \| \leq C (1 + \| x_{t_0}^c \|)$ . Since  $p_{t_0|0}^c$  is Gaussian with zero mean and covariance matrix  $(1 - \mathrm{e}^{-2t_0})$  Id, we get that for any  $p \in \mathbb{N}$ , there exists  $A_k \geq 0$  such that for any  $x_{t_0} \in \mathbb{R}^d$  with  $\| x_{t_0} \| \leq R$
$$
\int_ {\mathbb {R} ^ {d}} \| \nabla_ {x _ {t _ {0}}} \log p _ {t _ {0}} ([ x _ {t _ {0}}, x _ {t _ {0}} ^ {c} ]) \| ^ {p} p _ {t _ {0} | 0} ^ {c} \left(x _ {t _ {0}} ^ {c} \mid x _ {0} ^ {c}\right) \mathrm {d} x _ {t _ {0}} ^ {c} \leq A _ {k} (1 + \| x _ {0} ^ {c} \| ^ {p}). \tag {89}
$$
Therefore, using this result and the fact that for any  $s \geq 0$ ,  $\mathrm{e}^{-s} \geq 1 - s$ , we get that there exists  $B_k \geq 0$  such that for any  $k, p \in \mathbb{N}$  and for any  $x_{t_0} \in \mathbb{R}^d$  with  $\| x_{t_0} \| \leq R$
$$
\mathbb {E} \left[ \| \mathbf {X} _ {(k + 1) \gamma} - \mathbf {X} _ {k \gamma} \| ^ {p} \mid \mathbf {X} _ {k \gamma} = x \right] \leq B _ {k} \gamma^ {p / 2} \left(1 + \| x _ {0} ^ {c} \| ^ {p}\right). \tag {90}
$$
Therefore, combining this result and the Markov inequality, we get that for any  $x_{t_0} \in \mathbb{R}^d$  with  $\| x_{t_0} \| \leq R$  we have
$$
\lim  _ {\gamma \rightarrow 0} (1 / \gamma) \sup  \left\{\mathbb {P} \left(\left\| \mathbf {X} _ {(k + 1) \gamma} ^ {\gamma} - \mathbf {X} _ {k \gamma} ^ {\gamma} \right\| \geq \varepsilon \mid \mathbf {X} _ {k \gamma} = x\right) | x \in \mathbb {R} ^ {d}, \| x \| \leq R \right\} = 0, \tag {91}
$$
$$
\lim  _ {\gamma \rightarrow 0} (1 / \gamma) \| \mathbb {E} [ (\mathbf {X} _ {(k + 1) \gamma} ^ {\gamma} - \mathbf {X} _ {k \gamma} ^ {\gamma}) \mathbf {1} _ {\| \mathbf {X} _ {(k + 1) \gamma} ^ {\gamma} - \mathbf {X} _ {k \gamma} ^ {\gamma}} \| \geq 1 | \mathbf {X} _ {k \gamma} = x ] \| = 0, \tag {92}
$$
$$
\lim  _ {\gamma \rightarrow 0} (1 / \gamma) \| \mathbb {E} [ (\mathbf {X} _ {(k + 1) \gamma} ^ {\gamma} - \mathbf {X} _ {k \gamma} ^ {\gamma}) \mathbf {1} _ {\| \mathbf {X} _ {(k + 1) \gamma} ^ {\gamma} - \mathbf {X} _ {k \gamma} ^ {\gamma}} \| \geq 1 | \mathbf {X} _ {k \gamma} = x ] \| = 0 \tag {93}
$$
In addition, we have that for any  $x_{t_0} \in \mathbb{R}^d$  with  $R > 0$
$$
\begin{array}{l} \left| \left(2 / \gamma\right) \left(1 - \mathrm {e} ^ {- \gamma}\right) - 2 \right| \| \int_ {\mathbb {R} ^ {d}} \nabla_ {x _ {t _ {0}}} \log p _ {t _ {0}} \left(\left[ x _ {t _ {0}}, x _ {t _ {0}} ^ {c} \right]\right) p _ {t _ {0} | 0} \left(x _ {t _ {0}} ^ {c} \mid x _ {0} ^ {c}\right) \mathrm {d} x _ {t _ {0}} ^ {c} \| (94) \\ \leq A _ {1} \left(1 + \| x _ {0} ^ {c} \|\right) (2 / \gamma) \left| \mathrm {e} ^ {- \gamma} - 1 + \gamma \right|. (95) \\ \end{array}
$$
We also have that for any  $x_{t_0} \in \mathbb{R}^d$  with  $R > 0$
$$
\begin{array}{l} (4 / \gamma) | 1 - \mathrm {e} ^ {- \gamma} | ^ {2} \| \int_ {\mathbb {R} ^ {d}} \nabla_ {x _ {t _ {0}}} \log p _ {t _ {0}} \left(\left[ x _ {t _ {0}}, x _ {t _ {0}} ^ {c} \right]\right) ^ {\otimes 2} p _ {t _ {0} | 0} \left(x _ {t _ {0}} ^ {c} \mid x _ {0} ^ {c}\right) \mathrm {d} x _ {t _ {0}} ^ {c} \| (96) \\ \leq A _ {2} \left(1 + \left\| x _ {0} ^ {c} \right\| ^ {2}\right) (4 / \gamma) \left| 1 - \mathrm {e} ^ {- \gamma} \right| ^ {2}. (97) \\ \end{array}
$$
Combining this result, (91), the fact that  $\lim_{\gamma \to 0}(4 / \gamma)|1 - \mathrm{e}^{-\gamma}|^2 = 0$  and  $\lim_{\gamma \to 0}(2 / \gamma)|\mathrm{e}^{-\gamma} - 1 + \gamma| = 0$ , we get that  $\lim_{\gamma \to 0}\sup \{\|\Sigma_{\gamma}(x) - \Sigma(x)\| |x \in \mathbb{R}^d, \|x\| \leq R\} = 0$ . Similarly, using (91), (94) and the fact that  $\lim_{\gamma \to 0}(4 / \gamma)|1 - \mathrm{e}^{-\gamma}|^2 = 0$ , we get that  $\lim_{\gamma \to 0}\sup \{\|b_{\gamma}(x) - b(x)\| |x \in \mathbb{R}^d, \|x\| \leq R\} = 0$ .
We can now conclude the proof of Theorem E.1.
Proof. We have that  $x \mapsto b(x)$  and  $x \mapsto \Sigma(x)$  are continuous. Combining this result and Lemma E.2, we conclude the proof upon applying Stroock and Varadhan (2007, Theorem 11.2.3).
Theorem E.1 is a non-quantitative result which states what is the limit chain for the REPAINT procedure. Note that if we do not resample, we get that
$$
b ^ {\text {c o n d}} (x) = 2 \nabla_ {x _ {t _ {0}}} \log p _ {t _ {0}} \left(\left[ x _ {t _ {0}}, x _ {t _ {0}} ^ {c} \right]\right), \quad \Sigma (x) = 4 \operatorname {I d}. \tag {98}
$$
Recalling (88), we get that (98) is an amortised version of  $b^{\mathrm{cond}}$ . Similar convergence results can be derived in this case. Note that it is also possible to obtain quantitative discretization bounds between  $(\mathbf{X}_t)_{t\geq 0}$  and  $(\mathbf{X}_t^{1 / n})_{t\geq 0}$  under the  $\ell^2$  distance. These bounds are usually leveraged using the Girsanov theorem (Durmus and Moulines, 2017; Dalalyan, 2017). We leave the study of such bounds for future work.
We also remark that  $b(x_{t_0})$  is not given by  $\nabla \log p_{t_0}(x_{t_0}|x_0^c)$ . Denoting  $U_{t_0}$  such that for any  $x_{t_0} \in \mathbb{R}^d$
$$
U _ {t _ {0}} \left(x _ {t _ {0}}\right) = - \int_ {\mathbb {R} ^ {p}} \left(\log p _ {t _ {0}} \left(x _ {t _ {0}} \mid x _ {t _ {0}} ^ {c}\right)\right) p _ {t | 0} \left(x _ {t _ {0}} ^ {c} \mid x _ {0} ^ {c}\right) \mathrm {d} x _ {t _ {0}} ^ {c}, \tag {99}
$$
we have that  $\nabla U_{t_0}(x_{t_0}) = -b(x_{t_0})$ , under mild integration assumptions. In addition, using Jensen's inequality, we have
$$
\int_ {\mathbb {R} ^ {d}} \exp \left[ - U _ {t _ {0}} \left(x _ {t _ {0}}\right) \right] \mathrm {d} x _ {t _ {0}} \leq \int_ {\mathbb {R} ^ {d}} \int_ {\mathbb {R} ^ {p}} p _ {t _ {0}} \left(x _ {t _ {0}} \left| x _ {t _ {0}} ^ {c}\right)\right) p _ {t | 0} \left(x _ {t _ {0}} ^ {c} \mid x _ {0} ^ {c}\right) \mathrm {d} x _ {t _ {0}} \mathrm {d} x _ {t _ {0}} ^ {c} \leq 1. \tag {100}
$$
Hence,  $\pi_{t_0}$  with density proportional to  $x\mapsto \exp [-U_{t_0}(x)]$  defines a valid probability measure.
We make the following assumption which allows us to control the ergodicity of the process  $(\mathbf{X}_t)_{t\geq 0}$ .
Assumption 2. There exist  $\mathfrak{m} > 0$  and  $C\geq 0$  such that for any  $x_{t_0}\in \mathbb{R}^d$  and  $x_{t_0}^c\in \mathbb{R}^p$
$$
\left\langle \nabla_ {x _ {t}} \log p _ {t _ {0}} \left(\left[ x _ {t}, x _ {t} ^ {c} \right]\right), x _ {t} \right\rangle \leq - \mathfrak {m} \| x _ {t} \| ^ {2} + C \left(1 + \| x _ {t} ^ {c} \| ^ {2}\right). \tag {101}
$$
The following proposition ensures the ergodicity of the chain  $(\mathbf{X}_t)_{t\geq 0}$ . It is a direct application of Roberts and Tweedie (1996, Theorem 2.1).
Proposition E.3. Assume Assumption 1 and Assumption 2. Then,  $\pi_{t_0}$  is the unique invariant probability measure of  $(\mathbf{X}_t)_{t\geq 0}$  and  $\lim_{t\to 0}\| \mathcal{L}(\mathbf{X}_t) - \pi_{t_0}\|_{\mathrm{TV}} = 0$  where  $\mathcal{L}(\mathbf{X}_t)$  is the distribution of  $\mathbf{X}_t$ .
Finally, for any  $t_0 > 0$ , denoting  $\pi_{t_0}$  the probability measure with density  $U_{t_0}$  given for any  $x_{t_0} \in \mathbb{R}^d$  by
$$
U _ {t _ {0}} \left(x _ {t _ {0}}\right) = - \int_ {\mathbb {R} ^ {p}} \left(\log p _ {t _ {0}} \left(x _ {t _ {0}} \mid x _ {t _ {0}} ^ {c}\right)\right) p _ {t | 0} \left(x _ {t _ {0}} ^ {c} \mid x _ {0} ^ {c}\right) \mathrm {d} x _ {t _ {0}} ^ {c}. \tag {102}
$$
We show that the family of measures  $(\pi_{t_0})_{t_0 > 0}$  approximates the posterior with density  $x_0 \mapsto p(x_0 | x_0^c)$  when  $t_0$  is small enough.
Proposition E.4. Assume Assumption 1. We have that  $\lim_{t_0\to 0}\pi_{t_0} = \pi_0$  where  $\pi_0$  admits a density w.r.t. the Lebesgue measure given by  $x_0\mapsto p(x_0|x_0^c)$
Proof. This is a direct consequence of the fact that  $p_{t|0}(\cdot |x_0^c) \to \delta_{x_0^c}$ .
This last results show that even though we do not target  $x_{t_0} \mapsto p_{t_0|0}(x_{t_0}|x_0^c)$  using this corrector term, we still target  $p(x_0|x_0^c)$  as  $t_0 \to 0$  which corresponds to the desired output of the algorithm.
F Experimental details
Models, training and evaluation have been implemented in Jax (Bradbury et al., 2018). We used Python (Van Rossum and Drake Jr, 1995) for all programming, Hydra (Yadan, 2019), Numpy (Harris et al., 2020), Scipy (Virtanen et al., 2020), Matplotlib (Hunter, 2007), and Pandas (McKinney et al., 2010). The code is publicly available at https://github.com/cambridge-mlg/neural_diffusion_processes.
F.1 Regression 1d
F.1.1 Data generation
We follow the same experimental setup as Bruinsma et al. (2020) to generate the 1d synthetic data. It consists of Gaussian (Squared Exponential (SE), MATÉRN(5/2), WEAKLY PERIODIC) and non-Gaussian (SAWTOOTH and MIXTURE) sample paths, where MIXTURE is a combination of the other four datasets with equal weight. Fig. 9 shows samples for each of these dataset. The Gaussian datasets are corrupted with observation noise with variance  $\sigma^2 = 0.05^2$ . The left column of Fig. 9 shows example sample paths for each of the 5 datasets.
The training data consists of  $2^{14}$  sample paths while the test dataset has  $2^{12}$  paths. For each test path we sample the number of context points between 1 and 10, the number of target points are fixed to 50 for the GP datasets and 100 for the non-Gaussian datasets. The input range for the training and interpolation datasets is  $[-2, 2]$  for both the context and target sets, while for the extrapolation task the context and target input points are drawn from [2, 6].
Architecture. For all datasets, except SAWTOOTH, we use 5 bi-dimensional attention layers (Dutordoir et al., 2022) with 64 hidden dimensions and 8 output heads. For SAWTOOTH, we obtained better performance with a wider and shallower model consisting of 2 bi-dimensional attention layers with a hidden dimensionality of 128. In all experiment, we train the NDP-based models over 300 epochs using a batch size of 256. Furthermore, we use the Adam optimiser for training with the following learning rate schedule: linear warm-up for 10 epochs followed by a cosine decay until the end of training.
F.1.2 Ablation Limiting Kernels
The test log-likelihoods (TLLs) reported in App. F.1.3 for the NDP models target a white limiting kernel and train to approximate the preconditioned score  $K\nabla \log p_{t}$ . Overall, we found this to be the best performing setting. App. F.1.3 shows an ablation study for different choices of limiting kernel and score parametrisation. We refer to Table 3 for a detailed derivation of the score parametrisations.
The dataset in the top row of the figure originates from a Squared Exponential (SE) GP with lengthscale  $\ell = 0.25$ . We compare the performance of three different limiting kernels: white (blue), a SE with a longer lengthscale  $\ell = 1$  (orange), and a SE with a shorter lengthscale  $\ell = 0.1$  (green). As the dataset is Gaussian, we have access to the true score. We observe that, across the different parameterisations, the white limiting kernel performance best. However, note that for the White kernel  $K = I$  and thus the different parameterisations become identical. For non-white limiting kernels we see a reduction in performance for both the approximate and exact score. We attribute this to the additional complexity of learning a non-diagonal covariance.
In the bottom row of App. F.1.3 we repeat the experiment for a dataset consisting of samples from the Periodic GP with lengthscale 0.5. We draw similar conclusions: the best performing limiting kernel, across the different parametrisations, is the White noise kernel.
F.1.3 Ablation Conditional Sampling
Next, we focus on the empirical performance of the different noisng schemes in the conditional sampling, as discussed in Fig. 8. For this, we measure the Kullback-Leibler (KL) divergence between two Gaussian distributions: the true GP-based conditional distribution, and an distribution created by drawing conditional sampling from the model and fitting a Gaussian to it using the empirical mean and covariance. We perform this test on the 1D squared exponential dataset (described above) as this gives us access to the true posterior. We use  $2^{12}$  samples to estimate the empirical mean and covariance, and fix the number of context points to 3.
In Fig. 11 we keep the total number of score evaluations fixed to 5000 and vary the number of steps in the inner  $(L)$  loop such that the number of outer steps is given by the ratio  $5000 / L$ . From the figure, we observe that the particular choice of noising scheme is of less importance as long at least a couple  $(\pm 5)$  inner steps are taken. We further note that in this experiment we used the true score (available because of the Gaussianity of the dataset), which means that these results may differ if an approximate score network is used.
Figure 9: Visualisation of 1D regression experiment.
Figure 10: Ablation study targeting different limiting kernels and score parametrisations.
Table 5: Mean test log-likelihood (TLL)  $\pm 1$  standard error estimated over 4096 test samples are reported. Statistically significant best non-GP model is in bold. The NP baselines (GNP, ConvCNP, ConvNP and ANP) are quoted from Bruinsma et al. (2020). \* stands for a TLL below -10.
SE	MATÉRN-5/2	WEAKLY PER.	SAWTOOTH	MIXTURE
INTERPOLATION
GP (OPTIMUM)	0.70±0.00	0.31±0.00	-0.32±0.00	n/a	n/a
T(1)-GEOMNDP	0.72±0.03	0.32±0.03	-0.38±0.03	3.39±0.04	0.64±0.08
NDP*	0.71±0.03	0.30±0.03	-0.37±0.03	3.39±0.04	0.64±0.08
GNP	0.70±0.01	0.30±0.01	-0.47±0.01	0.42±0.01	0.10±0.02
CONVCNP	-0.80±0.01	-0.95±0.01	-1.20±0.01	0.55±0.02	-0.93±0.02
CONVNP	-0.46±0.01	-0.67±0.01	-1.02±0.01	1.20±0.01	-0.50±0.02
ANP	-0.61±0.01	-0.75±0.01	-1.19±0.01	0.34±0.01	-0.69±0.02
GENERALISATION
GP (OPTIMUM)	0.70±0.00	0.31±0.00	-0.32±0.00	n/a	n/a
T(1)-GEOMNDP	0.70±0.02	0.31±0.02	-0.38±0.03	3.39±0.03	0.62±0.02
NDP*	*	*	*	*	*
GNP	0.69±0.01	0.30±0.01	-0.47±0.01	0.42±0.01	0.10±0.02
CONVCNP	-0.81±0.01	-0.95±0.01	-1.20±0.01	0.53±0.02	-0.96±0.02
CONVNP	-0.46±0.01	-0.67±0.01	-1.02±0.01	1.19±0.01	-0.53±0.02
ANP	-1.42±0.01	-1.34±0.01	-1.33±0.00	-0.17±0.00	-1.24±0.01
Figure 11: Ablation noising schemes for conditional sampling.
F.2 Gaussian process vector fields
Data We create synthetic datasets using samples from two-dimensional zero-mean GPs with the following E(2)-equivariant kernels: a diagonal Squared-Exponential (SE) kernel, a zero curl (CURL-FREE) kernel and a zero divergence (DIV-FREE) kernel, as described in App. D.1. We set the variance to  $\sigma^2 = 1$  and the lengthscale to  $\ell = \sqrt{5}$ . We evaluate these GPs on a disk grid, created via a 2D grid with  $30 \times 30$  points regularly space on  $[-10, 10]^2$  and keeping only the points inside the disk of radius 10. We create a training dataset of size  $80 \times 10^3$ , and a test dataset of size  $10 \times 10^3$ .
Models We compare two flavours of our model GeomNDP. One with a non-equivariant attention-based score network (Figure C.1, Dutordoir et al., 2022), referred as  $\mathrm{NDP}^*$ . Another one with a  $\mathrm{E}(2)$ -equivariant score architecture, based on steerable CNNs (Thomas et al., 2018; Weiler et al., 2018). We rely on the e3nn library (Geiger and Smidt, 2022) for implementation. A knn graph  $\mathcal{E}$  is built with  $k = 20$ . The pairwise distances are first embedded into  $\mu(r_{ab})$  with a 'smooth_finite' basis of 50 elements via e3nn. soft_one_hot_1inspace, and with a maximum radius of 2. The time is mapped via a sinusoidal embedding  $\phi(t)$  (Vaswani et al., 2017). Then edge features are obtained as  $e_{ab} = \Psi^{(e)}(\mu(r_{ab})||\phi(t)) \forall (a,b) \in \mathcal{E}_k$  with  $\Psi^{(e)}$  an MLP with 2 hidden layers of width 64. We use 5 e3nn. FullyConnectedTensorProduct layers with update given by  $V_a^{k+1} = \sum_{b \in \mathcal{N}(a,\mathcal{E}_k)} V_a^k \otimes (\Psi^v(e_{ab}||V_a^k||V_b^k)) Y(\hat{r}_{ab})$  with  $Y$  spherical harmonics up to order  $2m \Psi^v$  an
MLP with 2 hidden layers of width 64 acting on invariant features, and node features  $V^{k}$  having irreps  $12 \times 0 \times e + 12 \times 0 \times e + 4 \times 1 \times e + 4 \times 1 \times e$ . Each layer has a gate non-linearity (Weiler et al., 2018).
We also evaluate two neural processes, a translation-equivariant CONVCNP (Gordon et al., 2020) with decoder architecture based on 2D convolutional layers (LeCun et al., 1998) and a  $\mathrm{C4} \ltimes \mathbb{R}^2 \subset \mathrm{E}(2)$ -equivariant STEERCNP (Holderrieth et al., 2021) with decoder architecture based on 2D steerable convolutions (Weiler and Cesa, 2021). Specific details can be found in the accompanying codebase https://github.com/PeterHolderrieth/Steerable_CNPs of Holderrieth et al. (2021).
**Optimisation.** Models are trained for  $80k$  iterations, via (Kingma and Ba, 2015) with a learning rate of  $5e - 4$  and a batch size of 32. The neural diffusion processes are trained unconditionally, that is we feed GP samples evaluated on the full disk grid. Their weights are updated via with exponential moving average, with coefficient 0.99. The diffusion coefficient is weighted by  $\beta : t \mapsto \beta_{\min} + (\beta_{\max} - \beta_{\min}) \cdot t$ , and  $\beta_{\min} = 1e - 4$ ,  $\beta_{\max} = 15$ .
As standard, the neural processes are trained by splitting the training batches into a context and evaluation set, similar to when evaluating the models. Models have been trained on A100-SXM-80GB GPUs.
Evaluation. We measure the predictive log-likelihood of the data process samples under the model on a held-out test dataset. The context sets are of size 25 and uniformly sampled from a disk grid of size 648, and the models are evaluated on the complementary of the grid. For neural diffusion processes, we estimate the likelihood by solving the associated probability flow ODE (53). The divergence is estimated with the Hutchinson estimator, with Rademacher noise, and 8 samples, whilst the ODE is solved with the 2nd order Heun solver, with 100 discretisation steps.
We also report the performance of the data-generating GP, and the same GP but with diagonal posterior covariance GP (DIAG.).
F.3 Tropical cyclone trajectory prediction
Data. The data is drawn from the International Best330 Track Archive for Climate Stewardship (IBTrACS) Project, Version 4 (Knapp et al., 2010; Knapp et al., 2018). The tracks are taken from the 'all' dataset covering the tracks from all cyclone basins across the globe. The tracks are logged at intervals of every 3 hours. From the dataset, we selected tracks of at least 50 time points long and clipped any longer to this length, resulting in 5224 cyclones.  $90\%$  was used fro training and  $10\%$  held out for evaluation. This split was changed across seeds. More interesting schemes of variable-length tracks or of interest, but not pursued here in this demonstrative experiment. Natively the track locations live in latitude-longitude coordinates, although it is processed into different forms for different models. The time stamps are processed into the number of days into the cyclone forming and this format is used commonly between all models.
Four models were evaluated.
The GP  $(\mathbb{R}\to \mathbb{R}^2)$  took the raw latitude-longitude data and normalised it. Using a 2-output RBF kernel with no covariance between the latitude and longitude and taking the cyclone time as input, placed a GP over the data. The hyperparameters of this kernel were optimised using a maximum likelihood grid search over the data. Note that this model places density outside the bounding box of  $[-90,90]\times [-180,180]$  that defines the range of latitude and longitude, and so does not place a proper distribution on the space of paths on the sphere.
The STEREOGRAPHIC GP  $(\mathbb{R}\to \mathbb{R}^2 /\{0\})$  instead transformed the data under a stereographic projection centred at the north pole, and used the same GP and optimisation as above. Since this model only places density on a set of measure zero that does not correspond to the sphere, it does induce a proper distribution on the space of paths on the sphere.
The NDP  $(\mathbb{R}\to \mathbb{R}^2)$  uses the same preprocessing as GP  $(\mathbb{R}\rightarrow \mathbb{R}^2)$  but uses a Neural Diffusion Process from (Dutordoir et al., 2022) to model the data. This has the same shortcomings as the GP  $(\mathbb{R}\to \mathbb{R}^2)$  in not placing a proper density on the space of paths on the sphere. The network used for the score function and the optimisation procedure is detailed below. A linear beta schedule was used with  $\beta_0 = 1e - 4$  and  $\beta_{1} = 10$ . The reverse model was integrated back to  $\epsilon = 5e - 4$  for numerical
stability. The reference measure was a white noise kernel with a variance 0.05. ODEs and SDEs were discretised with 1000 steps.
The GEOMNDP  $(\mathbb{R}\to S^2)$  works with the data projected into 3d space on the surface of the sphere. This projection makes no difference to the results of the model, but makes the computation of the manifold functions such as the exp map easier, and makes it easier to define a smooth score function on the sphere. This is done by outputting a vector for the score from the neural network in 3d space, and projecting it onto the tangent space of the sphere at the given point. For the necessity of this, see (De Bortoli et al., 2021). The network used for the score function and the optimisation procedure is detailed below. A linear beta schedule was used with  $\beta_0 = 1e - 4$  and  $\beta_{1} = 15$ . The reverse model was integrated back to  $\epsilon = 5e - 4$  for numerical stability. The reference measure was a white noise kernel with a variance 0.05. ODEs and SDEs were discretised with 1000 steps.
Neural network. The network used to learn the score function for both NDP  $(\mathbb{R}\to \mathbb{R}^2)$  and GEOMNDP  $(\mathbb{R}\rightarrow S^2)$  is a bi-attention network from Dutordoir et al. (2022) with 5 layers, hidden size of 128 and 4 heads per layer. This results in 924k parameters.
**Optimisation.** NDP  $(\mathbb{R} \to \mathbb{R}^2)$  and GEOMNDP  $(\mathbb{R} \to S^2)$  were both optimised using (correctly implemented) Adam for 250k steps using a batch size of 1024 and global norm clipping of 1. Batches were drawn from the shuffled data and refreshed each time the dataset was exhausted. A learning rate schedule was used with 1000 warmup steps linearly from 1e-5 to 1e-3, and from there a cosine schedule decaying from 1e-3 to 1e-5. With even probability either the whole cyclone track was used in the batch, or 20 random points were sub-sampled to train the model better for the conditional sampling task.
Conditional sampling. The GP models used closed-form conditional sampling as described. Both diffusion-based models used the Langevin sampling scheme described in this work. 1000 outer steps were used with 25 inner steps. We use a  $\psi = 1.0$  and  $\lambda_0 = 2.5$ . In addition at the end of the Langevin sampling, we run an additional 150 Langevin steps with  $t = \epsilon$  as this visually improved performance.