Geometric Neural Diffusion Processes
Authors: Emile Mathieu
Vincent Dutordoir
Michael John Hutchinson
Valentin De Bortoli
Yee Whye Teh
Richard E Turner
Venue: NeurIPS 2023 poster (2023)
Keywords: diffusion model
functional space
stochastic process
time-series
neural processes
Gaussian processes
random fields
invariance
equivariance
symmetries
stationarity
TL;DR: We extend diffusion models to functional data with arbitrary codomain and incorporating symmetry

Abstract: Denoising diffusion models have proven to be a flexible and effective paradigm for generative modelling.
Their recent extension to infinite dimensional Euclidean spaces has allowed for the modelling of stochastic processes.
However, many problems in the natural sciences incorporate symmetries and involve data living in non-Euclidean spaces.
In this work, we extend the framework of diffusion models to incorporate a series of geometric priors in infinite-dimension modelling.
We do so by a) constructing a noising process which admits, as limiting distribution, a geometric Gaussian process that transforms under the symmetry group of interest, and b) approximating the score with a neural network that is equivariant w.r.t. this group.
We show that with these conditions, the generative functional model admits the same symmetry.
We demonstrate scalability and capacity of the model, using a novel Langevin-based conditional sampler, to fit complex scalar and vector fields, with Euclidean and spherical codomain, on synthetic and real-world weather data.

Geometric Neural Diffusion Processes
Emile Mathieu University Of Cambridge
Vincent Dutordoir University Of Cambridge
Michael J. Hutchinson
University Of Oxford
Secondmind Labs
Valentin De Bortoli Center for Science of Data, ENS Ulm
Yee Whye Teh University Of Oxford
Richard E. Turner University Of Cambridge, Microsoft Research
Abstract
Denoising diffusion models have proven be flexible and effective paradigm for generative modelling. Their recent extension to infinite dimensional Euclidean spaces has allowed for the modelling of stochastic processes. However, many problems in the natural sciences incorporate symmetries and involve data living non Euclidean spaces. In this work, extend the framework of diffusion models modelling. We do
to incorporate a series of geometric priors in by a) constructing noising process which admits, as limiting distribution, geometric Gaussian process that transforms under the symmetry group of interest, and b) approximating the score with neural network that is equivariant w.r.t. this group. We show that with these conditions, the generative functional model admits the same symmetry. We demonstrate scalability and capacity of the model, using a novel Langevin-basse conditional sampler, to fit complex scalar and vector fields, with Euclidean and spherical codomain, on synthetic and real-world weather data.
Introduction
Traditional denoising diffusion models are defined on finite-dimension Euclidean spaces (Song and Ermon, 2019; Song et al., 2021; Ho al., 2020; Dhariwal and Nichol, 2021). Extensions have recently been developed for more exotic distributions. such as those supported on Riemannian manifolds (De Bortoli et al., 2022; Huang et al., 2022: Lou and Ermon, 2023; Fishman al., 2023), and on function spaces of the form f : IR" + Fd (Dutordoir et al., 2022: Kerrigan et al., 2022; Lim al., 2023a; Pidstrigach al., 2023; Franzese et al., 2023; Bond-Taylor and Willcocks, 2023) (i.e. stochastic processes). In this work, we extend diffusion models to further deal with distributions over functions that incorporate oon-Euelidean geometry in two different ways. This investigation of geometry also naturally leads to the consideration of symmetries in these distributions, and as such we also present methods for incorporating these into diffusion models.
Firstly, we look at tensor fields. Tensor fields are geometric objects that assign to all points on some manifold value that lives in some vector space V. Roughly speaking, these functions of the form f .M- + V. These objects are central to the study of physics as they form a generic mathematical framework for modelling natural phenomena. Common examples include the pressure of a fluid in motion as f I + IR, representing wind over the Earth's 's surface as f S + TS? where TS² is the tangent-space of the sphere, or modelling the stress deformed object as f: Object - TR' & 3 where & is the tensor-product of the tangent spaces. Given the inherent symmetry in the laws of nature, these tensor fields can transform symmetries. Any modelling of these laws may benefit from respecting these symmetries.
way that preserves these
37th Conference on Neural Information Processing Systems (NeurIPS 2023).
"Equal contribution

---

Secondly, we look at fields with manifold codomain, and in particular, at functions of the form f: IR M. The challenge in dealing with sanifold-valued output, arises from the lack of vector- space structure. In applications, these functions typically appear when modelling processes indexed by time that take values on a manifold. Examples include tracking the eye of eyclones moving on the surface of the Earth, or modelling the joint angles of robot as it performs tasks.
The lack of data or noisy measurements in the physical process of interest motivates probabilistic treatment of such phenomena, in addition to its functional nature. Arguably the most important framework for modelling stochastic processes are Gaussian Processes (GPs) (Rasmussen, 2003), as they allow for exact or approximate posterior prediction (Titsias, 2009; Rahimi and Recht, 2007; Wilson et al., 2020). In particular, when choosing equivariant mean and kernel functions, GPs are invariant (i.e. stationary) (Holderrieth et al., 2021; Azangulov et al.. 2022; Azangulov et al., 2023). Their limited modelling capacity and the difficulty in designing complex, problem-speciffe kernels motivates the development of neural processes (NPs) (Garnelo et al., 2018b), which learn to approximately model a conditional stochastic process directly from data. NPs have been extended to model translation invariant (scalar) processes (Gordon et al., 2020) and more generic E(n)- -invarirn processes (Holderrieth et al., 2021). Yet, the Gaussian conditional assumption of standard NPs still limits their flexibility and prevents such models from fitting complex processes. Diffusion models provide a compelling alternative for significantly greater modelling flexibility. In this work, we develop geometric diffusion neural processes which incorporate geometrical prior knowledge into functional diffusion models.
Our contributions are three-fold: (a) We extend diffusion models to more generic function spaces (i.e. tensor fields, and functions f X - MM) by defining a suitable noising process. (b) We incorporate group invariance of the distribution of the generative model by enforcing the covariance kernel and the score network to be group equivariant. (c) We propose a novel Langevin dynamics scheme for efficient conditional sampling.
Background
Denoising diffusion models. We briefly recall here the key concepts behind diffusion models on " and refer the readers to (Song et al., 2021) for more detailed introduction. We consider forward noising process (Yr)ı20 defined by the following Stochastic Differential Equation (SDE)
dY, - - Ydt + dB,, Yo Po:
where (Bz)120 is a d-dimensional Brownian motion and pa the data distribution. The process (Y,)22 is simply
Ornstein-Ulhenbed (OU) process which converges with geometric rate to ((,, Id) (Durmus and Moulines, 2017). Under mild conditions on po, the time-reversed process (Y,)220 = (Yr (YT-:)eejo.r) also satisfies an SDE (Cattiaux et al., 2021) given by
dY: + Y, + V (Y))r + dB,, YO~
where Pt denotes the density of Y. Unfortunately we cannot sample exactly from (2) as pr and the scores (V log are unavailable. First, Pr is substituted with the limiting distribution ((0, Id) as it converges towards it. Second, one can easily show (Kallenberg, 2021, Thm 5.1) that the score V log P: is the minimiser of (s) = E[[ls(Y,
Vy, log Y, |Y))l over functions s [0, T x Rd + Rd where the expectation over the joint distribution of Yo. Y,, and as such can readily be approximated by a neural network so (4, vt) by minimising this functional. Finally, discretisation of (2) is performed (e.g. MaruraMa to obtain approximate samples of po-
Steerable fields. We now define steerable feature fields which are collections of tensor fields. We focus on the Euclidean group G = E(d), which elements g admits a unique decomposition g wh where h E ((d) is d x d orthogonal matrix and u E ((d) a translation which can be identified as an element of IR for a vector E Rd. g I = hz + u denotes the action of y on x, with acting from the left on x by matrix multiplication. This special case simplifies the presentation, but can be extended to the general case is discussed in App. C.I.
f(x)
flg 1x)
p(h)f(g x))
Figure 1: Illustration of vector field f f: R² IR² with representation p(h) h being steered by group element h 90 é 0(2) C E(2).

---

We are interested in learning probabilistic model over functions of the form f: X + R' such that
a group G acts on X and Rd. We call feature field a tuple (f. with f : + R mapping
between input X to some feature f(x) with associated representation p: G + (L()))) (Scott and Serre, 1996). This feature field is said to be G-steerable if it is transformed for all x E g E G
as g f(x) (g x). In this setting, the action of E(d) = T(d) O(d) on the feature field yields g f(z) = (uh) f(x) a p(h),f (h (x - w)). Typical examples of feature fields include scalar fields with Puriv (g) _ I transforming as g f(x) f (g-'x) such as temperature fields, and vectors or potential fields with pra(g) _ h transforming as g (xx hf (g-1 illustrated in Fig. 1, such as wind or force fields.
For many natural phenomena. a priori we do not want to express preference for a particular conformation of the feature field and thus want a prior p to place the same density on all the transformed fields u(g f) n(),, Vy E G. Leveraging this symmetry can drastically reduce the amount of data required to learn from and reduce training time.
3 Geometric neural diffusion processes: GEOMNDPS
3.1 Continuous diffusion on function spaces
We construct a diffusion model on functions f X' - y/, with ) Rd, by defining a diffusion model for every finite set of marginals. Most prior works on infmiit-diimessionll diffusiions consider noising process on the space of functions (Kerrigan et al., 2022; Pidstrigach et al., 2023; Lim et al., 2023b). In theory, this allows the model to define a consistent distribution over all the finite marginals of the process being modelled. In practice, however, only finite marginals can be modelled computer and the score function needs to be approximated, and at this step lose consistency over the marginals. The only work to stay fully consistent in implementation is Phillips et al. (2022), at the cost of limiting functions that can be modelled to a fiite-dimensional subspace. With this in mind, eschew the technically laborious process of defining diffusions over the infinite-dimension space and work solely on the finite marginals following Dutordoir et al. (2022). We find that in practice consistency can be well learned from data see Sec. 5, and this allows for more flexible choices of score network architecture and easier training.
Noising process. We assume we are given data process (Yo(z))oex.. Given any (x),...,, x") E X", consider the following forward noising process (Y,(z))120 و (Y,(+)),.... Ye(x"))zo = (Y),.... Y2) E y defined by the following SDE
dY((x)= = (m(x) - Y,(x)) Bidt + (//"K(x, 1/BBB,,,
(3)
where K(x,)),, A(x²,²) with k :XX' x X + R a kernel and m : X + y. The process (Y)(z))e20 multivariate Ornstein-Uhlenbee
process-with with drift b(t, Y((z)) = m(x) converges with geometric rate to N(m(x), K(x,,)) Using Phillips et al. (2022), it can be shown that this convergence extends to the process (Y))20 which converges to the Gaussian Process with mean m and kernel k. denoted Y...
(x(x) and diffusion coefficient at x. Yx))) = ((z,
In the specific instance where k(x, x²) (x'), then the limiting process Yoe is simply Gaussian white noise, whilst other choices such as the squared-exponential or Matérn kernel would lead to the associated Gaussian limiting process Y Note that the while noise setting is not covered by the existing theory of functional diffusion models, as a Hilbert space and square integral kernel are required, see Kerrigan et al. (2022) for instance.
Denoising process. Under mild conditions over Yo², the time-reversal process (Y,(z)) satisfies
dY, (x) = (-
logpr (V)() 33 K x:) 1/2BB, (4)
with Y, GP k:) and pr the density of Y, w.r.t. Lebesgue. In practice, the log Pr- term- known as the Stein score, not tractable and must be approximated by a neural network. We then consider the generative stochastic process model defined by first sampling Y, GP k)) and then simulating the reverse diffusion (4) (e.g. via Euler- Maruyaaa discretisation..
Haausmaan and Pardoux (1986) assumes Lipschitz continuity of the drift and volatility matrix and Cattiaux al. (2021, Thm 4.9) which only assume finite entropy condition on the space of processes.
3

---

Manifold valued outputs. So far we have defined our generative model with y - R we can readily extend the methodology to manifold-valued functional models using Riemannian diffusion models such as De Bortoli et al. (2022) and Huang et al. (2022), see App. C. One of the main notable difference is that in the case where y is compact manifold, we replace the Ornstein-Uhlenbeck process by Brownian motion which targets the uniform distribution.
Training.
As the reverse SDE (4) involves the preconditioned score K(2, V log Pt. we directly approximate it with neural network (Ks)a [0, 7' x X" x ) + Ty" where TV is the tangent bundle of V,see App. C. The conditional score of the noising process (3) is given by
Yv, log Pe Pe (Y(()) Yo(r)) -- ( (Y(()) - majo) = - (x, 2) -)12,,
(5)
since Y, Mejo + 20
1// with ~ N(O, Id), and Enjo o K with dyo = (1 - e- So (12)
see App. B.1. We learn the preconditioned score (Ks) by minimising the following denoising score matching (DSM) loss (Vincent et al., 2010) weighted by A()) 0/o KTK
[(0; A())) [[ Y.)
(o) [ || (Ks)o(1, Y.) K1/2 [
(6)
where |==13 = T Az. Note that when targeting vart-nar white noise, then K = Id and the loss (6) reverts to the DSM loss with weighting A(t) = 1/o (Song et al., 2021). In App. B.2, we explore several preconditioning terms and associated weighting A(I). Overall, we found the preconditioned score KV log Pt parameterisation, in combination with the lz loss, to perform best, as shown by the ablation study in App. F.1.3.
3.2 Invariant neural diffusion processes
In this section, we show how we can incorporate geometrical constraints into the functional diffusion model introduced in the previous Sec. 3.1. In particular, given group G, we aim to build a generative model over stcerable feature fields as defined Sec. 2.
Invariant process. A stochastic process f ~ μ is said to be G-invariant if pاg A) = u((A) for any g E G. with μ E ((C(X)))) where P the space of probability measure on the space of continuous functions and A C C(A' measurable. From a sample perspective, this means that with input- output pairs C = [(x², y² }} and denoting the action of G on this set g-C - {(gxx' p(g)y' )> f f~ ja G-iinvvrinn
and only if g C has the same distribution as C. In what follows, we aim to derive sufficient conditions the model introduced in Sec. so that it satisfies this G-invariance property. First, we recall such a necessary and sufficient condition for Gaussian processes.
Proposition 3.1. Invariant (stationary) Gaussian process (Holderrieth et al., 2021). We have that Gaussian process (m k) is G-invariant if and only if its mean m. and covariance is are suitably G- G-equivariant that is, for all x, x' E Xx E EG
míg x) = p(g)m(() and kig- x. = p(9)
Trivial examples of Equuiaqquvariiant kernels include diagonal kernels k = ki Id with kn invari- ant (Holderrieth et al., 202 1), but see App. F.2 for non trivial instances introduced by Macêdo and Castro (2010). Building on Prop. 3.1, we then state that our introduced neural diffusion process is also invariant we additionally assume the score network to be G-equivariant.
Proposition 3.2. Invariant neural diffusion process (Yim et al., 2023, Prop 3.6). We denote by (Y (Yx(=))eexxreeo.) the process induced by the time-reversa SDE (4) where the score is approximated by a score network [0,T] X X" x y + Ty", and the limiting process is given by L(Yo) GP k). Assuming m and k are respectively G-equivariant per Prop. 3.1, if we additionally have that the score network is G-equivariant vector field, i.e. salt, g x, v.,(()))) p(g)sa x, v) for all x E X.g g E G. then for any t E [0, T) the process (Y,(z))zEX is G-imvariant.
This result can be proved in two ways, from the probability flow ODE perspective directly in terms of SDE via Fokker- Planck, see App. D.2. In particular, when modelling an invariant sealar data process (Yo(x))rEx such temperature field, we need the score network admit the invariance constraint So g x, x, = so(1, x., y).
Equivariant conditional process. Often precedence is given to modelling the predictive process

---

f(
f(x") 1g C
f(x")JC
f(x 119 c
Figure 2: Samples from equivariant neural diffusson processes conditioned context red) and evaluated on regular grid x" for scalar (Left) 2D vector (Right) fields. Same model is then conditioned transformed context C. with group element being a translation length 2 (Left) or 90" rotation (Right).
given a set of observations = {(x 3/) leec. In this context, the conditional process (Pollard, 2002, p.117) inherits the symmetry of the prior process in the following sense. A stochastic process with distribution y given a context C is said to be conditionally G E--quivariant the conditional satisfies ((Alg C) p(g A|C) for any g E G and A C(X,3) measurable, illustrated in Fig. 2.
Assume a stochastic process I ~ is G invariant. Then the conditional process f C given a of observations C is G-equivariant.
Proposition 3.3. Equivariant conditional process.
Originally stated Holderrieth et al. (2021) the case where the process over functions of the form f R" + IRd and marginals with density w.r.t. Lebesgue, we prove Prop. 3.3 for stochastic processes over generic fields on manifolds in terms only of the measure of the process (App. D.3).
3.3 Conditional sampling
There exist several methods to perform conditional sam- pling in diffusion models such as: replacement sampling, amortisation and conditional guidance, which we dis- cuss in App. E.1. Here we propose a new method for sampling from exact conditional distributions of NDPs using only the score network for the joint distribution. Using the fact that the conditional score can be writ- ten as V, log ((ply) = Va logp(z: p(x, y) - Va log p(y) = V, loggoo p(x, y) can therefore, for any point in the dif- fusion time, conditionally sample using Langevin dynam- ics, following the SDE dY, JKV PY
((Y,) pr
P2y
L(Y
Po
Pr
C Y
C(Y.) L(Y!)
Figure 3: Illustration of Langevin cor- rected conditional sampling. The black line represents the noising process dynamics (pr)ee The time reversal (i.c. predictor) step, is combined with a Langevin corrector step projecting back onto the dynamics.
vKdB,, by only applying the diffusion to the variables of interest and holding the others fixed. While we could sample directly the end time this proves difficult in practice. Similar to the motivation of Song and Ermon (2019), we sample along the reverse diffusion, taking a number of conditional Langevin steps at each time. addition, we apply the forward noising SDE to the conditioning points at each step, as this puts the combined context and sampling set in a region that the score function will be well learned in training. Our procedure is illustrated in Alg. 1. In App. E.1 we draw links with REPAINT of Lugmayr et al. (2022).
3.4 Likelihood evaluation
Similarly to Song et (2021), we can derive deterministic ODE which has the same marginal density as the SDE (3), which is given by the 'probability flow' Ordinary Differential Equation (ODE), see App. B. Once the score network learnt, we can thus use in conjunction with an ODE solver to compute the likelihood of the model. A perhaps more interesting task is to evaluate the predictive posterior likelihood p(y" [x.", (x<,y" "e)))))) given a context set {x yy" Jicc simple approach is to simply rely on the conditional probability rule evaluate p(y" lx",, (x* "",""" (iee)) p(y", (vr heela"" {x }ec)) (ccc)/((( kecc (x*)<<))))) ).c)). This can be done by solving two probability flow ODEs: one over the joint evaluation and context set, and another only over the context set.
5

---

4 Related work
Gaussian processes and the neural processes family. One important and powerful framework to construct distributions over functional spaces are Gaussian processes (Rasmussen, 2003). Yet, they are restricted in their modelling capacity and when using exact inference they scale poorly with the number of datapoints. These problems can be partially alleviated by using neural processes (Kim et al., 2019; Garnelo et al., 2018b; Garnelo et al., 2018a; Jha et al., 2022; Louizos et al., 2019; Singh et al., 2019), although they also assume a Gaussian likelihood. Recently introduced autoregressive NPs (Bruinsma et al., 2023) alleviate this limitation, but they disadvantaged by the fact that variables early in the auto-regressive generation only have simple distributions (typically Gaussian). Finally, (Dupont al., 2022) model weights of implicit neural representation using diffusion models.
Stationary stochastic processes. The most popular Gaussian process kernels (e.g. squared expo- nential, Matém) are stationary, that is, they are translation invariant. These lead to invariant Gaussian processes, whose samples when translated have the same distribution as the original ones. This idea can be extended to the entire isometry group of Euclidean spaces (Holderritth et al., 2021), allowing for modelling higher order tensor fields, such wind fields or incompressible fluid velocity (Macêdo and Castro, 2010). Later, Azangulov et al. (2022) and Azangulov et al. (2023) extended stationary kernels and Gaussian processes to large class of -adEEucideenn spaces, in partieular all compact spaces, and symmetric non compact spaces. In the context of neural processes, (Gordon et al., 2020) introduced CONVCNP so as to encode translation equivariance into the predictive process. They do so by embedding the context into a translation equivariant functional representation which is then decoded with convolutional neural network. Holderrieth et al. (2021) later extended this idea to construct neural processes that are additionally equivariant w.r.t. rotations or subgroup thereof.
Spatial structure in diffusion models. A variety of approaches have also been proposed to incorporate spatial correlation in the noising process of dinte-idimenioonll diffusion models leveraging the multiscale structure of data (Jing et al., 2022; Guth et al., 2022; Ho et al., 2022: Saharia et al., 2021; Hoogeboom and Salimans, 2022; Rissanen et al., 2022). Our methodology a principled way to modify the forward dynamics in classical denoising diffusion models. Hence, our contribution can be understood in the light of recent advances in generative modelling soft and cold denoising diffusion models (Daras et al., 2022; Bansal et al., 2022; Hoogeboom and Salimans, 2022). Several recent work explicitly introduced covariance matrix in the Gaussian noise, either on a choice of kernel (Biloš et al., 2022), based on Discrete Fourier Transform of images (Voleti et al., 2022), or via empirical second order statistics (squared pairwise distances and the squared radius of gyration) for protein modelling (Ingraham et al., 2022). Alternatively, (Guth et al., 2022) introduced correlation on images leveraging a wavelet basis.
also be seen as
Functional diffusion models. Infinite dimensional diffusion models have been investigated in the Euclidean setting in (Kerrigan et al., 2022; Pidstrigach et al., 2023; Lim et al., 2023b; Bond-Taylor and Willcocks, 2023; Hagemann et al., 2023; Franzese al., 2023; Dutordoir al., 2022; Phillips et al., 2022). Most of these works are based on an extension of the diffusion models techniques (Song al., 2021; Ho et al., 2020) the infiite-dimensional space, leveraging tools from the Cameron-Martin theory such as the Feldman-Häjek theorem (Kerrigan et al., 2022; Pidstrigach al., 2023) define
Gaussian measures and how they interact. We refer to (Da Prato and Zabczyk. 2014) for thorough introduction to Stochastic Differrntial Equations in infinite dimension. (Phillips al., 2022) consider another approach by defining countable diffusion processes in a basis. All these approaches amount to learn a diffusion model with spatial structure. Note that this induced correlation necessary for the theory of infinite dimensional SDE (Da Prato and Zabezyk, 2014) to be applied but is not necessary to implement diffusion models (Dutordoir et al., 2022). Several approaches have been considered for conditional sampling. (Pidstrigach et al., 2023; Bond-Taylor and Willcocks, 2023) modify the reverse diffusion to introduce a guidance term, while (Dutordoir et al., 2022; Kerrigan et al., 2022) use the replacement method, Finally (Phillips et al., 2022) amortise the score function w.r.t. the conditioning context.

---

Data Prior
Model Prior
Model Posterior
2.5
SE
True Posterior
2.5
Sawtooth
0.5
0.0
0.5
2
Figure 4: Prior and posterior samples (in blue) from the data process and the GeomNDP model, with context points in red and posterior mean in black.
Experimental results
5.1 ID regression over stationary scalar fields
We evaluate GEOMNDPs several synthetic 1D regression datasets. We follow the same experi- mental setup as Bruinsma et al. (2020) which we detail in App. F.1. In short, it contains Gaussian (Squared Exponential (SE), MATÉRN( WEAKLY PERIODIC) and Ganssaan (SAWTOOTH and MIXTURE) sample paths, where MIXTURE is combination of the other four datasets with equal weight. Fig. 9 shows samples for each of these dataset. The Gaussian datasets corrupted with observation noise with variance x² = 0.05². Table 1 reports the average log- log-likelihood p(y" |=y">>" C) across 4096 test samples, where the context set size uniformly sampled between 1 and 10 and the target has fixed size of 50. All inputs z", xx" are chosen uniformly within their input domain which is [-2, for the training data and 'interpolatton' evaluation and [2. 6] for the "generalisation' evaluation.
We compare the performance of GeomNDP to a GP with the true hyperparameters (when available), (convolutional) Gaussian NP (Bruinsma et al., 2020), convolutional NP (Gordon et al., 2020) and vanilla attention-based NDP (Dutordoir et al., 2022) which we reformulated the continuous diffusion process framework allow for log-likelihood evaluations and thus fair comparison- denoted NDP", We enforce translation invariance in the score network for GeomNDP by subtracting the centre of mass from the imput æ, inducing stationary scalar fields.
On the GP datasets, GNP, CoNvNPs and GeomNDP methods are able to fit the conditionals perfectly-matching the log-likelihood of the GP model. GNP's performance degrades the non- Gaussian datasets as it restricted by its conditional Gaussian assumption, whilst NDPs methods still performs well illustrated ability to generalise outside of the training input range æ E [-2, 21. and evaluate them on translated grid where context and target points are sampled from [2, 6]. Only convolutional NPs (GNP and CONvNP) and G(E)-GEOMNDP are able to model stationary processes and therefore to perform well in the interpolation task. The NDP". the contrary, drastically fails this task.
of Table 1, we assess the models Fig. 4. In the bottom
Table I: Mean test log- log-likeldhood (TLL) (1) + standard error estimated over 4096 test samples are reported. Statistically significant best non-GP model is in bold. '*> stands for TLL below -10. NP baselines from Bruinsma al. (2020). T(1) GCommDGo indicates our proposed method with translation invariant score.
SE
MATÉRN(?)
WEAKLY PER.
SAWTOOTH
MIXTURE
INTERPOLAT.
GP (OPTIMUM)
0.70
0.31
-0.3240.00
T(1)- GEOMNDP
0.72:Lo.03
0.32
-0.88
3.39
0.64
NDP*
0.71
0.30-Fo.os
-0.37-bo.03
3.39-to.04
0.64-1o.08
GNP
0.70=Fo.01
0.30+Fo.0
-0.17-b0.01
0.124o.01
0.10:bo.02
CONVNP
0.46-to.01
0.67
-1.02+00.01
1.20
0.50
GENERALISAT.
GP (OPTIMUM)
0.70
0.31
-0.3240.00
T((1) GEOMNDP
0.70
0.31
0.38
3.39
0.621o..02
NDP*
GNP
0.69
0.30+o.m
-0.47+0.00
0.42+o.m
0.10
CONVNP
0.46-to.01
-0.67
-1.02+o..0
1.19+o.00
-0.53

---

MODEL
SE
CURL-FREE
DIV-FREE
Log-likelihood
GP
0.56
0.66
0.66+4.00
0.6-
NDP
0.55+0.00
0.62+0.01
0.62+00.1
((2)-GEOMND)
0.5610.01
0.65
0.66+0.01
GP
NDP B(2)-GeomNDP
GP (DIAG.)
-1.5620.00
-1.47
-0.07
0.5
T(2)-CONV 10.01
-1.77=0.01
-1.76.0.00
( STEERCCE -1.61-0.00
-1.57=0.00
-1.57=0.01
10a
10a
10a
10a
10
Number of training samples
Figure 5: Quantitative results for experiments on GP vector fields. Mean predictive log rigliield (1) and confidence interval estimated over 5 random seeds. Left: Comparison with neural processes. Statistically significant results are in bold. Right: Ahlation study when varying the number of training data samples.
Non white kernels for limiting process. The NDP methods in the above experiment target the white kernel (xx x²) in the limiting process. In App. F.1.3, we explore different choices for the limiting kernel, such as SE and periodic kernels with short and long lengthscales, along with several score parameterisations, see App. B.3 for description of these. We observe that although choosing such kernels gives head start to the training, eventually yield slightly worse performance. We attribute this to the additional complexity of learning non-diagonal covariance. Finally, across all datasets and limiting kernels, we found the preconditioned score KV log to result in the best performance.
Conditional sampling ablation. We employ the SE dataset to investigate various configurations of the conditional sampler as we have access to the ground truth conditional distribution through the GP posterior. In Fig. 11 compute the Kullback-Leible divergence between the samples generated by GEOMNDP and the actual conditional distribution across different conditional sampling settings. Our results demonstrate the importance of performing multiple Langevin dynamics steps during the conditional sampling process. Additionally, we observe that the choice of noising scheme for the context values ye has relatively less impact on the overall outcome.
5.2 Regression Gaussian process vector fields
We now focus attention to modelling equivariant vector fields. For this, create datasets using samples from a two-dimensional zero-mean GP with one of the following (equ-vruuva kernels: diagonal Squared-Exponeniaal (SE) kernel, (DIV-FREE) kernel, as described in App. D.1.
curl (CURL-FREE) kernel and zero divergence
We equip our model, GeomNDP, with E(?)-equivariint score architecture, based on steerable CNNs (Thomas et al., 2018; Weiler and Cesa, 2021). We compare to NDP" with non-equivariant attention-based network (Dutordoir et al., 2022). We also evaluate two neural processes, a translation- cquivariant CONvCNI (Gordon et al., 2020) and a C4 R² C (())-equivvriant S STEERCNE (Holderricth et al., 2021). We also report the performance of the data-generating GP, and the same GP but with diagonal posterior covariance GP (DIAG.). We measure the predictive log- logllikelihood of the data process samples under the model on held-out test dataset.
We observe in Fig. 5 (Left), that the CNPs performance limited by their diagonal predictive covariance assumption, and as such cannot do better than the GP (DIAG.). We also see that although NDP" able to well GP posteriors, it does not reach the maximum fog-likelihoop value attained by the data GP, in contrast to its equivariant counterpart GFOMNDP To further explore gains brought by the built-in equivariance, we explore the data-efficiency Fig. 5 (Right), and notice that (22)GGFMNDP requires few data samples to fit the data process, since effectively the dimension of the (quotiented) state space is dramatically reduced.
5.3 Global tropical cyclone trajectory prediction
Finally, we assess our model on task where the domain the stochastic process is Euclidean manifold. We model the trajectories of cyclones over the earth, modelled as sample paths of the form + S coming from a stochastic process. The data is drawn from the International Best Track Archive for Climate Stewardship (IBTrACS) Project, Version 4 ((Knapp et al., 2010; Knapp et al., 2018)) and preprocessed as per App. F.3, where details on the implementation of the score function, the ODE/SDE solvers used for the sampling, and baseline methods can be found.
8

---

Figure 6: Left: 1000 samples from the training data. Right: 1000 samples from the trained model.
(a) Interpolation
(b) Extrapolation
Figure 7: Top: Examples of conditional trajectories sampled from the GeomNDP model. Blue: Conditioned sections the trajectory. Green: The actual trajectory of the cyclone. Red: conditional samples from the model. Purple: closest matching trajectories in dataset to conditioning data.
Fig. 6 shows some cyclone trajectories samples from the data process and from a trained GEOMNDP model. We also demonstrate how such trajectories can be interpolated or extrapolated using the conditional sampling method detailed in Sec. 3.3. Such conditional sample paths are shown in Fig. 7. Additionally, we report in Table the likelihood and MSE for series of methods. The interpolation task involves conditioning on the first and last 20% of the cyclone trajectory and predicting intermediary positions. The extrapolation task involves conditioning on the first 40% of trajectories and predicting future positions. We see that the GPs (modelled as f: R - R2, one on latinude/longiuude coordinates, the other via stereographic projection, using a diagonal RBF kernel with hyperparameters fitted with maximum likelihood) fail drastically given the high non- nn-Gaausianity of the data. In the interpolation task, the NDP performs as well as the GEOMNDP but the additional geometric structure of modelling the outputs living on the sphere appears to significantly help for extrapolation. See App. F.3 for more fine-grained results.
6 Discussion
In this work, we have extended diffusion models to model invariant stochastic processes over tensor fields. We did so by (a) constructing continuous noising process over function spaces which correlate input samples with an equivariant kernel, (b) parameterising the score with an equivariant neural network. We have empirically demonstrated the ability of our introduced model GEOMNDP to fit complex stochastic processes, and by encoding the symmetry of the problem hand, we show that it is more data efficient and better able to generalise.
We highlight below some current limitations and important research directions. First, evaluating the model is slow as it relies on costly SDE or ODE solvers, as existing diffusion models. Second, targeting white noise process appears to over-perform other Gaussian processes. In futuree work, we would like to investigate the practical influence of different kernels. Third, strict invariance may sometimes be too strong, we thus suggest softening it by amortising the score network over extra spatial information available from the problem at hand.
Table 2: Comparative results of different models on the cyclone datasct, comparing test set likelihood, interpola- tion likelihood and mean squared error (MSE), and extrapolation likelihood and mean squared error. These are estimated over 5 random seeds. We only report likelihoods of models defined w.r.t the uniform measure on S²
TEST DATA
INTERPOLATION
EXTRAPOLATION
MODEL
Likelihood
Likelihood
MSE (km) Likelihood MSE (km)
GEOMNDP (IR - S²)
802+5 393
535-4
162
536+4
496+14
STEREOGRAPHIC GP (IR + (()())))
26643
2619
245
6587
NDP (R + R²) R2)
166
769
GP (R -
6852 1:11
8138

---

Acknowledgements
We are grateful to Paul Rosa for helping with the proof, and to José Miguel Hernández--Lobtto for use- ful discussions. We thank the hydra (Yadan, 2019), jax (Bradbury al., 2018) and geomstats (Mi- olane et al., 2020) teams, as library is built on these great libraries. Richard E. Turner and Emile Mathicu are supported by an EPSRC Prosperity Partnership EP/T005386/1 between Microsoft Research and the University of Cambridge. Michael J. Hutchinson is supported by the EPSRC Centre for Doctoral Training in Modern Statistics and Statistical Machine Learning (EP/S023151/I).
References
I. Azangulov, Smolensky, Terenin, and Borovitskiy. Stationary Kernels and Gaussian Processes on Lie Groups and their Homogeneous Spaces I: the Compact Case, Aug. 2022. DOI: 48550/arX11 artvv. 2208. 14960. Cited on pages 6.
I. Azangulov, A. Smolensky, A. Terenin, and Borovitskiy. Stationary Kernels and Gaussian Processes on Lie Groups and their Homogeneous Spaces II: non-compact symmetric spaces, Apr. 2023. DOI: 10. 48550855/220 araiv. 2301 13088. Cited on pages
A. Bansal, Borgnia, H.-M. Chu, Li, H. Kazemi, F. Huang, M. Goldblum, J. Geiping, and T. Goldstein. Cold diffusion: inverting arbitrary image transforms without noise. arXiv preprint arXiv:2208.081 09392, 2022. Cited on page
M. Biloš, K. Rasul, A. Schneider, Nevmyvaka, and Günnemann. Modeling Temporal Data as Continuous Functions with Process Diffusion, 2022. URL: https://openrevat net/forum? id=VnJKUypQQ&RR Cited on page 6.
S. Bond-Taylor and C. G. Willcocks. Infinite-diff: infinite resolution diffusion with subsampled mollified states. arXiv preprint arXiv:2303. 288,::22.....222,, 2023. Cited on pages 1, 6.
Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q Zhang. JAX: composable transformations of Python+NumPy programs, version 0.2.5, 2018. URL: http: //gg//tttt///////////// com/ google/ jax. Cited on pages 10, 13.
W. Bruinsma, S. Markou, J. Requeima, A. Y. K. Foong, A. Vaughan, T. Andersson, A. Buonomo, S. Hosking, and R. E. Turner. Autoregressive Conditional Neural Processes. In International Conference on Learning Representations, Feb. 2023. URL: https: //opereeiiee net/forum? id=OAsXFPBTTBh. Cited on page 6.
W. Bruinsma, J. Requeima, A. Y. K. Foong, J. Gordon, and R. E. Turner. Gaussian Neural Processes. In 3rd Symposium on Advances in Approximate Bayesian Inference, Feb. 2020. Cited on pages 14, 16.
P. Cattiaux, G. Conforti, I. Gentil, and C. Léonard. Time reversal of diffusion processes under finite entropy condition. arXiv preprint arXiv:2104.07708, 2021. Cited on pages 5.
S. Chen, S. Chewi, J. Li, Y. Li, A. Salim, and A. R. Zhang. Sampling is easy learning the score: theory for diffusion models with minimal data assumptions. arXiv preprint 2rXiv:2209.12211 2022. Cited on page 11.
T. Cohen. Equivariant convolutional networks. PhD thesis, 2021. Cited page 4.
G. Da Prato and J. Zabezyk. Stochastic equations in infinite dimensions. Cambridge university press, 2014. Cited on page
A. S. Dalalyan. Theoretical guarantees for approximate sampling from smooth and log-concave densities. J. R. Stat. Soc. Ser. B. Stat. Methodol. 79(3):651-676, 2017. ISSN: 1369-7412, DOI: 10. 1111/rssb. 12183.
URL: ht attps:///det org/10, wttpg//...//./....eeeeeeee 1111/rssb. 12183. Cited on page 13.
G. Daras, M. Delbracio, H. Talebi, A. G. Dimakis, and P. Milanfar. Soft diffusion: score matching for general corruptions. arXiv preprint
2022. Cited on page 6.
10

---

V. De Bortoli. Convergence of denoising diffusion models under the manifold hypothesis. arXiv preprint arXiv:2208.05314, 2022. Cited on pages 10, 11.
V. De Bortoli, E. Mathieu, M. Hutchinson, Thortoo,, Y. W. Tch, and A. Doucet. Riemannian score-based generative modeling. arXiv preprint arXiv:2202.02766. 2022. Cited on pages 1, 4.
V. De Bortoli, J. Thortoon, J. Heng, and A. Doucet. Diffusion Schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 2021. Cited on pages 5, 18.
P. Dhariwal and A. Nichol. Diffusion models beat GAN on image synthesis. arXiv preprint AXXi::2055002333 2021. Cited on page
E. Dupont, H. Kim, Eslami, D. Rezende, and D. Rosenbaum. From data functa: your data point is function and you should treat it like one. arXiv preprint arXiv:2201.12200 2022. Cited on page 6.
A. Durmus and E. Moulines. High-dimensional Bayesian inference via the Unadjusted Langevin Algorithm. ArXiv e-prints, May 2016. arXiv: 1605. 01559 (math.ST] Cited on page 4.
A. Durmus and É. Moulines. Nonasymptotic convergence analysis for the unadjusted Langevin algorithm. The Annals of Applied Probability, 27(3) 2017. ISSN: 1050-5164. DOL: 10. 10111121111111 AA-1AAAP... URL: https: //ooi. org/10. 1214/16-AAP1238. Cited pages 2, 13.
V. Dutordoir, A. Saul, Z. Ghahramani, and F. Simpson. Neural Diffusion Processes, June 2022. DOT:
10. 4 arXiv. 2206. 03992. Cited on pages I, 3, 6-8, 14, 16-18.
N. Fishman, L. Klarner, V. D. Bortoli, E. Mathieu, and M. Hutchinson. Diffusion models for con- strained domains, 2023. arXiv: 2304. 053664 [cs.LG]. Cited on page 1.
G. Franzese, S. Rossi, D. Rossi, M. Heinonen, M. Filippone, and P. Michiardi. Continuous-time functional diffusion processes, arXiv preprint arXiv:2303.00806 2023. Cited on pages 1,
M. Garnelo, D. Rosenbaum. C. Maddison, Ramalho, D. Saxton, M. Shanahan, Y. W. Teh, D. Rezende, and S. M. A. Eslami. Conditional neural processes. In J. Dy and A. Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Pro- ceedings of Machine Learning Research, pages 1704-1713. PMLR, July 2018. URL: https: //proceedings. mlr. S.lr.press/v80 htt///onneellle/ html. Cited on page 6.
M. Garnelo, J. Schwarz, Rosenbaum, F. Viola, D. Rezende, Eslami, and W. Teh. Neural processes. arXiv preprini /80.001/ 2018 8. Cited on pages 2, 6.
M. Geiger and T. Smidt. E3nm: euclidean neural networks, 2022. DOI: 10. 48550/ARXIV. 2207. 09453. URL: https: /ttp::////xvvee/e wwwpriiwww.....eeoooee.eee httgg/abs/2201// 09453. Cited on page 16.
Gordon, W. Bruinsma, A. Y. Foong, Requeima, Y. Dubois, and R. E. Turner. Convolutional conditional neural processes. In International Conference on Learning Representations, 2020. URL:
https://openreviex net net/forummid=d oowmrod Swe..eeBss..eeeees..... Cited on pages 2, 6-8, 17.
Guth, S. Coste. V. De Bortoli, and Mallat. Wavelet score-based generative modeling. arXiv preprint wwXiv:2208.05003 2022. Cited on page 6.
P. Hagemann, L. Ruthotto, G. Steidl, and N. T. Yang. Multilevel diffusion: infinite dimensional score-based diffusion models for image generation. arXiv preprint arX1v:2303.04772, 2023. Cited on page 6.
C. R. Harris, K. J. Millman, S. J. van der Walt, R. Gommers, P. Virtanen, D. Cournapeau, E. Wieser, Taylor, S. Berg, N. Smith, R. Kern, M. Pieus, S. Hoyer, M. H. van Kerkwijk, M. Brett, A. Haldane, J. F. del Río, M. Wiebe, P. Peterson, P. Gérard-Marchant, K. Sheppard, T. Reddy, W. Weckesser, Abbasi, C. Gohlke, and T. E. Oliphant. Array programming with NumPy. Nature, 585(7825):357-366 Sept. 2020. Cited on page 13.
U. G. Haussmann and E. Pardoux. Time reversal of diffusions. The Annals of Probability, 14(4):1188- 1205, 1986. Cited on page

---

J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in Neural Informa- tion Processing Systems, 2020. Cited on pages
J. Ho, C. Saharia, W. Chan, D. J. Fleet, M. Norouzi, and T. Salimans. Cascaded diffusion models for high fidelity image generation. J. Mach. Learn. Res., 23:47-1, 2022. Cited on page 6.
P. Holderrieth, M. J. Hutchinson, and Y. W. Tch. Equivariant learning of stochastic fields: gaussian processes and steerable conditional neural processes. In International Conference on Machine Learning, 2021. Cited on pages 2, 4-6, 8, 17.
E. Hoogeboom and T. Salimans. Blurring diffusion models. arXiv preprint arXiv:2209.055577 2022. Cited on page
C.-W. Huang, M. Aghajohari, Bose, P. Panangaden, and A. C. Courville, Riemannian diffusion models. Advances in Neural Information Processing Systems, 35:2750-2761, 2022. Cited on pages 1, 4.
J. D. Hunter. Matplotlib: 2d graphics environment. Computing in Science Engineering, 9(3)990- 95, 2007. Cited on page 13.
M. Hutchinson, A. Terenin, V. Borovitskiy, S. Takao, Y. Teh, and M. Deisenroth. Vector-valued gaussian processes riemannian manifolds via gauge independent projected kernels. Advances in Neural Information Processing Systems, 34:17160-17169 2021. Cited on page 5.
J. Ingraham, M. Baranov, Z. Costello, Frappier, A. Ismail, Tie, W. Wang, V. Xue, F. Obermeyer, A. Beam, and G. Grigoryan. Illuminating protein space with programmable generative model. bioRxiv, 2022. DOT: 10 1 1101/ 12. 01 518682. eprint: https: : /www biorxiv org/ content early / 2022 12 /2 2022. 12. 01 518682 full pdf. URL: https /ww. biorxiv.orge
12. 01. 518682. Cited on page 6.
S. Jha, D. Gong. X. Wang, E. Turner, and L. Yao. The neural process family: survey, applications and perspectives. arXiv preprint arXiv:2209.0m 005.000577 17, 2022. Cited on page 6.
B. Jing, G. Corso, R. Berlinghieri, and T. Jaakkola. Subspace diffusion generative models. arXiv preprint arXiv:2205.01490 2022. Cited on page 6.
O. Kallenberg. Foundations of Modern Probability. Probability Theory and Stochastic Modelling. Springer, 2021. ISBN: ShJezgEACAAJ. Cited on page 2.
URL: hwtps::/t/ooooo google. co uk/booksiid=
T. Karras, M. Aittala, T. Aila, and Laine. Elucidating the Design Space of Diffusion-Basee Generative Models, Oct 2022. DOI: 10. 48550/ arxiv. www....20.................. 2206. 00364. Cited on page 3.
G. Kerrigan, J. Ley, and P. Smyth. Diffusion Generative Models in Infinite Dimensions, 2022, DOI: 10. 00.48550/arrtt art.v. 2212.22222222222 00886. Cited on pages 1, 3, 6.
H. Kim, Mnih, J. Schwarz, M. Garnel,, A. Eslami, D. Rosenbaum, O. Vinyals, and W. Teh. Attentive neural processes. In International Conference on Learning Representations, 2019, URL.: attps://openreviee/ net/ www.rrdddmmdorrooooooo J9KX. Cited on page 6.
D. Kingma and J. Ba. Adam: Method for Stochastic Optimization. urXiv: [es), 2015.
URL: httt:///arxia org/ abs/ /112/ 6980. Cited on page 17.
H. J. Knapp Kenneth R. Diamond, J. P. Kossin, M. C. Kruk, and C. J. Schreck. International Best Track Archive for Climate Stewardship (IBTrACS) Project, Version 4. Technical report, NOAA National Centers for Environmental Information, 2018. DOT: https : / //ooi. org/10 /0// 25921/ 92ty-9e16. Cited on pages 8, 17.
K. R. Knapp, M. C. Kruk, best track archive for climate stewardship (ibtracs): unifying tropical cyclone data. Bulletin of the American Meteorological Society, (1(3):363 2010. DOI: https: //ooi. org/ 10. 1175/ 2009BAMS2755. 1. Cited pages 8, 17.
Levinson, J. Diamond, and C. Neumann. The international
12

---

J. Köhler, L. Klein, and F. Noé. Equivariant Flows: exact likelihood generative learning for symmetric densities, arXiv:2006.0244 02425 [physics, stat], June 2020. URL: http: httt::///axiv.. org/ abs /206. 02425. Cited on page 6.
Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied to document recognition. In Proceedings of the IEEE, volume 86 of number 1, pages 2278 -2224, 1998. URL: http: citeseerx. ist. psu. edu/ve hot/vieddo///////o summary' ? doi = 10. 1 1. 42. 7665. Cited on page 17.
H. Lee, Lu, and Tan. Convergence of based generative modeling for general data distribu- tions. In International Conference on Algorithmic Learning Theory, pages 946-985. PMLR, 2023. Cited on page 11.
H. Lim, N. B. Kovachki, Baptista, C. Beckham, K. Azizzadenesheli, Kossaifi, V. Voleti, J. Song, K. Kreis, J. Kautz, C. Pal, A. Vahdat, and A. Anandkumar. Score-based diffusion models in function space, 2023. DOI: 10. 48550/ARXIV 2302. 07400. Cited on page 1.
H. Lim, N. B. Kovachki, R. Baptista, C. Beckham, K. Azizzadenesheli, Kossaifi, V. Voleti, J. Song, K. Kreis, J. Kautz, C. Pal, A. Vahdat, and A. Anandkumar. Score-based diffusion models in function space, 2023, Cited on pages
A. Lou and S. Ermon. Reflected diffusion models. arXiv preprint arXiv:2304.04744 2023. Cited page 1.
C. Louizos, X. Shi, K. Schutte, and M. Welling. The functional neural process. Advances in Neural Information Processing Systems, 32, 2019. Cited on page 6.
A. Lugmayr, M. Danelljan, A. Romero, F. Yu, R. Timofte, and L. Van Gool. Repaint: inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11461-1147, 2022. Cited on pages 8-11.
I. Macèdo and R. Castro. Learning Divergence- Free and Curl-Free Vector Fieldy with Matrix- Natrix-Valued Kernels. Pré-Publicações / A.: Pré-publicações. IMPA, 2010. Cited on pages 4, 6.
W. McKinney et al. Data structures for statistical computing in python. In Proceedings of the 9th Python in Science Conference, volume 445, pages 51-56. Austin, TX, 2010. Cited on page 13.
N. Miolane, N. Guigui, A. L. Brigant, J. Mathe, B. Hou, Y. Thanwerdas, S. Heyder, O. Peltre, N. Koep, H. Zaatiti, H. Hajri, Y. Cabanes, T. Gerald, P. Chauchat, C. Shewmake, D. Brooks, B. Kainz, C. Donnat, S. Holmes, and X. Pennec. Geomstats: python package for riemannian geometry in machine learning. Journal of Machine Learning Research, 21(223):1-9, 2020. URL.: http://jmlx wwwoooogg....oeeo... papppppppp///aaa1 http://v1111///// 027 htmml. Cited on page 10.
B. Øksendal. Stochastic differential equations. In Stochastic differential equations, pages 65-84. Springer, 2003. Cited on page4. 4.
G. Papamakarios, E. Nalisnick, J. Rezende, S. Mohamed, and B. Lakshminarayanan, Normalizing flows for probabilistic modeling and inference. arXiv preprint arXiv: 2019. Cited on page 6.
A. Phillips, T. Seror, M. Hutchinson, De Bortoli, A. Doucet, and E, Mathieu. Spectral Diffusion Processes, Nov. 2022. URL: http: /taxxi//axiiv...o trggabss/2000 14125. Cited on pages 3, 6, 1, 4, 5, 8.
Pidstrigach, Y. Marzouk, Reich, and S. Wang. Infinite-iimensional Diffusion Models for Function Spaces, Feb. 2023. URL: http: /ttt:i//axxv....//// htrg/abs/23022 10130. Cited pages 1, 6.
D. Pollard. User's Guide to Measure Theoretic Probability. Cambridge Series in Statistical and URL:
Probabilistic Mathematics. Cambridge University Press, 2002. ISBN:
ntttps:/bbokks. google. utk/bookssiid-B/cccccoooo c2G21MC Cited on page 5.
A. Rahimi and B. Recht. Random features for large-scale kernel machines. Advances neural information processing systems, 20, 2007. Cited on page 2.
13

---

C. E. Rasmussen. Gaussian processes in machine learning. In Summer school on machine learning. pages 63-71. Springer, 2003. Cited on pages 2.
S. Rissanen, M. Heinonen, and A. Solin. Generative Modelling With Inverse Heat Dissipation, July 2022. DOI: 10. 48550/a5555 arriv. 2206 13397. Cited on page 6.
G. O. Roberts and R. L. Tweedie. Exponential convergence of Langevin distributions and their discrete approximations. Bernoulli, 344333, 1996. ISSN: 1350-7265. DOE: 10. 2307/3318418, URL: https://doi org/10 10. 2307 /3318413. Cited on page 13.
C. Saharia, J. Ho, W. Chan, T. Salimans, D. J. Fleet, and M. Norouzi. Image super-resolution via iterative refinement. arXiv preprint arXiv:2104. 076.6666,, 2021. Cited on page
S. Särkkä and A. Solin. Applied Stochastic Differential Equations. Cambridge University Press, 1st edition, 2019. ISBN: 978-1-108- 118--853--55 978-1-316-51008- 978-1-316-64946-6. DOI: 10. 0017/978110818077
Cited on page
L. Scott and J. Serre. Linear Representations of Finite Groups. Graduate Texts in Mathematics. Springer New York, 1996. ISBN: 978-0-387-90190-9, URL: htttps://book...ge0 co. cu.uk/ books?idd-NGE7/r54114.. Cited on page
G. Singh, J. Yoon, Y. Son, and S. Ahn. Sequential neural processes. Advances Neural Information Processing Systems, 32, 2019. Cited on page 6.
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, 8.
Y. Song, J. Sohl-Dickstein. D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-based generative modeling through stochastic differential equations. International Conference on Learning Representations, 2021. Cited pages 1, 2, 6,
D. W. Stroock and S. S. Varadhan. Multidimennionaal diffusion processes. Springer, 2007. Cited on pages 11, 13.
N. Thomas, T. Smidt, Kearnes, L. Yang, L. Li, K. Kohlhoff, and P. Riley. Tensor field networks: Rotation- and tramslation-eqquivrrian neural networks for 3D point clouds. arXiv:1802.0821 [cs], May 2018. URL: //atpi//arxiv... org/
argg/bb//a abs/ /102..80...0 11209. 08219. Cited on pages 8, 16.
M. Titsias Variational learning of inducing variables in sparse gaussian processes. In D. van Dyk and M. Welling, editors, Proceedings of the Twelth International Conference on Artificial Intelligence and Statistics, volume 5 of Proceedings of Machine Learning Research, pages 567-574, Hilton Clearwater Beach Resort, Clearwater Beach, Florida USA. PMLR, 16-18 Apr 2009. URL: https: //proceedings. mlr teooss/v5/titsaaee
html. Cited on page 2.
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/ arxiv. 2206. 04119. Cited on page 8.
G. Van Rossum and F. L. Drake Jr. Python reference manual. Centrum voor Wiskunde en Informatica Amsterdam, 1995. Cited page 13.
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention all you need. In Advances Neural Information Processing Systems, 2017. Cited on page 16.
T. Vincent, L. Risser, and P. Ciuciu. Spatially adaptive mixture modeling for analysis of fMRI time series. IEEE Trans. Med. Imag., 29(4):1059-1074, Apr. 2010. ISSN: 0278-0062. DOI: 10. 1109/ TMI. 2010. 2042064. Cited on page 4.
P. Virtanen, R. Gommers, T. E. Oliphant, M. Haberland, T. Reddy, D. Cournapeau, E. Burovski, P. Peterson, W. Weckesser, J. Bright, J. van der Walt, M. Brett, Wilson, K. Millman, N. Mayorov, A. J. J. Nelson, E. Jones, R. Kern, E. Larson, C. Carey, i. Polat, Y. Feng. E. W. Moore, VanderPlas, Laxalde, J. Perktold, R. Cimrman, I. Henriksen, E, A. Quintero, C. R. Harris, A. M. Archibald, A. H. Ribeiro, F. Pedregosa, P. van Mulbregt, and SciPy 1.0 Contributors. SciPy
14

---

1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-272. 2020. Cited on page 13.
V. Voleti, C. Pal, and A. M. Oberman. Score-based Denoising Difftssinn with Non-Isotropic Gaussian Noise Models, 2022. URL: https:/ net/ id igcBc.K.bbb.. Cited on page
M. Weiler, P. Forré, E. Verlinde, M. Welling, et al. Coordinate independent convolutional networks: isometry and gauge equivariant convolutions on riemannian manifolds, 2021. Cited on pages 4, 5,
M. Weiler and G. Cesa. General EEiviS-EEuivvriian Steerable CNNs. arXiv:1911.0823 [cs, eess]. Apr. 2021. URL: http://arxivy wwtp://////.bb....o/ 1911. 08251. Cited on pages 8, 17.
M. Weiler, M. Geiger, M. Welling, W. Boomsma, and T. S. Cohen. 3D steerable CNNs: Learning rotationally equivariant features in volumetric data. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Biance and R. Gareet, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URl: https: http///proeeeen...eeee neurips cc/ pppp://084..///66600000000020222 Paper pdf. Cited on pages 16, 17.
Terenin, Mostowsky, and M. Deisenroth. Efficiently sampling functions from gaussian process posteriors. In ICML, pages 10292-10302. PMLR, 2020. Cited on page 2.
Wilson, BBoroviskiy, A.
O. Yadan. Hydra framework for elegantly configuring complex applications. Github, 2019. Cited on pages 10, 13.
Yim, B. L. Trippe, V. De Bortoli, E. Mathieu, A. Doucet, R. Barzilay, and T. Jaakkola, Se (3) diffusion model with application to protein backbone generation, arXiv preprint 0rX1:23302.02277, 2023. Cited on pages 4, 7.
15

---

Supplementary to:
Geometric Neural Diffusion Processes
A Organisation of appendices
In this supplementary, we first introduce in App. B an Omsttei Uhlenbeck process on function space (via finite marginals) along with several score approximations. Then in App. C, we show how this methodology extend to manifold-valued inputs or outputs. Later in App. D, we derive sufficient conditions for this introduced model yield group invariant process. What's more in App. E, we study some conditional sampling schemes. Eventually in App. F, we give thorough description of experimental settings along with additional empirical results.
B
Ornstein Uhlenbeck on function space
B.1 Multivariate Censtein-Uhlenben process
First, we aim show that we can define stochastic process on an infinite dimensional function space, by defining the joint finite marginals Y (z:) as the solution of multidimensional Orsstein Uhlenbeck process. In particular, for any set of input x = (æ1, .xR E X* we define the joint marginal as the solution of the following SDE
(8) Proposition B.1. (Phillips et al., 2022) We assume we are given a data process (Yo(z))ue DEx and we denote by G ~ GP(O, k) a Gaussian process with zero mean and covariance. Then let 's define
dύ.()) = (m(()- Yx(x))/2 d + (8,K(z,z)dB,
Y, 2 e 2:1i 9.do Yo + (1- Gade) Srde) m + (1-e- B.ds 1/2 G.
Then (Y(())) is a stochastic process (by virtue of being a linear combination of stochastic processes). We thus have that Y, a.S., Yo and Y, a.s., effectively (Y((z))EER, ,æEx interpolates between the data process and this limiting Gaussian process.
Yoo with Y ~ GP k), so
t-4
1-000
Additionally, C(Y, yu) GP(m;, with my = 2.5 B.ds you (1 e1- 3. 8, m and Σ c si-o da K. Furthermore, (Y (((()))ERR,.Ee is the solution of the SDE in (8).
Proof. We aim to compute the mean and covariance of the process (Y,) described by the SDE (3). First let's recall the time evolution of the mean and covariance the solution from a multivariate process given by
dY, f(Y,, t)dt + L(Y,, t.
know that the time evolution mean and the covariance are given respectively by Särkkä and Solin (2019)
dm E(f((Y,,))
(10)
dΣ,
Elf(Y,, t)
-
(11)
Plugging in the drift f(Y,,t) = 1/2 (m Y,) and diffusion term L(Y,, t) /BKK from (3),
we get
dm:
= 1/2 (m Y,)
(12)
dt
dΣp
= 3 3K(K KK -e].
(13)
dt

---

Solving these two ODEs we get
my = E - 3. Pods m0 + e Aodo m
Andomp
(14)
M, K e 1i Pode (200- K)
(15)
with my - Y and Σg - Col
Now let's compute the first two moments of (Y (2) We have
E Y = E e Simo ds Yo + (1 S: Foda) m + (1 61-6- 2:J Rodo) G
(16)
e 2 Pado mo + (1 e1-e-t Pode Bode) m
(17)
m,
(18)
Cov[Y,] = Cov 2i Bodo Yoo + Cov (1 -1-e- 1/2
(19)
= e -- Bizo Eo + (1 c Sizo
-Sizo Bods) Bizo
K
(00)
K = So==i=o 2ods (Eu K)
(21)
Σg
(22)
B.2
Conditional score
Hence, condition on Yo the score is the gradient of the log Gaussian characterised by mean majo = - and Enjo- = - e (1-e=B(()) with B(t) fo ((s)ds which be derived from the above marginal mean and covariance with mo Yo and Co = O.
V y, (oggo(YY./o) YY) Vv. log (Y, [meooo Egjo)
(23)
= Vx. - 1/2(Y, - m4jo) T Σ (Y, -mejo) +
(24)
-Ejo Yr m,jo)
(25)
-I1,j0 TT, I 20/0
(26)
--1,0
(27)
where Lejo denotes the Cholesky decomposition of Enjo Lojo and Y. = -4jo + L.joe
Then we can plugin our learnt (preconditioned) score into the backward SDE 4 which gives
dY, Jr [m(mx)) - Y.)/2 K(z,x)Vv, logpr-((t. Y.)] dt + v (8 (28)
B.3 Several score parametrisations
In this section, we discuss several parametrisations of the neural network and the objective.
For the sake of versatility, we opt to employ the symbol Do for the network instead of s, as mentioned in the primary text, it allows us approximate not only the score but also other quantities from which the score can be derived. In full generality, we use a residual connection, weighted by R. to parameterise the network
Cout: Cskip
(29) We recall that the input the network is time and the noised vector Y. Mrjo + n. where Prjo = e B and n ~ (0, Enjo with Edo = (1 e-B()) The gram matrix K corresponds to x) with & the limiting kemel. We denote by Lujo and S respectively the Cholesky decomposition of Enj L.joLiT and K SST.
Do(1, Y.) = (aii(() + Caun (1) Y.).
The denoising score matching loss weighted by A(t) is given by (())= = E |Do Y.) Vy, log IY
(30)

---

Table 3: Summary of different score parametrisations as well as the values for Cakip and Coul that we found to be optimal, hased the recommendation from Karras et al. (2022, Appendix B.6).
No precond.
Precond. K
Precond. S
Predict Yo
Cakp
(orj + -110-9-11
(oro + 10-3)-1 lon Do + SzlB2
(arj + 1-009--1 I Do + =112
Cout
Loss
1 Do 2|12
[| Yoll 2000 -e Da)
80
KV log
D
Do
SD,
No preconditioning By reparametrisation, let Y, = Mrjo + Liu where 2 ~ NN(0 the loss from Eq. (30) can be written as
( = [|Da(1, Y,) + 27j (Y, - - 1110)
(31)
E [|Da(t, Y.) + 2 Lu 11(0
(32)
= E [DDo(t. Y.) + Lnjo z 200]
(33)
(34)
Choosing A(t) = Σnj0 = L we obtain
[(0) = E [ILT Do Y.) + x1
(35)
E [Ilo:00 Do Y.) + =1
(36)
Preconditioning by K Altematively, one train the neural network to approximate the pre- conditioned score Dg x KVv, (Y, Y)) The loss, weighted by A = o2 I, then given by
((0) E Do((, Y,) K 10j0 =1/200)
(37)
E [Dt Y.) 210
Sz 11)0
(38)
= [o(o Y + Sz
(39)
Precondition by ST A variation of the previous one, is to precondition the score by the transpose Cholesky of the limiting kernel gram matrix, such that Dg = S Vv, logpt pe YY,
The loss, weighted by A = o2 becomes
(()) E
Dul(1 (D)( + 1 = =112(0)
(40)
Dottt Y, + 77j0 = =13200)
(41)
[log Y.) =1
(42)
Predicting Yo Finally, an alternative strategy to predict Yo from a noised version Y, this case, the loss takes the simple form
((0) E ||] (b Y.) Yolk
The score can be computed from the network's prediction following
V (og))(YY,|Yo) = 2njo (Y, Mejo)
(43)
-Y e
(44)
= -Eqj (Y, -e-(0-B(1/2 (tt)
(45)
(46)
Table summarises the different options for parametrising the score as well as the values for Gaip and Cout that we found to be optimal, based on the recommendation from Karras et al. (2022, Appendix B.6). In practice, we found the precondition by parametrisation produce the best results, but we refer to App. F.1.3 for more in-depth ablation study.

---

B.4 Exact (marginal) score in Gaussian setting
Interpolating between Gaussian processes (m Σe) and GP(m,
KV Y, log (Y,,) -KE, (m,)
(47)
= =-K(K + 6 J: Bads (E -)] "Y, -m+)
(48)
=-K(L,Li - -m4)
(49)
-KKL, - L, --, -m4)
(50)
(51)
with Σ, K + Sizo P.da (Eo K) L,L, obtained via Cholesky decomposition.
B.5
Langevin dynamics
Under mild assumptions V log Pr- (Durmus and Moulines, 2016) the following SDE
dY. KV log Po (Y.,ds + VdBB
(52)
admits a solution (Y,)s20 whose law ((Y,) converges with geometric rate to pr-t for any invertible matrix K.
B.6 Likelihood evaluation
Similarly to Song et al. (2021), we can derive a deterministic process which has the same marginal density as the SDE (3), which is given by the following Ordinary Differential Equation (ODE)- referred as the probability flow ODE
Y,(x) (lop
{m(x) - Y=(x) - (1oR).(((())) (-gaa
B {m((x) - Y,(x) - K(r,x)V logp:(Y 3, dt.
(53)
Once the score network is learnt, we can thus use it in conjunction with an ODE solver to compute the likelihood of the model.
B.7 Discussion consistency
So far have defined a generative model over functions via its finite marginals Y (z). These finite marginals were to arise from a stochastic process if, as per the Kolmogorov extension theorem (Øksendal, 2003), they satisfy exchangeability and consistency conditions. Exchangeability can be satisfied by parametrising the score network such that the score network is equivariant w.r.t permutation, i.e. so a x, y) o so x. y) for any a E E,, Additionally, we have that the noising process (Yr(x))zEx is trivially consistent for any E R4 since it is a stochastic process as per Prop. B, 1, and consequently so the score Se V log Pt, we lose the consistency over the generative process Y (x) as the constraint on the score network is non trivial to satisfy. This is actually really strong constraint on the model class, and as soon as one goes beyond linearity (of the posterior w.r.t. the context set), to enforce without directly parameterising a stochastic process, e.g. as Phillips et al. (2022). There thus seems to be strong trade-off between satisfying consistency, and the model's 's ability fit complex process and scale to large datasets.
(true) time-reversal (x)(x)) Yet, when approximating
non trivial
Manifold-valued diffusion process
C.1 Manifold-valued inputs
In the main text we dealt with simplified case of tensor fields where the tensor fields are over Euclidean space. Nevertheless, Significant work has been done on performing convolutions on feature fields on generic manifolds (a superset tensor fields on generic manifolds), core references being (Cohen, 2021) for the case of homogeneous spaces and (Weiler al., 2021) for more general Riemannian manifolds. We
certainly possible to apply our methods to these settings.

---

recommend these as excellent mathematical introductions to the topic and build on them to describe how to formulate diffusion models over these spaces.
Tensor fields as sections of bundles. Formally the fields we are interested in modelling are sections of associated tensor bundles of the principle G-bundle on manifold M. We shall denote such bundle BM and the space of sections (BM). The goal, therefore, is to model random elements from this space of sections. For clear understanding of this definition, please see Weiler et al. (2021, pages 73-95) for an introduction suitable to ML audiences. Prior work looking at this setting is (Hutchinson et al., 2021) where they construct Gaussian Processes over tensor fields on manifolds.
Stochastic processes on spaces of sections. Given we can see sections as maps σ M + BM, where an element in BM is tuple (m, b), m in the base manifold and b in the typical fibre, alongside the condition that the composition of the projection proj; (m,, b) H> m. with the section is the see distribution over sections stochastic processes with
identity, proj, oa = Td it is clear index set the manifold M, and output space a point in the bundle BM, with the projection condition satisfied. The projection onto finite marginals, i.c. a finite set of points the manifold, defined as Tm 19 (a) (mmm), o (m,)).
Noising process. To define noising process over these marginals, we can use Gaussian Processes defined in (Hutchinson et al., 2021) over the tensor fields. The convergence results of Phillips et al. (2022) hold still, and so using these Gaussian Processes as noising processes the marginals also defines a noising process on the whole section.
Reverse process. The results of (Cattiaux et al., 2021) are extremely general and continue to hold in this case of SDEs on the space of sections. Note we don't ' actually need this to be the case, we can just work with the reverse process on the marginals themselves, which are much simpler objects. It is good to know that it a valid process on full sections though should one want to try and parameterise score function the whole section akin to some other infinite-dimension diffusion models.
Score function. The last thing do therefore is parameterise the score function on the marginals. If we were trying to parameterise the score function over the whole section at once (akin to a number of other works on infinite dimension diffusions). this could present some problems in enforcing the smoothness of the score function. As we only deal with the score function on a finite set of marginals, however, we need not deal with this issue and this presents distinct advantage in simplicity for our approach. All we need to do is pick way of numerically representing points on the manifold and b) pick a basis for the tangent space of each point on the manifold. This lets us represent elements from the tangent space numerically, and therefore also elements from tensor space at each point numerically as well. This done, we can feed these to a neural network to learn to output a numerical representation of the score on the same basis at each point.
C.2 Manifold-valued outputs
In the setting, where one aim to model a stochastic process with manifold codomain Y+(x) = (Y,(r)),.... (x.)) E M things are less trivial manifolds do not have vector space structure which is necessary to define Gaussian processes. Fortunately, We can still target a know distribution marginally independently
each marginal, since this is well defined, and as such revert to the Riemannian diffusion models introduced in De Bortoli et al. (2021) with n independent Langevin noising processes
(54) are applied to each marginal. Hence in the limit - co. (x()) has density (assuming it exists) which factors as dp/ (VVolM((((11)) ... .())))) x II. e ()), For compact manifolds, we can target the uniform distribution by setting U'xx) = 0. The reverse time process will have correlation between different marginals, and so the score function still needs to be a function of all the points in the marginal of interest.
d' = (xx)( Bid + v BBdB....
D Invariant neural diffusion processes
D.1 E(n)-equivariant kernels
kernel k: IRd x IRd -> Rdxd is equivariant if it satisfies the following constraints: (a) k is stationnary, that is if for all x, R
5

---

k(a,=²) = k(x x²) = k( x¹)
(55)
and if (b) it satisfies the angular constraint for any h E H
kehrr, ho') (x) (p()T
(56)
trivial example such an equivariant kernel is the diagonal kernel k(z, x²) = ko(2, x²)I (Hold- errieth et al., 2021), with ky stationnary. This kernel can be understood has having d independent Gaussian process uni-di dimemeiinionnl output, that is, there is no anter-dmmensionl correlation.
Less trivial examples, are the E(n) equivariant kernels proposed in Macêdo and Castro (2010).
Namely curl-free and divergence-free kernels, allowing for instance to model electric or magnetic fields. Formally we have keurl = kig and kaiv kg B with ko stationary, e.g. squared exponential
kernel ko (2t x') = =0² exp
I and A and B given by 21
(r (2)( -
x')T
A(r, =
(57)
12
B(x, x') (x x') (I x')T
r x' 112
+
11 11
I.
(58)
12
12
See Holderrieth et al. (Appendix C, 2021) for proof.
D.2 Proof of Prop. 3.2
Below give proofs for the group invariance of the generative process, one via the probability flow ODE and one directly via Fokker-Planck.
Proof. Reverse ODE. The reverse probability flow associated with the forward SDE (3) with approx- imate score So((,)) R V log is given by
dY [-m(a)+ (- + Y, + ((z, - 1, Y.) de
(59)
- SODE (t, x. D,)dt
(60)
This ODE induces flow of X" y - TY" for given integration time t, which is said to be G-equivariant if the vector field is G G-equivrriant itself, i.e. b(l, a, Y.)) = p(g)b(l, Y.). We have that for any g G
bODE (1,g x. p(g) V,)
mmgg 2) + +()) + K(g x. x) so g x. p(g) Y,)]
(61)
!1)
pl) (20 so(1, x. p(g)))))))))
(62)
2 (==()))(z)++p()) p pl) p(g) sa Y.)
(63)
3 p(() | (z) + Y, + K(x, x) ss (t, Y.)
(64)
p(g) (t, Y.)
(65)
with (1) from the G-invariant prior GP conditions on and k. (2) assuming that the score network is G-equivariant and (3) assuming that p(g) E O(n). To prove the opposite direction, we can simply follow these computations backwards. Finally, know that with a G-invariant probability measure Pref and G-equivariant map o, the pushforward probability measure Prei is also G- invariant (Köhler et al., 2020; Papamakarios et al., 2019). Assuming G-invariant prior GP, and a G-equivariant score network, we thus have that the generative model from Sec. 3 defines marginals that are G-invariant.
Proof. Reverse SDE. The reverse SDE associated of the forward SDE (3) with approximate score So (l, V log is given by
dÝ{/== = [m(m(x) (-((() Y.)/2 ((x, (T) - x. Y.) dt + (3,K(z, x)dB, A bspeett, Y,)dt + (t, dB,-
(66)
(67)

---

As for the probability flow drift bODE, we have that ÓSDE is similarly G-equivariant, that is bODE gt.g z, p())Y.) = p(g)bsDE (t, x, Y,) for any g G. Additionally, we have that diffusion matrix is also G-equivariant as for any g E G we have [ g x) = 3.K(g I, g x) (z, p(g))((,=))({)) since K is the gram matrix of G-equivariant kemel k.
Additionally assuming that SSDE and Σ are bounded, Yim et al. (Proposition 3.6, 2023) says that the distribution of Y, is G-invariant, and in in particular L())).
D.3 Equivariant posterior maps
Theorem D.1 (Invariant prior stochastic process implies an equivariant posterior map). Using the language of Weiler et al. (2021) our tensor fields are sections of an associated vector bundle of manifold M with a G structure. Let ISoMIG be the group of G-structure preserving isometries on M. The action of this group on section of the bundle J E (AA) is given by
f=f o
(Weiler et al., 2021). f P. a distribution over the space of section. Let o > P be the law of of d f. Let Vs= ((f())) Tz P, the law of f evaluated at point, where na is the canonical projection operator onto the marginal at x. # the pushforward operator in the measure theory sense, x E M and y in the fibre of the associated bundle. Let 5ez a (f)( = 3²) = n a of
Tz (2))
y'),, the conditional law of the process when given [xx') y'.
Assume that the prior is invariant under the action of Isome:M, i.e. that D Mex (( -r
Then the conditional measures are equivariant, in the sense that Φ 12 af =) #
of (x) = 2 ((2):0..a1)) 1ta (oo af)
Proof. VA, B test functions, @ E ISOMGM,
B(B((=)))A((O>> [(²)) = -IB
D( r('']l
E E(B((CE) Al)((n))) AWWCew)
M(x')
1(4) (dy)
= | B(J') / AlW)(o ((0) (dy)
= BB')) / A (o (dy)
By invariance this quantity is also equal to E[B(((-1) > f)'') A(( - (x)) F[B(( (
E (B(FaF(o')) / Alw)a o(e') ofa'). (do)
= f BWW) / Aly) ( (dv) (0..1) it (oter)dd)))
f B(V)) / A(v)) do M) (dw) Ho (dg'))
Hence
(bb wit Me)) (du)) = go(o' go(l' (dg) (dyy)
7

---

By the stated invariance x- > Mar = Mo', hence
( 15 s(t's (dy) (10ll'100dl)) (dy) a.e. yf
So
ò D u 122 ,(x')
11 160
(68)
as desired.
E
Langevin corrector and the iterative procedure of REPAINT (Lugmayr et al., 2022)
E.1
Langevin sampling scheme
Several previous schemes exist for conditional sampling from Diffusion models. Two different types of conditional sampling exist. Those that try sample conditional on some part of the state space over which the diffusion model has been trained, such as in-painting or extrapolation tasks, and those that post-hoc attempt to condition on something outside the state space that the model has been trained on.
This first category is the one we are interested in, and in it we have:
Replacement sampling (Song et al., 2021), where the reverse ODE SDE is evolved but by fixing the conditioning data during the rollout. This method does produce visually coherent sampling in some cases, but is not exact conditional sampling method.
SMC-based methods (Trippe et al., 2022), which are an exact method up to the particle filter assumption. These can produce good results but can suffer from the usual SMC methods downsides on highly multi-model data such as particle diversity collapse.
The RePaint scheme of (Lugmayr et al., 2022). While not originally proposed as an exact sampling scheme, will show later that it can in fact be shown that this method is doing specific instantiation of our newly proposed method, and is therefore exact.
Amortisation methods, e.g. Phillips et al. (2022). While they can be effective, these methods can never perform exact conditional sampling, by definition.
Our goal is to produce exact sampling scheme that does not rely on SMC-based methods. Instead, we base our method on Langevin dynamics. If we have a score function trained over the state space [x" 2²), where xf are the points wish to condition and æ, points we wish to sample, exploit the following score breakdown:
Vn logp(x" = Va- logr(("", x) Vn. logp(**) Va- p(x)
If we have access to the score on the joint variables, we, therefore, have access to the conditional score by simply only taking the gradient of the joint score for the variable we are not conditioning on.
Given we have learnt (t, x) R Va log Pe(x), we could use this to perform Langevin dynamics at t= c, some time very close 0. Similar (Song and Ermon, 2019) however, this produces the twin issues of how to initialise the dynamics, given a random initialisation will start the sampler in a place where the score has been badly learnt, producing slow and inaccurate sampling.
Instead, we follow a scheme of tempered Langevin sampling detailed in Alg. 1. Starting at / = T we sample an initialisation of y" based on the reference distribution. Progressing from / = T towards f = we alternate between running a series Lavgevin corrector steps to sample from the distribution pr (y" (,)) and single backwards SDE step sample from Pa (1),1111))) with step size y. At each inner and outer step, we sample noised version of the conditioning points y based forward SDE applying noise to these context points, Pt, "f For the exactness of this scheme, all that matters is that at the end of the sampling scheme, we are sampling from P= (y" (yy) (up to the away from zero clipping of the SDE). The rest of the scheme is designed to map from the initial sample at t = T of y to a viable sample through regions where the score has been learnt well.
Given the noising scheme applied to the context points does not actually play into the theoretical exactness of the scheme, only the practical difficulty of staying near regions of well-learnt score, we could make a series of different choices for how to noise the context set at each step.
8

---

Fully independent noise, every inner step
Fully independent noise, every outer step
SDE path noise, every outer step
Fixed context set
Outer step
Inner step
Figure Comparison of different context noising schemes for the conditional sampling.
Table 4: Comparison of complexity of different noise sampling schemes for the context
Scheme
Closed-form noise
Simulated noise
Re-sample noise every inner step
(NNI
O(N
Re-sample noise at every outer step Sampling an SDE path on the context No noise applied
((N)
O(N
((N)
O(N)
The choices that present themselves are
1. The initial scheme of sampling context noise from the SDE every inner and outer step.
2. Only re-sampling the context noise every outer step, and keeping fixed to this for each inner step associated with the outer step.
3. Instead of sampling independent marginal noise at each outer step, sampling single noising trajectory of the context set from the forward SDE and use this as the noise at each time.
4. Perform noising all. Effectively the replacement method with added Langevin sam- pling.
These are illustrated Fig. 8. The main trade-off of different schemes the speed at which the noise can be sampled sample diversity. In the Euclidean case, we have a closed form for the evolution the marginal density of the context point under the forward SDE. this case sampling the noise at given time is (l1) cost. On the other hand, in some instances such nosing SDEs on general manifolds, we have to simulate this noise by discretising the forward SDE. In this case, is (n) cost, where is the number of discretisation steps in the SDE. For N outer steps and / inner steps, the complexity of the different noising schemes is compared Table 4. Note the conditional sampling scheme other than the noise sampling is O(NI) complexity.
E.2 REPAINT (Lugmayr et al., 2022) correspondance
In this section, we show that:

---

Algorithm Conditional sampling with Langevin dynamics.
Require: Score network 8o (t, s. y). conditioning points (æ", y²), query locations x²
u="x""
Augmented inputs set Sample initial noise
Pr N(mi(= ) k(x", x=="x"")) x²))
for e {T, T 7 .... c)
do
y: ~ pt. V. v6)
Noise context outputs P Sample tangent noise
2 ~ No, ld)
L. 87 Ø7- Vi: V. + + (m(x) - [bC (fl. + K( t)so(t, a, (v D.) + ) 2)
Euler-Maruyama step
for E {1,....
L do
Vi Vi-y 7t-7,0" (yi v6)
Noise context ouiputs b Sample tangent noise
Z ~ (0,10
L Vini
Vi +3K(0, ())(t,)) ~, [V + vTK(a, a) Z Langevin step
y -i-,- = Vi-7
return V:
(a) Alg. and Alg. Repaint from (Lugmayr et al., 2022) are equivalent in specific setting.
(b) There exists a continuous limit (SDE) for both procedures. This SDE targets a probability density which does not correspond to p( (zz).
(c) When to - 0 this probability measure converges to p( (z)) which ensures the correctness the proposed sampling scheme.
We begin by recalling the conditional sampling algorithm we study in Alg, 1 and Alg. 2.
Algorithm 2 REPAINT (Lugmayr et al., 2022).
Require: Score network (t æ, y). conditioning points (æ", ,y",, query locations
@ [ x²]
Augmented inputs set P Sample initial noise
[yr, VT ~ N(m(@), k( Ø))
for f [7.77 7,.... c) do
V. y
E [1,.... I.} do
#f ~ N(m (x' y (a ("","";""")) w")""))
b. Noise context outputs
Z ~ N(0, Id)
& Sample tangent noise
[. _ I 3el +7 (- (m(() - (19( + K(æ, x s)soltt [v =(1) + /FK(2 (x, x) 1/2
& Reverse step
Z ~ N(0, Id)
Sample tangent noise b Forward step
J. = Vi-n
y 1 (m(x') - x=)
Vi-, Vi-9
return y.
First, we start by describing the RePaint algorithm (Lugmayr al., 2022). We consider (2 Zz. Z()KEN sequence of independent Gaussian random variable such that for any k E N. Zt and z are d-dimensional Gaussian random variables with zero mean and identity covariance matrix and 인 is p-dimensional Gaussian random variable with zero mean and identity covariance matrix. We assume that the whole sequence be inferred is of size while the context is of size p. For simplicity. we only consider the Euclidean setting with K Id. The proofs can be adapted to cover the case K * Id without loss of generality.
Let us fix a time to E T). We consider the chain (Xe)aeN given by Xo E Rdd and for any k N, we define
N Xx+1/2 Xx + 2(6 1)V mn log (1x,, Xii) + 1)
(69)
where X 1o X6 - - 22o 1/270. Finally, we consider
X = Xx+11/2 + - -2)) 1/22
(70)
Note that (69) corresponds to one step backward SDE integration and (70) corresponds one step of forward SDE integration. In both cases have used the exponential integrator, (De
10

---

Bortoli, 2022) for instance. While we use the exponential integrator in the proofs for simplicity other integrators such as the classical Euler- Maruyama integration could have been used. Combining (69) and (70), we get that for any & E N we have
X = xx + 2(1 -e 7) Pog ([ X + (1 - 2:))1/222 4
(71)
Remarking that (Zk)KEN = (( + variables with zero mean and identity covariance matrix, we get that for any k N
is a family of d-dimensional Gaussian random
X6+1 Xx 2(1 c >V
Va,, n2 log Pt. ((x x e 2-)1-/2-
-
(72)
where we recall that X - to X (1 e - / Ze. Note that the process (72) is another version of the Repaint algorithm (Lugmayr et al., 2022), where we have concatenated the denoising and noising procedure. With this formulation, it is clear that Repaint is equivalent to Alg. 1. In what follows, we identify the limitating SDE of this process.
In what follows, we describe the limiting behavior of (72) under mild assumptions on the target distribution. In what follows, for any Fto E Rd, we denote
MFre = 2 2R Vza 21g
Pog (1 Ft, Pr.l) (z daf.
(73)
We emphasize that 6/2 # Vxio seo p(lp))))))..) In particular, using Tweedic's 's identity, we have that for any Tio E IRd
V Ptg (z)) STP Vzs, logp(( 111 [ro 2
(74)
We introduce the following assumption. Assumption 1. There exist L, C 2 m such that for any ii Vi E IPP and Ito 31 E Rd
|| log Pto (1 o1, V log Pro (Ive, vfIIIl S L([ vell + 1 vlll).
(75)
Assumption ensures that there exists a unique strong solution to the SDE associated with (72). Note that conditions under which log Pt. is Lipschitz are studied in De Bortoli (2022). In the theoretical literature on diffusion models the Lipschitzness assumption is classical, see Lee et al. (2023) and Chen et al. (2022).
We denote ((X)) (20) 7200 the family of processes such that for any & N and y 0, we have for any E [kry, (k + 1)m), X7 = (1 - (t ky)/n)x7 + (t k>)/>//// (K+1) and
X = x,, + 2(1 - - V)V log (xx key xx=") + 2/ - 21 key
(76)
where RERB is family of independent d-dimensional Gaussian random variables with zero mean and identity covariance matrix and for any à E N, Y > 0. by = e to x + (1 e (1-e-22a) ) /2
where (Z47) keN,y> is a family of independent p-dimensional Gaussian random variables with zero mean and identity covariance matrix. This is a linear interpolation of the Repaint algorithm in the form of (72).
Finally, we denote (X,) such that
dx, = b(X,)dt 2B,
X,
(77)
We recall that depends on to but to is fixed here. This means that are at time to in the diffusion and consider a corrector at this stage. The variable does not corresponds to the backward evolution but the forward evolution in the corrector stage. Under Assumption (77) admits unique strong solution. The rest of the section is dedicated to the proof of the following result.
Theorem E.1. Assume Assumption 1. Then lim (X. 1220 = (X))120-
This result an application of Stroock and Varadhan (2007, Theorem 11.2.3). explicits what the continuous limit of the Repaint algorithm (Lugmayr et al., 2022).

---

In what follows, we verify that the assumptions of this result hold in our setting, For any y > 0 and I E R.d we define
by(2)= = -2-))( Sge Van, log Pto (1re. 1P of (2% 1=))
-
(78)
- -(1/7)B[(X((+=1) x2y) )1 [x/ -x, 121
X = =]]
(79)
(x) - (4/h)(1 - e 7) Sar Vai log pro (Ireg 2fo))
8
Pco Ptolo( ( + (2/7)(1 27) ld
(80)
- (IX (1+1)) X2,) 1
x X/411
x2,1121 X,y = z).
(81)
Note that for any 7 > 0 and x E R4, we have
(x) = =B/1ixqann - XL, (K+1)7 - -Xxx) [X] by = x]
(82)
Ex(x) = [[ 11 Xz, ISI (X (1((114177 - -x 82 X = x]]
(83)
(84)
Lemma E.2, Assume Assumption 1. Then, we have that for any R, & > 0 and 7 E (0,
||||| 1 R] = 0, lim supfllE, - E(=)]| |r E R.,
(85)
7-30
lim supfllb, (x) - b(=) |x E R4, |||||| - R} = 0, 7-+0
(86)
lim (1/ X (k+1) 1 2 E
X = x)) 10 R" ||>||| $ R) = 0.
(87)
7-00
Where we recall that for any x E Rd
b(x) = 2 Sgr aso log (10t # Drtio (#f (=))
Σ(r) = 4 4Td.
(88)
Proof. Let R > O and y E I0. Using Assumption 1, there exists C' > 0 such that for any
Ito E Rd with ||ooo|| S R. we have V/V,, log Prggro([Ero DI S C(1 + |2 I).. Since pi is
(1
Gaussian with zero mean and covariance matrix (1 -e -20) Id, we get that for any p N, there exists A& 2 0 such that for any Ft, E Rd with ||= S R
SR VE,, s1, log Pto ([rto 21, P. (0%, 1x)), S AAx11 + |1=61T)
(89)
Therefore, using this result and the fact that for any 2 0, e 2 1 8, we get that there exists
Bg 2 such that for any k, E N and for any Tta E Rdd with | || S R
E|IIX(+117 X II' [X,, = x] S 2/ (1 + |k@
(90)
Therefore, combining this result and the Markov inequality, we get that for any Ito E Rd with | l| S R we have
lim,--011/)) sup (PP (8+1)7 -X,,1 2 & Xxy = x) Jar E Rdd, , |||||| S R] = 0,
(91)
lim, (1m))=(())))E(((XX (K+1) -1x -x7, 12. X = =x]1=00. 00 0,
(92)
=lm, (K+1)*
-X, 121 112
Xby = =x]1 =
(93)
In addition, we have that for any It, E Rdd with R >> 0
((2/x)(1 - (->)) - -111 JR. V ang log Pto (1 - 1) j0 (2f 1z)) |I
(94)
< Ai (1 (=]
1 1+7/-
(95)
We also have that for any Ito E Rd with R > 0
(4/7)11---71 See Va,, Pro (1T 11 82
Peo (2f (2) |I
(96)
< A1 + (3 (4//)) 1 - - - 1
(97)
Combining this result, (91), the fact that (0/7) - 1-
= and
(2/))ol - 1 11++ y = 0, we get that lim (upp |x Rd, || < R}
Similarly, using (91), (94) and the fact that lim, (44=)) = 0. we get
that
sop (lb b(z))l ja E Rdd, |||||| < R) = 0.
12

---

We can now conclude the proof of Theorem E.1.
Proof. We have that x + b(x)) and + Exx) 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. 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
grond (r) 2V log Pto (I ot)),,
(x) = 4 ld.
(98)
Recalling (88), we get that (98) is an amortised version of goond, Similar convergence results can be derived in this case. Note that also possible to obtain quantitative discretization bounds between (X)):20 and (x!/")(2 ")22 under the f 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(x1)) is not given by V log (Tt (z)). Denoting Uto such that for any Ita E Ra
Uza (Ita) = - SRr (log Pt. (Ita 10% )P: a< (=f)dx%
(99)
we have that VUr (Ir.) -b(x).), under mild integration assumptions. In addition, using Jensen's inequality, we have
J expl [ddtt S S Jax Pto (ato (=f (= la dxf, 1.
(100)
Hence, Toa with density proportional to I + expl-Uta( defines valid probability measure. We make the following assumption which allows us control the ergodicity of the process (X1)
Assumption 2. There exist m > and C 2 0 such that for any Ito E Rdd and a E IRP
v x, log (I [I),, m²) -m|1=.11 C(1 + ||=;|]]).
(101)
The following proposition ensures the ergodicity of the chain (X1)120- It is direct application of Roberts and Tweedie (1996, Theorem 2.
Proposition E.3. Assume Assumption 1 and Assumption 2. Then, Teg is the unique invariant probability measure of (Xt)120 and lim,- ||(())) Tio | = 0, where L(X.) is the distribution X
Finally, for any to > 0, denoting Te, the probability measure with density Un, given for any 2ι, by
Rd
Ung (1rg)) Sgr (log Pto (I jef, ))px o (2f, (=g)dxz
(102)
We show that the family of measures (Too)ss> approximates the posterior with density To p(zu (x() when to is small enough.
+
Proposition E.4. Assume Assumption 1. We have that lim Tea To where no admits density w.r.t the Lebesgue measure given by xo + p(zol=z).
Proof. This a direct consequence the fact that (-106))))) - 8ug.
This last results shows that even though do not target Ita Ptojo(=e, (z)) using this corrector term, we still target p(zol1r5 as to - O 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 (McK- inney et al., 2010). The code is publicly available at https ://g::://iii h/tptt/ttttt/tt/ com cawbrrdgeemeenneee neural www.diisionoorreeessse.
13

---

F1
Regression 1d
F111 Data generation
We follow the same experimental setup as Bruinsma et al. (2020) to generate the 1d synthetic data. consists of Gaussian (Squared Exponential (SE), MATÉRN 9), WEAKLY PERIODIC) and non- Gaussian (SAWTOOTH and MIXTURE) sample paths, where MIXTURE is 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 o² 0.05². The left column of Fig. 9 shows example sample paths for each of the 5 datasets.
The training data consists of 2 sample paths while the test dataset has 2 paths. For each test path we sample the number of context points between 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 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 oigllikelhooods (TLLs) reported in App. F.1.3 for the NDP models target white limiting kernel and train to approximate the preconditioned score K'V log 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 detailed derivation of the score parametrisations.
The dataset in the top row of the figure originates from Squared Exponential (SE) GP with lengthscale / = 0.25. We compare the performance of three different limiting kernels: white (blue), (orange), and a SE with shorter lengthscale { 0.1 (green).
a SE with longer lengthscale 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 reduction performance for both the approximate and exact score. We attribute this to the additional complexity of learning a non-diagone non-diagonal covariance.
In the bottom row of App. F.1.3 we repeat the experiment for 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 noising schemes in the conditional sampling, as discussed Fig. 8. For this, we measure the 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 ID squared exponential dataset (described above) as this gives us access to the true posterior. We use 212 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 (++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.
14

---

Data Prior
Model Prior
Model Posterior
SE
Tree Posterior
Matém- 3
2.5
0.0
2.5
Weakly Periodic
2.5
0.0
2.5
Sawtooth
0.5
2.5
Mixture
0.0
-2.5
2
Figure 9: Visualisation ID regression experiment.
White
SE (& = 1.0)
SE (E = 0.1)
exact scure
approximate scure
0.75
0.50
2
Precondition by K
Precondition by
No preconditioning
Predict
(a) Squared Exponential dataset with lengthscale = 0.25
White
Periodic 1.0)
Periodic (E 0.1)
exact approximate score
1.00
0.00
3
5
3
0
3
Precondition by K
Precondition by S
No preconditioning
Predict yo
(b) Periodic datasct with lengthscale é = 0.25
Figure 10: Ablation study targeting different limiting kernels and score parametrisations.
15

---

Table 5: Mean test log-likelihood (TLL) + standard error estimated over 4096 test samples reported. Statistically significant best non-GP model is in bold. The NP baselines (GNP, ConvCNP, ConvNP and ANP) are quoted from Bruinsma al. (2020). '*" stands for TLL below -10.
SE
MATÉRN-
WEAKLY PER.
SAWTOOTH
MIXTURE
INTERPOLATION
GP (OPTIMUM)
0.70
t0.00
-0.32.0.000
n/a
n/a
T(1) ---OMNDD
0.72-t0.03
0.32:
-0.38-t..03
1..9
0.64-t0.08
NDP
0. 0.71:t0.03
0.30=to.0
-0.37=0.03
3.39
0.64-t0.08
GNP
0.70:10.01
0.30-to.o1
-0.47.to.01
0.42-oo.o
0.10:Eo0.02
CONVCNP
0.80
0.95
-1.20:to.01
0.55-bo.02
-0.93-:0.02
CONVNP
0.46:
0.67
-1.02.oo.01
1.20
0.500
ANP
0. 00.01
0.75
-1.19:0.00
0.34
-0.69100002
GENERALISATION
GP (OPTIMUM)
0. 0.700o0000
0.31
0.32-t0.00
n/a
n/a
T(1) GEOMNDP
0.70.
0.01
0.38-1o.0
3.39 Lv.03
0.62
NDP
GNP
b..9
0.30
-0.471o..0
0.42.oo.01
0.10:to.02
CONVCNP
0..1
-0.95
-1.20.0..01
0.53.Lo.uz
-0.96.o..zz
CONVNP
-0.46=0.01
-0.67
-1.02+00.01
1.19:t..01
-0.53+00022
ANP
-1.42+0.01
-1.34+0.01
-1.33+00.0
-0.1740.00
-1.24+0..0
10
method
independent
dilependettcconexxttt _ociett_ooeee____e waceeeewihttttteeeeeeeeeee resampiigs
sample_pod context_noise
raw_com hawoooowithhreeeeeeeeeeeeeeeeeee resampingg
independent_contexx__oo
10a
++++
++++
++++
+++.
10a
0
10
20
100
num aaa___nnne___ppp
Figure 11: Ablation noising schemes for conditional sampling.
F.2
Gaussian process vector fields
Data We create synthetic datasets using samples from two-dimensionall zero-mean GPs with the following E(2)--quivariant kernels: diagonal Squared-Exxpoeettie (SF) kernel, zero curl (CURL- FREE) kernel and zero divergence (DIV-FREE) kernel, as described in App. D.I. We set the variance V5. We evaluate these GPs on a disk grid, created via 2D
2² and the lengthscale grid with 30 x 30 points regularly space on [- 10, 10]² and keeping only the points inside the disk of radius 10. We create training dataset of size 80 103. and test dataset size 10 103.
Models We compare two flavours of our model GeomNDP. One with -on-vauiiaaraat attention- based score network (Figure C.1, Dutordoir et al., 2022), referred as NDP". Another one with a E(2)-equivariant score architecture, based steerable CNNs (Thomas al., 2018; Weiler al., 2018). We rely on the e3nn library (Geiger and Smidt, 2022) for implementation. A knn graph is built with A 20. The pairwise distances first embed into μ(rab) with 'smooth, finite' basis of 50 elements via e3nn. soft_ -nnettoo_____e hot_linspace and with maximum radius 2. The time mapped via sinusoidal embedding Φ(t) (Vaswani al., 2017). Then edge features are obtained as eab w(e) (((ra)))) V(a,b) E with y(e) an MLP with hidden layers of width 64. We use e3mn Fww.. Sea..ualllyyoomcettttoooorrddrdd EvEN(O.EL) Vi & (v" V= (f.ss with Y spherical harmonics up to order 2m 4 an
layers with update given by y
16

---

MLP with 2 hidden layers of width 64 acting on invariant features, and node features V having irreps 12x0e
12x0o 4x1e 4x1o. Each layer has gate non- non-linearity (Weiler et al., 2018).
We also evaluate two neural processes, a translation-equiiarian CONVCNP (Gordon et al., 2020) with decoder architecture based on 2D convolutional layers (LeCun et al., 1998) and a C4 R (2))-- cquivariant 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:// haaub.com/eteeroodderreee wwweeeetteeeeeeeeeeeeeeee CNNPs of Holderrieth et al. (2021).
Optimisation. Models are trained for 80k iterations, via (Kingma and Ba, 2015) with a learning rate of 5e -4 4 and a batch size of 32. The neural diffusion processes are trained unconditionally, that 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 8 t 3min + (8max Bmin) t, and 3min = le 4, Amax = 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 11100SXM-800-- GPUs.
Evaluation. We measure the predictive log- likelihood of the data process samples under the model on held-out test dataset. The context sets are of size 25 and uniformly sampled from 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 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 he 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 selected tracks of at least 50 time points long
at intervals of every hours. From the dataset, 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 laituue-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.
Models.
Four models were evaluated.
The GP (R. + IR2) took the raw latitud--loggitue data and normalised it. Using 2- 2-output RBF kernel with no covariance between the latitude and longitude and taking the cyclone time as input, placed GP over the data. The hyperparameters of this kernel were optimised using maximum likelihood grid search over the data. Note that this model places density outside the bounding box of [-90, 90 x 180, 180) that defines the range of latitude and longitude, and does place proper distribution on the space of paths on the sphere.
The STEREOGRAPHIC GP (R + R² (00) instead transformed the data under a sterographicc projection centred at the north pole, and used the same GP and optimisation as above. Since this model only places density on 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 (R + R uses the same preprocessing as GP (R + RR2) but uses Neural Diffusion Process from (Dutordoir et al., 2022) to model the data. This has the same shortcomings as the GP (R + R2) 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 detailed below. linear beta schedule was used with 30 1E- 4 and B 10. The reverse model was integrated back to 5e for numerical
17

---

stability. The reference measure was white noise kernel with a variance 0.05. ODEs and SDEs were discretised with 1000 steps.
The GEOMNDP (R + S²) 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 casier to define smooth score function on the sphere. This is done by outputting 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 detailed below. A linear beta schedule was used with 8n- = le 4 and 31 = 15. The reverse model was integrated back to 5e for numerical stability. The reference measure was a white noise kernel with variance 0.05. ODEs and SDEs were discretised with 1000 steps.
Neural network. The network used learn the score function for both NDPP (R RR, and GEOMNDP (IR ++ S²) is a bi-attention network from Dutordoir et al. (2022) with layers, hidden size of 128 and heads per layer. This results in 924k parameters.
Optimisation. NDP (R + IR²) and GEOMNDP (RR + S²) were both optimised using (correctly implemented) Adam for 250k steps using 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. learning rate schedule was used with 1000 warmup steps linearly from le-5 to 1e-3, and from there cosine schedule decaying from 1e-3 1e-5. With even probability either the whole cyclone track was used in the batch, 20 random points were sub-sampled to train the model better for the conditional sampling task.
Conditional sampling. The GP models used closed-fome form conditional sampling as described. Both diffusion-basse models used the Langevin sampling scheme described in this work. 1000 outer steps were used with 25 inner steps. We use v 1.0 and λ 2.5. In addition at the end of the Langevin as this visually improved performance.
sampling, we run an additional 150 Langevin steps with
Evaluation. For the model (conditional) log probabilities the GP models were computed in closed form. For the diffusion-based models, they were computed using the auxiliary likelihood ODE discretised over 1000 steps. The conditional probabilities were computed via the difference between the log-likeliihn ligllikeldoood of the whole trajectory and the logelikeld of the context set only. The mean squared errors were computed using the geodesic distance between 10 conditionally sampled trajectories, described above.
18


Reviews:
[{"review_id": "SUvH3lLE0V", "reviewer": "Reviewer_uLvV", "summary": "The present thesis deals with the problem of learning geometric-sensitive Denoising Diffusion Models (DDMs) for modeling equivariant conditional processes in (non)-Euclidean observation spaces. In doing so, the authors extend existing work that defines diffusions over function spaces with infinite dimensions using a (non-)Euclidean score-based modeling approach. They incorporate a series of geometric priors by constructing a noise process that allows a geometric Gaussian process as the limiting distribution transforming under the symmetry group of interest, and approximate the score as part of the denoising process by a neural network that is equivariant with respect to this group. They show that the generative function model admits the same symmetry under these conditions. The scalability and capacity of the proposed model is evaluated in the context of a novel Langevin-based conditional sampler based on synthetic and real data in the Euclidean and non-Euclidean scenario.", "questions": "- Regarding the usage of numerical solvers: From my point of view, numerical (SDE) solvers play a key role in the proposed approach, as it involves the simulation of (reverse) diffusions defined via SDEs. The authors use an Euler-Maruyama discretisation scheme, which is a very naive solver. Despite the fact that more advanced systems would drastically reduce computational efficiency, it would be important to address their benefits. Last but not least, the vanilla Euler-Maruyama scheme is based on a vector space structure. Therefore, in the case of manifold valued SDEs, there are more advanced options that take these intrinsic structures into account, e.g., a geometric Euler-Maruyama scheme. How could this improve the simulation process?\n- l. 66 & 72 (4) -> (2) ?\n- l. 173 admits -> admit\n- l. 179 $A \\subset C$ and how is $x^\\ast$ defined in Fig. 2?", "limitations": "Limitations are addressed (e.g., cf. Chapter 6) and relate to a computational overhead and modeling weaknesses.\nNo evidence was found that the authors addressed potential negative societal impact. Please make the effort to include this mandatory information in your work.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "I greatly appreciate any attempt to incorporate geometrically sensitive analysis into machine learning. From my perspective it is advantageous, to assume that data incorporates structural information and to give models the ability to address them.\n\nThe proposed method builds on a novel combination of well-known techniques, and is original to the best of my knowledge. Further, the paper is quite well positioned to existing work, the content is well organized and, with a few exceptions, has no spelling or grammatical flaws. Empirical evidence is accessed through experiments on toy and real world data and the authors make effort to extend the small-scale experimental setting.", "weaknesses": "I would rate the overall added value of the contribution as rather moderate. The biggest shortcomings relate to the clarity of the submission and the technical soundness which I will address with examples in the following.\n- Some general points:\n  - I'm not very familiar with the theory of Gaussian processes in detail, but I have a solid knowledge of stochastic processes in general and SDEs and Bayesian inference. From this point of view, it is sometimes difficult to keep track of the present work.\n  - For me a stochastic process $\\mathcal{X}$ is a family of random variables $(X(t,\\omega))_{t\\in T}$ defined on a common probabillity space $(\\Omega, \\mathcal{F}, \\mathbb{P})$, indexed by some set $T$; for each $t\\in T$, $X(t,\\bullet):\\Omega \\to \\mathbb{R}^n$ is a measurable function and for $\\omega \\in \\Omega$, $X(\\bullet, \\omega): T \\to \\mathbb{R}^{n}$ is a deterministic function, commonly known as path of the process. One can obtain stochastic processes as solutions to stochastic differential equations using the Itô or Stratonovich calculus; a well known example with a closed form solution is the (univariate) Ornstein-Uhlenbeck process, which is also a Gaussian process.\n  - Following, e.g., Dutordoir et al. 2022 Sec. 2.1 then a Gaussian process is (univariate) stochastic process $(f(\\omega,x))_{x\\in\\mathbb{R}^D}$ with $f:\\Omega \\times \\mathbb{R}^D \\to \\mathbb{R}$ such that, for any finite collection of points from $\\mathbb{R}^D$ the resulting random vector follows a multivariate normal distribution.\n  - In this setting we can model Gaussian processes by using the Kolmogorov extension theorem; but according to Phillips et al. 2022 this comes with difficulties.\n  - However if we switch the codomain of the function(s) $f$ and consider $\\mathbb{R}^n$ for $n > 1$ we get into trouble how to handle the Gaussian distribution condition for a finite collection of points in this case. Therefore we need a different approach to define Gaussian processes and fortunately the SDE approach turns out to be the right one combined with some convergence results.\n  - Last but not least this setting even succeeds in case of non-Euclidean codomain such as compact manifolds.\n  - What I want to share with this quick sketch is an example of how I would have liked to have been introduced to the general setting in general. In particular Sec. 3.1 would then have gained in structure.\n- Background\n  - l. 63 Please provide a convergence result reference.\n  - l. 69 Please provide a Can-easily-be-shown result reference.\n  - l. 70 Over functions $\\mathbf{s}(\\mathbf{Y}_t)$ implies $\\mathbf{s}:\\mathbb{R}^d \\to \\mathbb{R}^{d}$ or do you mean to minimize over $\\mathbf{s}\\circ \\mathbf{Y}_t$\n  - l. 83 Please check wording.\n\n  - l. 73 Can you provide a motivation for considering steerable (feature) fields?\n  - l. 89 What exactly is the feature field in (2)?\n  - l. 90 What is $p$?\n  - l. 93 Can you specify $\\rho(h)$ in Fig. 1?\n  - l. 94f. I don't understand this comment, can you rephrase it?\n- Section 3.1\n  - As mentioned at the beginning please highlight the importance of the SDE approach that becomes the right perspective in case of diffusion models for function with codomain $\\mathbb{R}^n$ and non-euclidean subspaces\n  -  How is a data process defined; can you define the observation model in more detail?\n  -  Why is it suddenly important to specify the OU process as multivariate? We already introduced it in Section 2. and what is $\\beta_t$ in Eq. (3) and (4)?\n  - Please highlight the interpolation approach (Phillips et al. 2022) from your supplementary material Sec. B.1; in my opinion that's a key result at this point. Also please introduce the used convergence results in more detail.\n  - l. 120 Do you mean __Gaussian__ white noise?\n  - l. 125 Please add a reference for what kind of mild conditions.\n  - l. 128 Please add a reference for Stein score.\n  - l. 136ff This part (i.e. Training) including the conditional score is unclear to me! Why does the noising process include a score?\n  - What numerical solver is used in the manifold case for simulating the reverse diffusion?\n- Section 3.2\n  - l. 150 Sec. 2 include the definition of a feature field; how is this connected to a tensor field\n  - l. 158 (Prop. 3.1) what about the stationary conditional of Theorem 1 (Holderrieth et al.)?\n  - l. 165 I was not able to detect the result of Prop. 3.2 in Yim et al. 2023; can you please be more specific?\n  - l. 181 Please define or give reference for a definition of a conditional process.\n  - l. 184 Please define SP.\n- Section 3.3\n  - I find it difficult to connected the described method to Alg. 1. Especially, what are detailed reasons why sampling directly at the end time proves difficult. How does Fig. 3 contributes to this section?\n  - l. 218 Please check wording?\n- Section 4\n  - The present work relates to results from De Bortoli et al. 2022, Huang et al. 2022 and Yim et al. 2023 but these are not included into this related work section, why?\n- Section 5.1\n  - l. 271 Please define GEOMNDP and you are using GeomNDP a few times instead.\n  - What about comparing models from Holderrieth et al. and Phillips et al. which are closer conceptional?\n  - l. 292 Please define T(1)-GEOMNDP.\n  - l. 295 Please define $\\mathbb{I}(x=x’)$ or provide reference.\n  - l. 305 In my opinion, Fig. 10 is insufficiently evaluated. For example, changing from zero to a single inner step results in a drastic improvement, but then the quality stays nearly the same for 5 and 10 steps; and after that the error increases; Can you elaborate on this observation?\n\nImportant side note: Including supplementary material with unfinished notes is not of advantage!\n\nHowever, I would like to reiterate my appreciation for the authors' interest in expanding the field of neural diffusion models from a geometrical point of view. I'm sure that by working on the deficiencies mentioned, the work will gain in value and will deliver a valuable contribution to the community.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to produce such a thorough review.\n\nIn particular, thanks for your concrete suggestion on how to improve section 3.1 - we have updated our manuscript accordingly. Furthermore, due to strong space constraints, we had to limit the number of references and definitions. This may have led to the feeling of a lack of technical soundness. We have incorporated our reply to your comments in the main paper in order to rectify this.\n\n**Background**\n\n- L63: Reference to Bakry et al (2014) and Durmus and Moulines (2015) for Kullback-Leibler and Wasserstein-based convergence results.\n- L69: Reference to Theorem 5.1 of Kallenberg (Foundations of Modern Probability) which can directly be applied.\n- L70: $s:\\mathbb{R}^d \\rightarrow \\mathbb{R}^d$.\n- L73: We added to the \"Steerable fields\" paragraph that \"steerable feature fields\" are collections of tensor fields. Cref answer to L150.\n- L89: $g \\cdot f(x) = (uh) \\cdot f(x) \\triangleq \\rho(h) f(h^{-1} (x - u))$ is the special case of a $G$-transformed feature field $g \\cdot f(x) = \\rho(g) f(g^{-1} \\cdot x)$ when $G = \\mathrm{E}(d) = \\mathrm{T}(d) \\rtimes \\mathrm{O}(d)$.\n- L90: $p$ -> $u$, the translation element of $\\mathbb{E}(d)$.\n- L93: We have updated the caption adding that $\\rho(h) = h$, with $h$ a 90-degree rotation matrix.\n- L94: Updated to: \"For many natural phenomena, a priori we do not want to express a preference for a particular conformation of the feature field and thus want a prior $p$ to place the same density on all the transformed fields: $p(g \\cdot f) = p(f), \\forall g \\in G$.\"\n\n**Section 3.1**\n- We assume an additive homoscedastic Gaussian observation model and have elaborated on this in the manuscript.\n- You are correct, we will remove \"multivariate\" to prevent any confusion.\n- The $\\beta$-schedule can be seen as rescaling time as $s = \\int_{t=0}^s \\beta_t \\mathrm{d}t$.\n- Phillips et al. 2022: we have highlighted their result in the main text.\n- L125: We reference Haussman & Pardoux (1986) which assume Lipschitz continuity of the drift and volatility matrix and Cattiaux et al. (2021, Theorem 4.9) which only assume a finite entropy condition on the space of processes.\n- L128: Following Song & Ermon (2019) we refer to $\\nabla_x \\log p_t$ as the _Stein score_.\n- L136: We apologize for the confusion. We have clarified that the noising (forward) process does not include the score.\n- Solver: We rely on the projective solver introduced by Hairer (2011) to enforce the manifold constraint along the trajectory.\n\n**Section 3.2**\n- L150: We apologize for using the terms \"tensor field\" and \"feature field\" interchangeably. A tensor field would be defined by the representation $\\rho_p(g)=\\rho_{\\mathrm{triv}}(g)^{\\otimes p}$ ($p$ tensor product), e.g. a scalar field would have a type $p=0$, whilst a feature field would refer to a collection of tensors, with possibly a multiplicity of tensors with the same types, and as such would have a representation of the form $\\rho = \\sum_p a_p \\rho_p(g)$.\n- L165: Proposition 3.6 of Yim et al. 2023.\n- L181: We refer to Pollard (A User's Guide to Measure Theoretic Probability, p. 117) for the definition of a conditional process.\n- SP = Stochastic process.\n\n**Section 3.3**\n\nWe agree that the sentence \"While we could sample directly at the end time this proves difficult in practice\" is confusing. This refers to the idea of only using the estimate of the score $s_\\theta(t, \\cdot)$ at time $t=\\epsilon$ with $\\epsilon \\ll 1$ and then applying Langevin dynamics. This approach fails due to both the highly non-convex nature of $\\log p_\\epsilon$ (slow mixing) and the fact that $s_\\theta(t, \\cdot)$ can only be well estimated where $p_\\epsilon$ has mass.\nIn Alg. 1., the outer loop is the discretisation of the time-reversal process (i.e. predictor step) going from $t \\rightarrow t - \\gamma$, as illustrated in blue in Fig. 3, while the inner loop is the discretisation of the Langevin dynamics (corrector step) targeting the (conditional) marginal at time $t - \\gamma$ as illustrated in red.\n\n**Section 4**\n\nDue to space constraints, we cut part of the paragraph where we were discussing De Bortoli et al., etc., as we believed that referring to them in specific contexts (within Section 3) may be enough, yet we agree that it is actually worth properly discussing.\n\n**Section 5.1**\n- Comparison to Holderrieth et al.\nIn the setting where the group of interest is the translation group, then the method of Holderrieth et al. is exactly a convolutional NPs as introduced in Gordon et al.\n- Fig. 10. We discuss this figure in Appendix F.1.3 of the (proper) supplementary material. The _total_ number of predictor and corrector calls is kept fixed, hence the higher number of corrector steps, the smaller the number of predictor steps. It appears that without any Langevin corrector step, the conditional samples are of a poor fit, yet for a high number of corrector steps, the number of predictor steps becomes too small (the discretisation step becomes too large).\n\n**Questions:**\n- Choice of SDE numerical solver: As the time-reversal SDE is not a generic SDE, one can complement the time-reversal step (following a standard SDE solver), with 'corrector' steps. As shown by Song et al. (2019) there is a trade-off between spending computational budget between a higher order predictor (time-reversal steps) and more corrector steps. Regarding the manifold co-domain, indeed there is no vector space structure, we, therefore, leverage geometric random walk (see E. Jørgensen, 1975), which relies on the exponential map to preserve manifold constraint.\n- In fig. 2, $x^*$ is a regular 1D grid for the left-hand side plots, and regular 2D grid restricted to the interior of a disk for the right-hand side plots. We will add this clarification to the caption.\n\nWe hope that these clarifications addressed your concerns with regard to the clarity and soundness of the manuscript and that you will be able to update your score accordingly."}, {"review_id": "9EgpxpbEEb", "reviewer": "Reviewer_NoKK", "summary": "This paper proposes GeomNDP, a continuous-time functional diffusion model which can incorporate symmetries of the problem of interest. In particular, this model is constructed via a collection of finite-dimensional Ornstein-Uhlenbeck processes (one for each finite-dimensional marginal), which target a given prior Gaussian process. In addition, the authors study (Prop. 3.2, 3.3) conditions under which the resulting process is invariant to a given group. A conditional sampling method is also proposed (Section 3.3) which generalizes the RePAINT method (Lugmayr et al.) to the functional setting. ", "questions": "- Could you elaborate on the differences between the proposed conditional sampler (Sect. 3.3) and the \"replacement\" method? E.g. previously proposed conditional functional diffusion models [1-2] work in a manner that, from the description starting on Line 199, in a very similar manner. \n- It was unclear to me in what sense (or if at all) defining the noising process on each finite-dimensional marginal corresponds to a well-defined (and reversible) process over the underlying functional measure of interest. Could you provide additional details regarding this correspondence?\n\n\n\n[1] Neural Diffusion Processes. Dutordoir, Saul, Ghahramani, Simpson, 2023.\n\n[2] Diffusion Generative Models in Infinite Dimensions. Kerrigan, Ley, Smyth, 2023.\n", "limitations": "The limitations are adequately addressed within the submission.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written, clear, and generally has a high-quality presentation.\n- This work contributes to the growing literature on functional diffusion models by incorporating group invariances into their design. This is likely to be of interest to the community. \n- The proposed model shows competitive performance against several baselines across multiple tasks. In general, the experimental evaluation is clear and convincing. \n- The conditional sampling method achieves empirically strong performance (Figure 7), and is an important contribution to this literature. As noted by the authors, previous conditional functional diffusion models use ad-hoc methods which do not target the correct distribution.\n- The work is well-contextualized within the relevant literature. ", "weaknesses": "- The exposition of the conditional sampler (Section 3.3) was somewhat unclear (see questions), and required a careful read of the pseudocode in order to understand what exactly was being proposed.\n- It would have been interesting to compare against recently proposed function-space diffusion models (e.g. those discussed in the related work) as additional baselines -- although many of these works are likely to be contemporaneous with the submission, so this is not a major weakness.\n\n\n### Minor comments\n- The submission was not properly anonymized (see the Appendix in the main submission)\n- Line 85: it was somewhat unclear what exact assumptions are made on the domain $X$  ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your feedback and the positive review. Below, we respond to your main concerns.\n\n**Minor comments:**\n>The submission was not properly anonymized\n\nWe apologise for mistakenly attaching an early version of the appendix with the main pdf.\n\n>Line 85: it was somewhat unclear what exact assumptions are made on the domain $X$\n\nThe assumptions on $X$ are pretty weak, a metric space would be sufficient for instance. We will update Section 2 to clarify this.\n\n**Questions**\n\n> Clarification on Conditional Sampler\n\nBased on your comment, we have improved the exposition of the conditional sampler in section 3.3. We included the following points: \n- We develop the continuous formulation of the replacement method by Lugmayr et al. (2022), which we theoretically show does not target the correct conditional distribution (as detailed in E.2).\n- We suggest a correction of the algorithm targeting the true conditional distribution using Langevin corrector steps. We theoretically and empirically show this contribution to be crucial for good results.\n- We discuss several initialization schemes for the context variables, which further improve the empirical performance of the conditional sampling algorithm.\n\n>It was unclear to me in what sense (or if at all) defining the noising process on each finite-dimensional marginal corresponds to a well-defined (and reversible) process over the underlying functional measure of interest.\n\nThis follows from Kolmogorov Extension Theorem (see, for instance, Øksendal, 2003). One can show that the exchangeability and consistency conditions are satisfied for each finite set of marginals, both in the forward and (true) reverse process. As a result, the finite marginals originate from a stochastic process. We can thus define the (de)noising process on the finite-dimensional marginals without having to explicitly deal with the underlying functional measure of interest. Thank you for the question - we have also clarified this in the manuscript.\n\nOnce more, we appreciate the time you've dedicated to reviewing our manuscript. We hope that you convey your enthusiasm for our work to the other reviewers during the forthcoming discussion phase."}, {"review_id": "Il0fzfbwMA", "reviewer": "Reviewer_zVLR", "summary": "In this paper, the authors propose a generative model over vector fields. The proposed approach is a combination of ideas from neural diffusion processes [1], equivariant learning of stochastic fields [2] and equivariant scores [3]. The authors represent vector fields as vectors in finite sets of points and use Euclidean groups to incorporate desired symmetries on vector fields. The proposed diffusion model defined via Ornstein–Uhlenbeck process converges to Gaussian Process with equivariant mean and kernel. The model is trained on the denoising score matching loss with score parametrized by an equivariant neural network. This methodology can be extended to manifold-valued functions using Riemannian diffusion models. Further, the authors explore conditional sampling from the proposed model.\n\n[1] Dutordoir, V., Saul, A., Ghahramani, Z., and Simpson, F. (2022). Neural Diffusion Processes.\n[2] Holderrieth, P., Hutchinson, M. J., and Teh, Y. W. (2021). Equivariant learning of stochastic fields: Gaussian processes and steerable conditional neural processes. In International Conference on Machine Learning.\n[3] Yim, J., Trippe, B. L., De Bortoli, V., Mathieu, E., Doucet, A., Barzilay, R., and Jaakkola, T. (2023). Se (3) diffusion model with application to protein backbone generation. arXiv preprint\narXiv:2302.02277.", "questions": "Regarding Figure 5: The default NDP fails to achieve the same level of log-likelihood even with large amount of training data. Would it reach the same log-likelihood as E(2)-GeomNDP when the number of training examples tends to “infinity” (> 10^5)? Does it need more parameters? Or is there something else that prevents it from adequately fitting the data?", "limitations": "Yes", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well written and provides all necessary information for understanding. Main distinguishing features of proposed model are well demonstrated in experiments. The proposed model needs much less training examples to reach high log-likelihood. In several small-scale experiments the proposed model outperforms existing approaches.", "weaknesses": "The experiments are small-dimensional. It would be interesting to see how the proposed approach scales to higher dimensionality, where equivariant models could become even more beneficial.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for your insightful comments and the favourable evaluation of our manuscript. In the following, we address your primary concerns, aiming to further enhance your appreciation of our work.\n\n**1. Scaling to Higher Dimensionality**\n\nWe agree that extending our experiments to higher dimensionalities could indeed yield intriguing results, especially given the potential benefits of equivariant models in such contexts. Your suggestion certainly aligns with our overall aim of comprehensively exploring the potential of the proposed approach. However, such an extension represents a substantial piece of work, likely beyond the scope of this particular paper.\n\nIn the current study, we consciously chose to focus on low-dimensional but complex datasets, such as the one-dimensional sawtooth, mixture, and the cyclone dataset on the sphere. These datasets, despite being low-dimensional, include intricate complexities and irregularities. We believe that this selection underlines not just the robustness of our methodology, but also its versatility. Furthermore, the datasets used in Table 1 are the de-facto standard in the neural process (NP) literature (Gordon et al. (2020), Bruinsma et al. (2023), Dutordoir et al. (2023)) allowing us to compare our method to a variety of NP based models.\n\nRegarding scalability, the main computational bottleneck of our method is the number of data points in the dataset due to the quadratic scaling of the attention architecture. We share this constraint with previous models such as Kim et al. (2019) and Dutordoir et al. (2023). Scaling on the input and output dimensions, on the contrary, is linear as a result of the equivariant networks. This makes our method scale favourable with both the input and output dimensions.\n\n**2. Comparison with Default NDP in Figure 5.**\n\nThat's indeed a really good remark, we do not believe that more data would help as it appears that the performance is already plateauing. Similarly, we do not believe it comes from the optimisation, as on the contrary equivariant models tend to potentially induce harder optimisation objectives. Hence, it is likely an issue of model capacity. Yet, the non equivariant 'default NDP' already has much more parameters (\\~1M) than the equivariant $\\mathrm{SE}(3)-\\textrm{GeomNDP}$ (\\~80k).\n\n\nAgain, thank you for taking the time to review our manuscript. We hope that you can share your enthusiasm for our work with the other reviewers in the upcoming discussion period."}, {"review_id": "lnZcHDqbyp", "reviewer": "Reviewer_fmuk", "summary": "This paper proposes several extensions to functional diffusion models. In particular, the focus is on extending the process to cover various geometric inputs (such as manifold values and symmetries), although there are some other components (such as a different formulation of the diffusion that allows for more general kernels and consistency).", "questions": "N/A", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The proposed methodology presents a good extension of functional diffusion models. In particular, the more general formulation allows for geometric data, which can be useful for modeling Earth science data. Furthermore, the experiments show that this and the equivariance construction are useful when the datasets exhibit the natural conditions.", "weaknesses": "* The main weakness is the experiments. The results in table 1 imply that the alternative diffusion formulation doesn't produce a tangible benefit over the previous formulation when applied to standard data. Furthermore, the other experiments are generally quite toy, which makes it hard to ascertain certain limitations of the method (for example, for S^2 and manifolds in general the kernel is not given in closed form, which can make it hard to extend this methodology to high dimensions).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your thoughtful feedback and the positive assessment of our manuscript.\n\n**Experimental benefits over the previous formulation**\nThe goal of Table 1 (bottom part) is to highlight a key advantage of our geometric approach to neural diffusion processes (NDP): the ability to generalize in a data-efficient way to group actions such as roto-translations. This is a feature that previous NDP formulations do not possess. This claim is further evidenced in Figure 5, where our geometric approach requires orders of magnitude less data to obtain superior results compared to the default model. It is also worth noting that the default NDP requires much more parameters (\\~1M) than our equivariant model (\\~80k), which can further hinder training. Therefore, while the differences between the alternative and previous formulations may not be as pronounced on standard data, it's precisely the capability to generalize and perform efficiently in varied conditions that set the geometric approach apart.\n\n**Toyish experiments**\nOur choice to feature low-dimensional but complex datasets, such as the one-dimensional sawtooth, mixture, and the cyclone dataset on the sphere, was deliberate. These datasets, despite being low-dimensional, include intricate complexities and irregularities.  We believe that this selection underlines not just the robustness of our methodology, but also its versatility when dealing with disparate datasets. While we understand your concerns about potential limitations in extending the methodology to higher dimensions, we feel that our choice of datasets effectively illustrates the strengths of our approach. Finally, with regard to Table 1, we opted to use common datasets from the neural process literature [1,2,3] to empirically assess our model and compare to a range of other neural process-like models.\n\nThank you for taking the time to review our manuscript. We hope that we have addressed your main concerns in this rebuttal and that you might consider updating your score accordingly.\n\n\n**References**\n\n[1] Foong et al. Meta-Learning Stationary Stochastic Process Prediction with Convolutional Neural Processes. NeurIPS 2020.\n\n[2] Gordon et al. Convolutional conditional Neural Processes. ICLR 2020.\n\n[3] Bruinsma et al. Autoregressive Conditional Neural Processes. ICLR 2023."}]

Author rebuttal:
Dear Reviewers,

We thank the reviewers for taking the time and effort to provide feedback on our submitted manuscript.

We are glad that reviewers found our proposed method sound, well-motivated, well-written and to deliver an interesting and valuable contribution to the community.

In light of the feedback, we've identified "Clarity and presentation" to be the primary area that required our attention. We have made the following improvements:

- To enhance the paper's self-containedness and technical depth, we incorporated definitions and key references that were initially omitted due to space limitations. For a comprehensive list, please refer to reviewer uLvV's reply.
- We've refined the exposition of the conditional sampling algorithm in Section 3.3, providing greater clarity to our novel contributions.
- In our experimental evaluation, we now more effectively underscore the advantages of the geometrical approach over the default case.

Lastly, we genuinely regret and apologize for the oversight regarding the preliminary version of the appendix bundled with the main PDF.

Once again, we want to thank all reviewers for your diligent review, and the overall positive evaluation of our manuscript.


Meta-review: This paper proposes an algorithm for diffusion generative modelling over function spaces. Their proposed algorithm also incorporates symmetries in the model, and can be extended to manifold-valued functional models. Finally, a conditional sampling method is proposed. 

The authors demonstrate clear advantage of their algorithm on a number of smaller-scale experiments, both synthetic and real-world. These provide a convincing demonstration of the value of incorporating group invariances into functional diffusion models for geometric data.

On the other hand, as noted by some reviewers, the paper's claims would be strengthened if authors can provide experimental validation on more complex and higher-dimensional problems.
Decision: Accept (poster)