Title: Geometric Neural Diffusion Processes
Venue: NeurIPS 2023
Accepted: yes
Keywords: diffusion model

functional space

stochastic process

time-series

neural processes

Gaussian processes

random fields

invariance

equivariance

symmetries

stationarity
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

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.

1 Introduction

Traditional denoising diffusion models are defined on finite-dimension Euclidean spaces (Song and Ermon, 2019; Song et al., 2021; Ho et 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 et al., 2023), and on function spaces of the form f : R n → R d (Dutordoir et al., 2022; Kerrigan et al., 2022; Lim et al., 2023a; Pidstrigach et 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 non-Euclidean 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 a value that lives in some vector space V . Roughly speaking, these are 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 : R 3 → R , representing wind over the Earth's surface as f : S 2 → T S 2 , where T S 2 is the tangent-space of the sphere, or modelling the stress in a deformed object as f : Object → T R 3 ⊗ T R 3 , where ⊗ is the tensor-product of the tangent spaces. Given the inherent symmetry in the laws of nature, these tensor fields can transform in a way that preserves these symmetries. Any modelling of these laws may benefit from respecting these symmetries.

37th Conference on Neural Information Processing Systems (NeurIPS 2023).

Secondly, we look at fields with manifold codomain, and in particular, at functions of the form f : R →M . The challenge in dealing with manifold-valued output, arises from the lack of vectorspace 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 cyclones moving on the surface of the Earth, or modelling the joint angles of a robot as it performs tasks.

The lack of data or noisy measurements in the physical process of interest motivates a 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-specific 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 ) -invariant 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 → M ) 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.

2 Background

Denoising diffusion models. We briefly recall here the key concepts behind diffusion models on R d and refer the readers to (Song et al., 2021) for a more detailed introduction. We consider a forward noising process ( Y t ) t ≥ 0 defined by the following Stochastic Differential Equation (SDE)

where ( B t ) t ≥ 0 is a d -dimensional Brownian motion and p 0 is the data distribution. The process ( Y t ) t ≥ 0 is simply an Ornstein-Ulhenbeck (OU) process which converges with geometric rate to N(0 , Id) (Durmus and Moulines, 2017). Under mild conditions on p 0 , the time-reversed process ( ¯ Y t ) t ≥ 0 = ( Y T -t ) t ∈ [0 ,T ] also satisfies an SDE (Cattiaux et al., 2021) given by

where p t denotes the density of Y t . Unfortunately we cannot sample exactly from (2) as p T and the scores ( ∇ log p t ) t ∈ [0 ,T ] are unavailable. First, p T is substituted with the limiting distribution N(0 , Id) as it converges towards it. Second, one can easily show (Kallenberg, 2021, Thm 5.1) that the score ∇ log p t is the minimiser of ℓ t ( s ) = E [ ∥ s ( Y t ) -∇ y t log p t | 0 ( Y t | Y 0 ) ∥ 2 ] over functions s : [0 , T ] × R d → R d where the expectation is over the joint distribution of Y 0 , Y t , and as such can readily be approximated by a neural network s θ ( t, y t ) by minimising this functional. Finally, a discretisation of (2) is performed (e.g. Euler-Maruyama) to obtain approximate samples of p 0 .

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 = uh where h ∈ O( d ) is a d × d orthogonal matrix and u ∈ T( d ) is a translation which can be identified as an element of R d ; for a vector x ∈ R d , g · x = hx + u denotes the action of g on x , with h 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.1.

f

(

x

)

f

(

g

-

1

x

)

ρ

(

h

)

f

(

g

-

1

x

)

We are interested in learning a probabilistic model over functions of the form f : X → R d such that a group G acts on X and R d . We call a feature field a tuple ( f, ρ ) with f : X → R d a mapping between input x ∈ X to some feature f ( x ) with associated representation ρ : G → GL( R d ) (Scott and Serre, 1996). This feature field is said to be G -steerable if it is transformed for all x ∈ X , g ∈ G as g · f ( x ) = ρ ( g ) f ( g -1 · x ) . In this setting, the action of E( d ) = T( d ) ⋊ O( d ) on the feature field f yields g · f ( x ) = ( uh ) · f ( x ) ≜ ρ ( h ) f ( h -1 ( x -u ) ) . Typical examples of feature fields include scalar fields with ρ triv ( g ) ≜ 1 transforming as g · f ( x ) = f ( g -1 x ) such as temperature fields, and vectors or potential fields with ρ Id ( g ) ≜ h transforming as g · f ( x ) = hf ( g -1 x ) as illustrated in Fig. 1, such as wind or force fields.

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 µ ( g · f ) = µ ( f ) , ∀ g ∈ 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 Y = R d , by defining a diffusion model for every finite set of marginals. Most prior works on infinite-dimensional diffusions consider a 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 on a 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 finite-dimensional subspace. With this in mind, we 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 a data process ( Y 0 ( x )) x ∈X . Given any x = ( x 1 , . . . , x n ) ∈ X n , we consider the following forward noising process ( Y t ( x )) t ≥ 0 ≜ ( Y t ( x 1 ) , . . . , Y t ( x n )) t ≥ 0 = ( Y 1 t , . . . , Y n t ) t ≥ 0 ∈ Y n defined by the following SDE

where K( x, x ) i,j = k ( x i , x j ) with k : X × X → R a kernel and m : X → Y . The process ( Y t ( x )) t ≥ 0 is a multivariate Ornstein-Uhlenbeck process-with drift b ( t, x, Y t ( x )) = m ( x ) -Y t ( x ) and diffusion coefficient σ ( t, x, Y t ( x )) = K( x, x ) -which converges with geometric rate to N( m ( x ) , K( x, x )) . Using Phillips et al. (2022), it can be shown that this convergence extends to the process ( Y t ) t ≥ 0 which converges to the Gaussian Process with mean m and kernel k , denoted Y ∞ .

In the specific instance where k ( x, x ′ ) = δ x ( x ′ ) , then the limiting process Y ∞ 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 white noise setting is not covered by the existing theory of functional diffusion models, as a Hilbert space and a square integral kernel are required, see Kerrigan et al. (2022) for instance.

Denoising process. Under mild conditions over Y 0 2 , the time-reversal process ( ¯ Y t ( x )) t ≥ 0 satisfies

with ¯ Y 0 ∼ GP( m,k ) and p t the density of Y t ( x ) w.r.t. Lebesgue. In practice, the log p T -t termknown as the Stein score, is not tractable and must be approximated by a neural network. We then consider the generative stochastic process model defined by first sampling ¯ Y 0 ∼ GP( m,k ) and then simulating the reverse diffusion (4) (e.g. via Euler-Maruyama discretisation).

Manifold valued outputs. So far we have defined our generative model with Y = R d , 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 a compact manifold, we replace the Ornstein-Uhlenbeck process by a Brownian motion which targets the uniform distribution.

Training. As the reverse SDE (4) involves the preconditioned score K( x, x ) ∇ log p t , we directly approximate it with a neural network (Ks) θ : [0 , T ] ×X n ×Y n → T Y n , where T Y is the tangent bundle of Y ,see App. C. The conditional score of the noising process (3) is given by

since Y t = m t | 0 +Σ 1 / 2 t | 0 ε with ε ∼ N(0 , Id) , and Σ t | 0 = σ 2 t | 0 K with σ t | 0 = (1 -e -∫ t 0 β ( s )d s ) 1 / 2 , 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 Λ( t ) = σ 2 t | 0 K ⊤ K

where ∥ x ∥ 2 Λ = x ⊤ Λ x . Note that when targeting a unit-variance white noise, then K = Id and the loss (6) reverts to the DSM loss with weighting λ ( t ) = 1 /σ 2 t | 0 (Song et al., 2021). In App. B.2, we explore several preconditioning terms and associated weighting Λ( t ) . Overall, we found the preconditioned score K ∇ log p t parameterisation, in combination with the ℓ 2 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 a group G , we aim to build a generative model over steerable feature fields as defined in Sec. 2.

Invariant process. A stochastic process f ∼ µ is said to be G -invariant if µ ( g · A ) = µ ( A ) for any g ∈ G , with µ ∈ P (C( X , Y )) , where P is the space of probability measure on the space of continuous functions and A ⊂ C( X , Y ) measurable. From a sample perspective, this means that with inputoutput pairs C = { ( x i , y i ) } n i =1 , and denoting the action of G on this set as g ·C ≜ { ( g · x i , ρ ( g ) y i ) } n i =1 , f ∼ µ is G -invariant if and only if g · C has the same distribution as C . In what follows, we aim to derive sufficient conditions on the model introduced in Sec. 3 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 a Gaussian process GP( m,k ) is G -invariant if and only if its mean m and covariance k are suitably G -equivariant-that is, for all x, x ′ ∈ X , g ∈ G

Trivial examples of E( n ) -equivariant kernels include diagonal kernels k = k 0 Id with k 0 invariant (Holderrieth et al., 2021), 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 if 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 t ( x )) x ∈X ,t ∈ [0 ,T ] the process induced by the time-reversal SDE (4) where the score is approximated by a score network s θ : [0 , T ] ×X n ×Y n → T Y n , and the limiting process is given by L ( ¯ Y 0 ) = GP( m,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. s θ ( t, g · x, ρ ( g ) y ) = ρ ( g ) s θ ( t, x, y ) for all x ∈ X , g ∈ G , then for any t ∈ [0 , T ] the process ( ¯ Y t ( x )) x ∈X is G -invariant.

This result can be proved in two ways, from the probability flow ODE perspective or directly in terms of SDE via Fokker-Planck, see App. D.2. In particular, when modelling an invariant scalar data process ( Y 0 ( x )) x ∈X such as a temperature field, we need the score network to admit the invariance constraint s θ ( t, g · x, y ) = s θ ( t, x, y ) .

Equivariant conditional process. Often precedence is given to modelling the predictive process

2

0

2

f

(

x

∗

)

|C

0

2

4

f

(

x

∗

)

|

g

· C

f

(

x

∗

)

|C

f

(

x

∗

)

|

g

· C

-

given a set of observations C = { ( x c , y c ) } c ∈ C . 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 µ given a context C is said to be conditionally G -equivariant if the conditional satisfies µ ( A | g · C ) = µ ( g · A |C ) for any g ∈ G and A ∈ C( X , Y ) measurable, as illustrated in Fig. 2.

Proposition 3.3. Equivariant conditional process. Assume a stochastic process f ∼ µ is G -invariant. Then the conditional process f |C given a set of observations C is G -equivariant.

Originally stated in Holderrieth et al. (2021) in the case where the process is over functions of the form f : R n → R d 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 sampling in diffusion models such as: replacement sampling, amortisation and conditional guidance, which we discuss 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 written as ∇ x log p ( x | y ) = ∇ x log p ( x, y ) - ∇ x log p ( y ) = ∇ x log p ( x, y ) we can therefore, for any point in the diffusion time, conditionally sample using Langevin dynamics, following the SDE d Y s = 1 2 K ∇ log p T -t ( Y s )d s + √ Kd B s , by only applying the diffusion to the variables

p

0

L

( ¯

Y

0

) =

p

T

L

( ¯

Y

0

γ

)

L

( ¯

Y

1

γ

)

· · ·

p

γ

L

( ¯

Y

0

2

γ

)

L

( ¯

Y

1

2

γ

)

p

2

γ

of interest and holding the others fixed. While we could sample directly at 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. In 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 al. (2021), we can derive a 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 is learnt, we can thus use it 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 i , y i } i ∈ C ) given a context set { x i , y i } i ∈ C . A simple approach is to simply rely on the conditional probability rule evaluate p ( y ∗ | x ∗ , { x i , y i } i ∈ C ) = p ( y ∗ , { y i } i ∈ C | x ∗ , { x i } i ∈ C ) /p ( { y i } i ∈ C |{ x i } i ∈ 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.

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 are disadvantaged by the fact that variables early in the auto-regressive generation only have simple distributions (typically Gaussian). Finally, (Dupont et al., 2022) model weights of implicit neural representation using diffusion models.

Stationary stochastic processes. The most popular Gaussian process kernels (e.g. squared exponential, Matérn) 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 (Holderrieth et al., 2021), allowing for modelling higher order tensor fields, such as 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 a large class of non-Euclidean spaces, in particular 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 a 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 finite-dimensional 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 can also be seen as 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 on soft and cold denoising diffusion models (Daras et al., 2022; Bansal et al., 2022; Hoogeboom and Salimans, 2022). Several recent work explicitly introduced a 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.

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 et al., 2023; Dutordoir et al., 2022; Phillips et al., 2022). Most of these works are based on an extension of the diffusion models techniques (Song et al., 2021; Ho et al., 2020) to the infinite-dimensional space, leveraging tools from the Cameron-Martin theory such as the Feldman-Hájek theorem (Kerrigan et al., 2022; Pidstrigach et al., 2023) to define infinite-dimensional Gaussian measures and how they interact. We refer to (Da Prato and Zabczyk, 2014) for a thorough introduction to Stochastic Differential Equations in infinite dimension. (Phillips et 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 is necessary for the theory of infinite dimensional SDE (Da Prato and Zabczyk, 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.

-

2

.

5

0

.

0

2

.

5

SE

Data Prior

Model Posterior

True Posterior

2

1

0

1

2

-

0

.

5

0

.

0

0

.

5

Sawtooth

2

1

0

1

2

2

1

0

1

2

-

-

-

-

-

-

5 Experimental results

5.1 1D regression over stationary scalar fields

We evaluate GEOMNDPS on several synthetic 1D regression datasets. We follow the same experimental setup as Bruinsma et al. (2020) which we detail in App. F.1. In short, it contains Gaussian (Squared Exponential (SE), MATÉRN ( 5 2 ) , WEAKLY PERIODIC) and non-Gaussian (SAWTOOTH and MIXTURE) sample paths, where MIXTURE is a combination of the other four datasets with equal weight. Fig. 9 shows samples for each of these dataset. The Gaussian datasets are corrupted with observation noise with variance σ 2 = 0 . 05 2 . Table 1 reports the average log-likelihood p ( y ∗ | x ∗ , C ) across 4096 test samples, where the context set size is uniformly sampled between 1 and 10 and the target has fixed size of 50 . All inputs x c , x ∗ are chosen uniformly within their input domain which is [ -2 , 2] for the training data and 'interpolation' evaluation and [2 , 6] for the 'generalisation' evaluation.

We compare the performance of GeomNDP to a GP with the true hyperparameters (when available), a (convolutional) Gaussian NP (Bruinsma et al., 2020), a convolutional NP (Gordon et al., 2020) and a vanilla attention-based NDP (Dutordoir et al., 2022) which we reformulated in the continuous diffusion process framework to allow for log-likelihood evaluations and thus a fair comparisondenoted NDP * . We enforce translation invariance in the score network for GeomNDP by subtracting the centre of mass from the input x , 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 on the nonGaussian datasets as it is restricted by its conditional Gaussian assumption, whilst NDPS methods still performs well as illustrated on Fig. 4. In the bottom rows of Table 1, we assess the models ability to generalise outside of the training input range x ∈ [ -2 , 2] , and evaluate them on a translated grid where context and target points are sampled from [2 , 6] . Only convolutional NPs (GNP and CONVNP) and T(1) -GEOMNDP are able to model stationary processes and therefore to perform as well as in the interpolation task. The NDP * , on the contrary, drastically fails at this task.

MODEL

SE

CURL-FREE

DIV-FREE

GP

0

.

56

±

0

.

00

0

.

66

±

0

.

00

0

.

66

±

0

.

00

NDP *

0

.

55

±

0

.

00

0

.

62

±

0

.

01

0

.

62

±

0

.

01

E(2)

-GEOMNDP

0

.

56

±

0

.

01

0

.

65

±

0

.

01

0

.

66

±

0

.

01

GP (DIAG.)

-

1

.

56

±

0

.

00

-

1

.

47

±

0

.

00

-

1

.

47

±

0

.

00

T(2)

-CONVCNP

-

1

.

71

±

0

.

01

-

1

.

77

±

0

.

01

-

1

.

76

±

0

.

00

E(2)

-STEERCNP

-

1

.

61

±

0

.

00

-

1

.

57

±

0

.

00

-

1

.

57

±

0

.

01

10

1

10

2

10

3

10

4

10

5

Number of training samples

0

.

5

0

.

6

Log-likelihood

GP

NDP

*

E(2)-GeomNDP

Non white kernels for limiting process. The NDP methods in the above experiment target the white kernel ✶ ( x = 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 a description of these. We observe that although choosing such kernels gives a head start to the training, it eventually yield slightly worse performance. We attribute this to the additional complexity of learning a non-diagonal covariance. Finally, across all datasets and limiting kernels, we found the preconditioned score K ∇ log p t 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 we compute the Kullback-Leibler 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 y c has relatively less impact on the overall outcome.

5.2 Regression over Gaussian process vector fields

We now focus our attention to modelling equivariant vector fields. For this, we create datasets using samples from a two-dimensional zero-mean GP with one of the following E(2) -equivariant kernels: a diagonal Squared-Exponential (SE) kernel, a zero curl (CURL-FREE) kernel and a zero divergence (DIV-FREE) kernel, as described in App. D.1.

We equip our model, GeomNDP, with a E(2) -equivariant score architecture, based on steerable CNNs (Thomas et al., 2018; Weiler and Cesa, 2021). We compare to NDP * with a non-equivariant attention-based network (Dutordoir et al., 2022). We also evaluate two neural processes, a translationequivariant CONVCNP (Gordon et al., 2020) and a C4 ⋉ R 2 ⊂ E(2) -equivariant STEERCNP (Holderrieth 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-likelihood of the data process samples under the model on a held-out test dataset.

We observe in Fig. 5 (Left), that the CNPs performance is limited by their diagonal predictive covariance assumption, and as such cannot do better than the GP (DIAG.). We also see that although NDP * is able to fit well GP posteriors, it does not reach the maximum log-likelihood value attained by the data GP, in contrast to its equivariant counterpart GEOMNDP . To further explore gains brought by the built-in equivariance, we explore the data-efficiency in Fig. 5 (Right), and notice that E(2)-GEOMNDP 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 a task where the domain of the stochastic process is a non-Euclidean manifold. We model the trajectories of cyclones over the earth, modelled as sample paths of the form R → S 2 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.

(a) Interpolation

(b) Extrapolation

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 2 the likelihood and MSE for a 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 → R 2 , one on latitude/longitude coordinates, the other via a stereographic projection, using a diagonal RBF kernel with hyperparameters fitted with maximum likelihood) fail drastically given the high non-Gaussianity 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 a 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 at 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 a white noise process appears to over-perform other Gaussian processes. In future 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.

References

Azangulov, A. Smolensky, A. Terenin, and V. Borovitskiy. Stationary Kernels and Gaussian Processes on Lie Groups and their Homogeneous Spaces I: the Compact Case, Aug. 2022. DOI: 10.48550/arXiv.2208.14960 . Cited on pages 2, 6.

Azangulov, A. Smolensky, A. Terenin, and V. Borovitskiy. Stationary Kernels and Gaussian Processes on Lie Groups and their Homogeneous Spaces II: non-compact symmetric spaces, Apr. 2023. DOI: 10.48550/arXiv.2301.13088 . Cited on pages 2, 6.

Bansal, E. Borgnia, H.-M. Chu, J. S. Li, H. Kazemi, F. Huang, M. Goldblum, J. Geiping, and T. Goldstein. Cold diffusion: inverting arbitrary image transforms without noise. arXiv preprint arXiv:2208.09392 , 2022. Cited on page 6.

Biloš, K. Rasul, A. Schneider, Y. Nevmyvaka, and S. Günnemann. Modeling Temporal Data as Continuous Functions with Process Diffusion, 2022. URL: https://openreview.net/forum? id=VmJKUypQ8wR . Cited on page 6.

Bond-Taylor and C. G. Willcocks. Infinite-diff: infinite resolution diffusion with subsampled mollified states. arXiv preprint arXiv:2303.18242 , 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://github.com/google/jax . Cited on pages 10, 13.

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://openreview.net/forum? id=OAsXFPBfTBh . Cited on page 6.

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 7, 14, 16.

Cattiaux, G. Conforti, I. Gentil, and C. Léonard. Time reversal of diffusion processes under a finite entropy condition. arXiv preprint arXiv:2104.07708 , 2021. Cited on pages 2, 3, 5.

Chen, S. Chewi, J. Li, Y. Li, A. Salim, and A. R. Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. arXiv preprint arXiv:2209.11215 , 2022. Cited on page 11.

Cohen. Equivariant convolutional networks . PhD thesis, 2021. Cited on page 4.

Da Prato and J. Zabczyk. Stochastic equations in infinite dimensions . Cambridge university press, 2014. Cited on page 6.

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: https://doi.org/10.1111/rssb.12183 . Cited on page 13.

Daras, M. Delbracio, H. Talebi, A. G. Dimakis, and P. Milanfar. Soft diffusion: score matching for general corruptions. arXiv preprint arXiv:2209.05442 , 2022. Cited on page 6.

De Bortoli. Convergence of denoising diffusion models under the manifold hypothesis. arXiv preprint arXiv:2208.05314 , 2022. Cited on pages 10, 11.

De Bortoli, E. Mathieu, M. Hutchinson, J. Thornton, Y. W. Teh, and A. Doucet. Riemannian score-based generative modeling. arXiv preprint arXiv:2202.02763 , 2022. Cited on pages 1, 4.

De Bortoli, J. Thornton, J. Heng, and A. Doucet. Diffusion Schrödinger bridge with applications to score-based generative modeling. In Advances in Neural Information Processing Systems , 2021. Cited on pages 5, 18.

Dhariwal and A. Nichol. Diffusion models beat GAN on image synthesis. arXiv preprint arXiv:2105.05233 , 2021. Cited on page 1.

Dupont, H. Kim, S. Eslami, D. Rezende, and D. Rosenbaum. From data to functa: your data point is a function and you should treat it like one. arXiv preprint arXiv:2201.12204 , 2022. Cited on page 6.

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.

Durmus and É. Moulines. Nonasymptotic convergence analysis for the unadjusted Langevin algorithm. The Annals of Applied Probability , 27(3):1551-1587, 2017. ISSN: 1050-5164. DOI: 10.1214/16-AAP1238 . URL: https://doi.org/10.1214/16-AAP1238 . Cited on pages 2, 13.

Dutordoir, A. Saul, Z. Ghahramani, and F. Simpson. Neural Diffusion Processes, June 2022. DOI: 10.48550/arXiv.2206.03992 . Cited on pages 1, 3, 6-8, 14, 16-18.

Fishman, L. Klarner, V. D. Bortoli, E. Mathieu, and M. Hutchinson. Diffusion models for constrained domains, 2023. arXiv: 2304.05364 [cs.LG] . Cited on page 1.

Franzese, S. Rossi, D. Rossi, M. Heinonen, M. Filippone, and P. Michiardi. Continuous-time functional diffusion processes. arXiv preprint arXiv:2303.00800 , 2023. Cited on pages 1, 6.

Garnelo, D. Rosenbaum, C. Maddison, T. 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 Proceedings of Machine Learning Research , pages 1704-1713. PMLR, July 2018. URL: https: //proceedings.mlr.press/v80/garnelo18a.html . Cited on page 6.

Garnelo, J. Schwarz, D. Rosenbaum, F. Viola, D. J. Rezende, S. Eslami, and Y. W. Teh. Neural processes. arXiv preprint arXiv:1807.01622 , 2018. Cited on pages 2, 6.

Geiger and T. Smidt. E3nn: euclidean neural networks, 2022. DOI: 10.48550/ARXIV.2207. 09453 . URL: https://arxiv.org/abs/2207.09453 . Cited on page 16.

Gordon, W. P. Bruinsma, A. Y. K. Foong, J. Requeima, Y. Dubois, and R. E. Turner. Convolutional conditional neural processes. In International Conference on Learning Representations , 2020. URL: https://openreview.net/forum?id=Skey4eBYPS . Cited on pages 2, 6-8, 17.

Guth, S. Coste, V. De Bortoli, and S. Mallat. Wavelet score-based generative modeling. arXiv preprint arXiv:2208.05003 , 2022. Cited on page 6.

Hagemann, L. Ruthotto, G. Steidl, and N. T. Yang. Multilevel diffusion: infinite dimensional score-based diffusion models for image generation. arXiv preprint arXiv:2303.04772 , 2023. Cited on page 6.

R. Harris, K. J. Millman, S. J. van der Walt, R. Gommers, P. Virtanen, D. Cournapeau, E. Wieser, J. Taylor, S. Berg, N. J. Smith, R. Kern, M. Picus, 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, H. Abbasi, C. Gohlke, and T. E. Oliphant. Array programming with NumPy. Nature , 585(7825):357-362, Sept. 2020. Cited on page 13.

G. Haussmann and E. Pardoux. Time reversal of diffusions. The Annals of Probability , 14(4):11881205, 1986. Cited on page 3.

Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 2020. Cited on pages 1, 6.

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.

Holderrieth, M. J. Hutchinson, and Y. W. Teh. 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.

Hoogeboom and T. Salimans. Blurring diffusion models. arXiv preprint arXiv:2209.05557 , 2022. Cited on page 6.

C.-W. Huang, M. Aghajohari, J. 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.

D. Hunter. Matplotlib: a 2d graphics environment. Computing in Science & Engineering , 9(3):9095, 2007. Cited on page 13.

Hutchinson, A. Terenin, V. Borovitskiy, S. Takao, Y. Teh, and M. Deisenroth. Vector-valued gaussian processes on riemannian manifolds via gauge independent projected kernels. Advances in Neural Information Processing Systems , 34:17160-17169, 2021. Cited on page 5.

Ingraham, M. Baranov, Z. Costello, V. Frappier, A. Ismail, S. Tie, W. Wang, V. Xue, F. Obermeyer, A. Beam, and G. Grigoryan. Illuminating protein space with a programmable generative model. bioRxiv , 2022. DOI: 10.1101/2022.12.01.518682 . eprint: https://www.biorxiv.org/ content/early/2022/12/02/2022.12.01.518682.full.pdf . URL: https://www. biorxiv.org/content/early/2022/12/02/2022.12.01.518682 . Cited on page 6.

Jha, D. Gong, X. Wang, R. E. Turner, and L. Yao. The neural process family: survey, applications and perspectives. arXiv preprint arXiv:2209.00517 , 2022. Cited on page 6.

Jing, G. Corso, R. Berlinghieri, and T. Jaakkola. Subspace diffusion generative models. arXiv preprint arXiv:2205.01490 , 2022. Cited on page 6.

Kallenberg. Foundations of Modern Probability . Probability Theory and Stochastic Modelling. Springer, 2021. ISBN: 978-3-030-61872-8. URL: https://books.google.co.uk/books?id= 6hJezgEACAAJ . Cited on page 2.

Karras, M. Aittala, T. Aila, and S. Laine. Elucidating the Design Space of Diffusion-Based Generative Models, Oct. 2022. DOI: 10.48550/arXiv.2206.00364 . Cited on page 3.

Kerrigan, J. Ley, and P. Smyth. Diffusion Generative Models in Infinite Dimensions, 2022. DOI: 10.48550/arXiv.2212.00886 . Cited on pages 1, 3, 6.

Kim, A. Mnih, J. Schwarz, M. Garnelo, A. Eslami, D. Rosenbaum, O. Vinyals, and Y. W. Teh. Attentive neural processes. In International Conference on Learning Representations , 2019. URL: https://openreview.net/forum?id=SkE6PjC9KX . Cited on page 6.

P. Kingma and J. Ba. Adam: A Method for Stochastic Optimization. arXiv:1412.6980 [cs] , 2015. URL: http://arxiv.org/abs/1412.6980 . Cited on page 17.

J. Knapp Kenneth R. Diamond, J. P. Kossin, M. C. Kruk, and C. J. I. Schreck. International Best Track Archive for Climate Stewardship (IBTrACS) Project, Version 4. Technical report, NOAA National Centers for Environmental Information, 2018. DOI: https://doi.org/10.25921/ 82ty-9e16 . Cited on pages 8, 17.

R. Knapp, M. C. Kruk, D. H. Levinson, H. J. Diamond, and C. J. Neumann. The international best track archive for climate stewardship (ibtracs): unifying tropical cyclone data. Bulletin of the American Meteorological Society , 91(3):363-376, 2010. DOI: https://doi.org/10.1175/ 2009BAMS2755.1 . Cited on pages 8, 17.

Köhler, L. Klein, and F. Noé. Equivariant Flows: exact likelihood generative learning for symmetric densities. arXiv:2006.02425 [physics, stat] , June 2020. URL: http://arxiv.org/abs/2006. 02425 . Cited on page 6.

LeCun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied to document recognition. In Proceedings of the IEEE , volume 86 of number 11, pages 2278-2324, 1998. URL: http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.42.7665 . Cited on page 17.

Lee, J. Lu, and Y. Tan. Convergence of score-based generative modeling for general data distributions. In International Conference on Algorithmic Learning Theory , pages 946-985. PMLR, 2023. Cited on page 11.

H. Lim, N. B. Kovachki, R. Baptista, C. Beckham, K. Azizzadenesheli, J. 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, J. 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 3, 6.

Lou and S. Ermon. Reflected diffusion models. arXiv preprint arXiv:2304.04740 , 2023. Cited on page 1.

Louizos, X. Shi, K. Schutte, and M. Welling. The functional neural process. Advances in Neural Information Processing Systems , 32, 2019. Cited on page 6.

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-11471, 2022. Cited on pages 5, 8-11.

Macêdo and R. Castro. Learning Divergence-Free and Curl-Free Vector Fields with Matrix-Valued Kernels . Pré-Publicações / A.: Pré-publicações. IMPA, 2010. Cited on pages 4, 6.

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.

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: a python package for riemannian geometry in machine learning. Journal of Machine Learning Research , 21(223):1-9, 2020. URL: http://jmlr.org/papers/v21/19-027.html . Cited on page 10.

Øksendal. Stochastic differential equations. In Stochastic differential equations , pages 65-84. Springer, 2003. Cited on page 4.

Papamakarios, E. Nalisnick, D. J. Rezende, S. Mohamed, and B. Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. arXiv preprint arXiv:1912.02762 , 2019. Cited on page 6.

Phillips, T. Seror, M. Hutchinson, V. De Bortoli, A. Doucet, and E. Mathieu. Spectral Diffusion Processes, Nov. 2022. URL: http://arxiv.org/abs/2209.14125 . Cited on pages 3, 6, 1, 4, 5, 8.

Pidstrigach, Y. Marzouk, S. Reich, and S. Wang. Infinite-Dimensional Diffusion Models for Function Spaces, Feb. 2023. URL: http://arxiv.org/abs/2302.10130 . Cited on pages 1, 3, 6.

Pollard. A User's Guide to Measure Theoretic Probability . Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2002. ISBN: 978-0-521-00289-9. URL: https://books.google.co.uk/books?id=B7Ch-c2G21MC . Cited on page 5.

Rahimi and B. Recht. Random features for large-scale kernel machines. Advances in neural information processing systems , 20, 2007. Cited on page 2.

E. Rasmussen. Gaussian processes in machine learning. In Summer school on machine learning , pages 63-71. Springer, 2003. Cited on pages 2, 6.

Rissanen, M. Heinonen, and A. Solin. Generative Modelling With Inverse Heat Dissipation, July 2022. DOI: 10.48550/arXiv.2206.13397 . Cited on page 6.

O. Roberts and R. L. Tweedie. Exponential convergence of Langevin distributions and their discrete approximations. Bernoulli , 2(4):341-363, 1996. ISSN: 1350-7265. DOI: 10.2307/3318418 . URL: https://doi.org/10.2307/3318418 . Cited on page 13.

Saharia, J. Ho, W. Chan, T. Salimans, D. J. Fleet, and M. Norouzi. Image super-resolution via iterative refinement. arXiv preprint arXiv:2104.07636 , 2021. Cited on page 6.

Särkkä and A. Solin. Applied Stochastic Differential Equations . Cambridge University Press, 1st edition, 2019. ISBN: 978-1-108-18673-5 978-1-316-51008-7 978-1-316-64946-6. DOI: 10. 1017/9781108186735 . Cited on page 1.

Scott and J. Serre. Linear Representations of Finite Groups . Graduate Texts in Mathematics. Springer New York, 1996. ISBN: 978-0-387-90190-9. URL: https://books.google.co.uk/ books?id=NCfZgr54TJ4C . Cited on page 3.

Singh, J. Yoon, Y. Son, and S. Ahn. Sequential neural processes. Advances in Neural Information Processing Systems , 32, 2019. Cited on page 6.

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, 5, 8.

Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations , 2021. Cited on pages 1, 2, 4-6, 8.

W. Stroock and S. S. Varadhan. Multidimensional diffusion processes . Springer, 2007. Cited on pages 11, 13.

Thomas, T. Smidt, S. Kearnes, L. Yang, L. Li, K. Kohlhoff, and P. Riley. Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds. arXiv:1802.08219 [cs] , May 2018. URL: http://arxiv.org/abs/1802.08219 . Cited on pages 8, 16.

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.press/v5/titsias09a.html . Cited on page 2.

L. Trippe, J. Yim, D. Tischer, T. Broderick, D. Baker, R. Barzilay, and T. Jaakkola. Diffusion probabilistic modeling of protein backbones in 3D for the motif-scaffolding problem, June 2022. DOI: 10.48550/arXiv.2206.04119 . Cited on page 8.

Van Rossum and F. L. Drake Jr. Python reference manual . Centrum voor Wiskunde en Informatica Amsterdam, 1995. Cited on page 13.

Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems , 2017. Cited on page 16.

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.

Virtanen, R. Gommers, T. E. Oliphant, M. Haberland, T. Reddy, D. Cournapeau, E. Burovski, P. Peterson, W. Weckesser, J. Bright, S. J. van der Walt, M. Brett, J. Wilson, K. J. Millman, N. Mayorov, A. R. J. Nelson, E. Jones, R. Kern, E. Larson, C. J. Carey, ˙ I. Polat, Y . Feng, E. W. Moore, J. VanderPlas, D. 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

1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods , 17:261-272, 2020. Cited on page 13.

Voleti, C. Pal, and A. M. Oberman. Score-based Denoising Diffusion with Non-Isotropic Gaussian Noise Models, 2022. URL: https://openreview.net/forum?id=igC8cJKcb0Q . Cited on page 6.

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, 7.

Weiler and G. Cesa. General $E(2)$-Equivariant Steerable CNNs. arXiv:1911.08251 [cs, eess] , Apr. 2021. URL: http://arxiv.org/abs/1911.08251 . Cited on pages 8, 17.

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-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 31. Curran Associates, Inc., 2018. URL: https://proceedings.neurips.cc/ paper/2018/file/488e4104520c6aab692863cc1dba45af-Paper.pdf . Cited on pages 16, 17.

Wilson, V. Borovitskiy, A. Terenin, P. Mostowsky, and M. Deisenroth. Efficiently sampling functions from gaussian process posteriors. In ICML , pages 10292-10302. PMLR, 2020. Cited on page 2.

Yadan. Hydra - a 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 arXiv:2302.02277 , 2023. Cited on pages 4, 7.

Supplementary to: Geometric Neural Diffusion Processes

A Organisation of appendices

In this supplementary, we first introduce in App. B an Ornstein 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 to yield a group invariant process. What's more in App. E, we study some conditional sampling schemes. Eventually in App. F, we give a thorough description of experimental settings along with additional empirical results.

B Ornstein Uhlenbeck on function space

B.1 Multivariate Ornstein-Uhlenbeck process

First, we aim to show that we can define a stochastic process on an infinite dimensional function space, by defining the joint finite marginals Y ( x ) as the solution of a multidimensional Ornstein-Uhlenbeck process. In particular, for any set of input x = ( x 1 , · · · , x k ) ∈ X k , we define the joint marginal as the solution of the following SDE

Proposition B.1. (Phillips et al., 2022) We assume we are given a data process ( Y 0 ( x )) x ∈X and we denote by G ∼ GP(0 , k ) a Gaussian process with zero mean and covariance. Then let's define

Then ( Y t ( x )) x ∈X is a stochastic process (by virtue of being a linear combination of stochastic processes). We thus have that Y t a.s. - - - → t → 0 Y 0 and Y t a.s. - - - → t →∞ Y ∞ with Y ∞ ∼ GP( m,k ) , so effectively ( Y t ( x )) t ∈ R + ,x ∈X interpolates between the data process and this limiting Gaussian process. Additionally, L ( Y t | Y 0 = y 0 ) = GP( m t , K t ) with m t = e -1 2 · ∫ t s =0 β s ds y 0 + ( 1 -e -1 2 · ∫ t s =0 β s ds ) m and Σ t = ( 1 -e -∫ t s =0 β s ds ) K . Furthermore, ( Y t ( x )) t ∈ R + ,x ∈X is the solution of the SDE in (8) .

Proof. We aim to compute the mean and covariance of the process ( Y t ) t ≥ 0 described by the SDE (3). First let's recall the time evolution of the mean and covariance of the solution from a multivariate Ornstein-Uhlenbeck process given by

We know that the time evolution of the mean and the covariance are given respectively by Särkkä and Solin (2019)

Plugging in the drift f ( Y t , t ) = 1 / 2 · ( m -Y t ) β t and diffusion term L ( Y t , t ) = √ β t K from (3), we get

Solving these two ODEs we get

with m 0 ≜ E[ Y 0 ] and Σ 0 ≜ Cov[ Y 0 ] .

Now let's compute the first two moments of ( Y t ( x )) x ∈X . We have

B.2 Conditional score

Hence, condition on Y 0 the score is the gradient of the log Gaussian characterised by mean m t | 0 = e -1 2 B ( t ) Y 0 and Σ t | 0 = (1 -e -B ( t ) )K with B ( t ) = ∫ t 0 β ( s ) ds which can be derived from the above marginal mean and covariance with m 0 = Y 0 and Σ 0 = 0 .

where L t | 0 denotes the Cholesky decomposition of Σ t | 0 = L t | 0 L ⊤ t | 0 , and Y t = m t | 0 +L t | 0 ϵ .

Then we can plugin our learnt (preconditioned) score into the backward SDE 4 which gives

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 D θ for the network instead of s θ as mentioned in the primary text, as it allows us to 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 c out , c skip : R → R , to parameterise the network

We recall that the input to the network is time t , and the noised vector Y t = µ t | 0 + n , where µ t | 0 = e -B ( t ) / 2 Y 0 and n ∼ N (0 , Σ t | 0 ) with Σ t | 0 = (1 -e -B ( t ) ) K . The gram matrix K corresponds to k ( X,X ) with k the limiting kernel. We denote by L t | 0 and S respectively the Cholesky decomposition of Σ t | 0 = L t | 0 L ⊤ t | 0 and K = SS ⊤ .

The denoising score matching loss weighted by Λ( t ) is given by

No preconditioning By reparametrisation, let Y t = µ t | 0 + L t | 0 z , where z ∼ N ( 0 , I ) , the loss from Eq. (30) can be written as

Choosing Λ( t ) = Σ t | 0 = L t | 0 L ⊤ t | 0 we obtain

(34)

Preconditioning by K Alternatively, one can train the neural network to approximate the preconditioned score D θ ≈ K ∇ Y t log p t ( Y t | Y 0 ) . The loss, weighted by Λ = σ 2 t | 0 I , is then given by

Precondition by S ⊤ A variation of the previous one, is to precondition the score by the transpose Cholesky of the limiting kernel gram matrix, such that D θ ≈ S ⊤ ∇ Y t log p t ( Y t | Y 0 ) .

The loss, weighted by Λ = σ 2 t | 0 I , becomes

Predicting Y 0 Finally, an alternative strategy is to predict Y 0 from a noised version Y t . In this case, the loss takes the simple form

The score can be computed from the network's prediction following

(46)

Table 3 summarises the different options for parametrising the score as well as the values for c skip and c out 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 K parametrisation to produce the best results, but we refer to App. F.1.3 for a more in-depth ablation study.

B.4 Exact (marginal) score in Gaussian setting

Interpolating between Gaussian processes GP ( m 0 , Σ 0 ) and GP ( m, K)

(51)

with Σ t = K + e -∫ t s =0 β s ds (Σ 0 -K) = L t L t ⊤ obtained via Cholesky decomposition.

B.5 Langevin dynamics

Under mild assumptions on ∇ log p T -t (Durmus and Moulines, 2016) the following SDE

admits a solution ( Y s ) s ≥ 0 whose law L ( Y s ) converges with geometric rate to p T -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

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 we have defined a generative model over functions via its finite marginals ¯ Y θ T ( x ) . 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. s θ ( t, σ ◦ x, σ ◦ y ) = σ ◦ s θ ( t, x, y ) for any σ ∈ Σ n . Additionally, we have that the noising process ( Y t ( x )) x ∈X is trivially consistent for any t ∈ R + since it is a stochastic process as per Prop. B.1, and consequently so is the (true) time-reversal ( ¯ Y t ( x )) x ∈X . Yet, when approximating the score s θ ≈ ∇ log p t , we lose the consistency over the generative process ¯ Y θ t ( x ) as the constraint on the score network is non trivial to satisfy. This is actually a really strong constraint on the model class, and as soon as one goes beyond linearity (of the posterior w.r.t. the context set), it is non trivial to enforce without directly parameterising a stochastic process, e.g. as Phillips et al. (2022). There thus seems to be a strong trade-off between satisfying consistency, and the model's ability to fit complex process and scale to large datasets.

C Manifold-valued diffusion process

C.1 Manifold-valued inputs

In the main text we dealt with a simplified case of tensor fields where the tensor fields are over Euclidean space. Nevertheless, it is certainly possible to apply our methods to these settings. Significant work has been done on performing convolutions on feature fields on generic manifolds (a superset of tensor fields on generic manifolds), core references being (Cohen, 2021) for the case of homogeneous spaces and (Weiler et al., 2021) for more general Riemannian manifolds. We

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 a manifold M . We shall denote such a bundle BM and the space of sections Γ( BM ) . The goal, therefore, is to model random elements from this space of sections. For a 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 a tuple ( m,b ) , m in the base manifold and b in the typical fibre, alongside the condition that the composition of the projection proj i : ( m,b ) ↦→ m with the section is the identity, proj i ◦ σ = Id it is clear we can see distribution over sections as stochastic processes with 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.e. a finite set of points in the manifold, is defined as π m 1 ,...,m n ( σ ) = ( σ ( m 1 ) , ..., σ ( m n )) .

Noising process. To define a 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 on 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 is a valid process on full sections though should one want to try and parameterise a score function on the whole section akin to some other infinite-dimension diffusion models.

Score function. The last thing to 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 a distinct advantage in simplicity for our approach. All we need to do is pick a 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 t ( x ) = ( Y t ( x 1 ) , · · · , Y t ( x n )) ∈ M n , things are less trivial as manifolds do not have a vector space structure which is necessary to define Gaussian processes. Fortunately, We can still target a know distribution marginally independently on 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

are applied to each marginal. Hence in the limit t →∞ , Y t ( x ) has density (assuming it exists) which factors as dp/d Vol M (( y ( x 1 ) , · · · , y ( x n ))) ∝ ∏ k e -U ( y ( x n ))) . For compact manifolds, we can target the uniform distribution by setting U ( x ) = 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 Invariant neural diffusion processes

D.1 E( n ) -equivariant kernels

A kernel k : R d × R d → R d × d is equivariant if it satisfies the following constraints: (a) k is stationnary , that is if for all x, x ′ ∈ R n

and if (b) it satisfies the angular constraint for any h ∈ H

A trivial example of such an equivariant kernel is the diagonal kernel k ( x, x ′ ) = k 0 ( x, x ′ )I (Holderrieth et al., 2021), with k 0 stationnary. This kernel can be understood has having d independent Gaussian process uni-dimensional output, that is, there is no inter-dimensional 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 k curl = k 0 A and k div = k 0 B with k 0 stationary, e.g. squared exponential kernel k 0 ( x, x ′ ) = σ 2 exp ( ∥ x -x ′ ∥ 2 2 l 2 ) , and A and B given by

D.2 Proof of Prop. 3.2

Below we give two 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 approximate score s θ ( t, · ) ≈ ∇ log p t is given by

This ODE induces a flow ϕ b t : X n × Y n → T Y n for a given integration time t , which is said to be G -equivariant if the vector field is G -equivariant itself, i.e. b ( t, g · x, ρ ( g ) ¯ Y t ) = ρ ( g ) b ( t, x, ¯ Y t ) . We have that for any g ∈ G

with (1) from the G -invariant prior GP conditions on m and k , (2) assuming that the score network is G -equivariant and (3) assuming that ρ ( g ) ∈ O ( n ) . To prove the opposite direction, we can simply follow these computations backwards. Finally, we know that with a G -invariant probability measure p ref and G -equivariant map ϕ , the pushforward probability measure p -1 ref ◦ ϕ is also G -invariant (Köhler et al., 2020; Papamakarios et al., 2019). Assuming a 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 s θ ( t, · ) ≈ ∇ log p t is given by

As for the probability flow drift b ODE, we have that b SDE is similarly G -equivariant, that is b SDE ( t, g · x, ρ ( g ) ¯ Y t ) = ρ ( g ) b SDE ( t, x, ¯ Y t ) for any g ∈ G . Additionally, we have that diffusion matrix is also G -equivariant as for any g ∈ G we have Σ( t, g · x ) = β t K( g · x, g · x ) = β t ρ ( g )K( x, x ) ρ ( g ) ⊤ = ρ ( g )Σ( t, x ) ρ ( g ) ⊤ since K is the gram matrix of an G -equivariant kernel k .

Additionally assuming that b SDE and Σ are bounded, Yim et al. (Proposition 3.6, 2023) says that the distribution of ¯ Y t is G -invariant, and in in particular L ( ¯ Y 0 ) .

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 A of a manifold M with a G structure. Let Isom GM be the group of G -structure preserving isometries on M . The action of this group on a section of the bundle f ∈ Γ( A ) is given by

(Weiler et al., 2021). Let f ∼ P , P a distribution over the space of section. Let ϕ ▷ P be the law of of ϕ ▷ f . Let µ x = L ( f ( x )) = π x # P , the law of f evaluated at a point, where π x is the canonical projection operator onto the marginal at x , # the pushforward operator in the measure theory sense, x ∈ M and y is in the fibre of the associated bundle. Let µ x ′ ,y x = L ( f ( x ) | f ( x ′ ) = y ′ ) = π x µ x ′ ,y ′ = π x # L ( f | f ( x ′ ) = y ′ ) , the conditional law of the process when given f ( x ′ ) = y ′ .

Assume that the prior is invariant under the action of Isom GM , i.e. that

Then the conditional measures are equivariant, in the sense that

Proof. ∀ A,B test functions, ϕ ∈ Isom GM ,

By invariance this quantity is also equal to

Hence

By the stated invariance ϕ -1 ▷ µ x ′ = µ x ′ , hence

So 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 to 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 or 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 an 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, we will show later that it can in fact be shown that this method is doing a 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 an 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 = [ x c , x ∗ ] , where x c are the points we wish to condition on and x s points we wish to sample, we exploit the following score breakdown:

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 s θ ( t, x ) ≈ ∇ x log p t ( x ) , we could use this to perform Langevin dynamics at t = ϵ , some time very close to 0 . Similar to (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 = T we sample an initialisation of y ∗ based on the reference distribution. Progressing from t = T towards t = ϵ we alternate between running a series of Lavgevin corrector steps to sample from the distribution p t, x ∗ ( y ∗ | y c ) , and a single backwards SDE step to sample from p x ( y t -γ | y t ) with a step size γ . At each inner and outer step, we sample a noised version of the conditioning points y c based forward SDE applying noise to these context points, p t, x c ( y c t | y c ) . For the exactness of this scheme, all that matters is that at the end of the sampling scheme, we are sampling from p x ∗ ( y ∗ | y c ) (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.

Fixed context set

SDE path noise,

every outer step

Fully independent noise,

every outer step

Fully independent noise,

every inner step

Outer step

Inner step

The choices that present themselves are

The initial scheme of sampling context noise from the SDE every inner and outer step.

Only re-sampling the context noise every outer step, and keeping it fixed to this for each inner step associated with the outer step.

Instead of sampling independent marginal noise at each outer step, sampling a single noising trajectory of the context set from the forward SDE and use this as the noise at each time.

Perform no noising at all. Effectively the replacement method with added Langevin sampling.

These are illustrated in Fig. 8. The main trade-off of different schemes is the speed at which the noise can be sampled vs sample diversity. In the Euclidean case, we have a closed form for the evolution of the marginal density of the context point under the forward SDE. In this case sampling the noise at a given time is O (1) cost. On the other hand, in some instances such as nosing SDEs on general manifolds, we have to simulate this noise by discretising the forward SDE. In this case, it is O ( n ) cost, where n is the number of discretisation steps in the SDE. For N outer steps and I inner steps, the complexity of the different noising schemes is compared in 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 1 Conditional sampling with Langevin dynamics.

(a) Alg. 1 and Alg. 2 Repaint from (Lugmayr et al., 2022) are equivalent in a specific setting.

(b) There exists a continuous limit (SDE) for both procedures. This SDE targets a probability density which does not correspond to p ( x t 0 | x c 0 ) .

(c) When t 0 → 0 this probability measure converges to p ( x 0 | x c 0 ) which ensures the correctness of 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).

̸

First, we start by describing the RePaint algorithm (Lugmayr et al., 2022). We consider ( Z 0 k , Z 1 k , Z 2 k ) k ∈ N a sequence of independent Gaussian random variable such that for any k ∈ N , Z 1 k and Z 2 k are d -dimensional Gaussian random variables with zero mean and identity covariance matrix and Z 0 k is a p -dimensional Gaussian random variable with zero mean and identity covariance matrix. We assume that the whole sequence to be inferred is of size d 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 t 0 ∈ [0 , T ] . We consider the chain ( X k ) k ∈ N given by X 0 ∈ R d and for any k ∈ N , we define

where X c k = e -t 0 X c 0 +(1 -e -2 t 0 ) 1 / 2 Z 0 k . Finally, we consider

Note that (69) corresponds to one step of backward SDE integration and (70) corresponds to one step of forward SDE integration. In both cases we have used the exponential integrator, see (De

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 k ∈ N we have

Remarking that ( Z k ) k ∈ N = (( Z 1 k + Z 3 k ) / √ 2) k ∈ N is a family of d -dimensional Gaussian random variables with zero mean and identity covariance matrix, we get that for any k ∈ N

where we recall that X c k = e -t 0 X c 0 +(1 -e -2 t 0 ) 1 / 2 Z 0 k . 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 x t 0 ∈ R d , we denote

̸

We emphasize that b/ 2 = ∇ x t 0 log p ( ·| x c 0 ) . In particular, using Tweedie's identity, we have that for any x t 0 ∈ R d

We introduce the following assumption.

Assumption 1. There exist L , C ≥ 0 , m > 0 such that for any x c t 0 , y c t ∈ R p and x t 0 , y t ∈ R d

Assumption 1 ensures that there exists a unique strong solution to the SDE associated with (72). Note that conditions under which log p t 0 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 γ t ) t ≥ 0 ) γ> 0 the family of processes such that for any k ∈ N and γ > 0 , we have for any t ∈ [ kγ, ( k +1) γ ) , X γ t = (1 -( t -kγ ) /γ ) X γ kγ +( t -kγ ) /γ X γ ( k +1) γ and

where ( Z γ kγ ) k ∈ N ,γ> 0 is a family of independent d -dimensional Gaussian random variables with zero mean and identity covariance matrix and for any k ∈ N , γ > 0 , X c,γ kγ = e -t 0 x c 0 +(1 -e -2 t 0 ) 1 / 2 Z 0 ,γ kγ , where ( Z 0 ,γ kγ ) k ∈ N ,γ> 0 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 t ) t ≥ 0 such that

We recall that b depends on t 0 but t 0 is fixed here. This means that we are at time t 0 in the diffusion and consider a corrector at this stage. The variable t does not corresponds to the backward evolution but to the forward evolution in the corrector stage . Under Assumption 1, (77) admits a 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 n → + ∞ ( X 1 /n t ) t ≥ 0 = ( X t ) t ≥ 0 .

This result is an application of Stroock and Varadhan (2007, Theorem 11.2.3). It explicits what is 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 γ > 0 and x ∈ R d , we define

Note that for any γ > 0 and x ∈ R d , we have

(84)

Lemma E.2. Assume Assumption 1 . Then, we have that for any R,ε > 0 and γ ∈ (0 , 1)

Where we recall that for any x ∈ R d ,

Proof. Let R,ε > 0 and γ ∈ (0 , 1) . Using Assumption 1, there exists C > 0 such that for any x t 0 ∈ R d with ∥ x t 0 ∥ ≤ R , we have ∥∇ x t 0 log p t 0 ([ x t 0 , x c t 0 ]) ∥ ≤ C (1 + ∥ x c t 0 ∥ ) . Since p c t 0 | 0 is Gaussian with zero mean and covariance matrix (1 -e -2 t 0 ) Id , we get that for any p ∈ N , there exists A k ≥ 0 such that for any x t 0 ∈ R d with ∥ x t 0 ∥ ≤ R

Therefore, using this result and the fact that for any s ≥ 0 , e -s ≥ 1 -s , we get that there exists B k ≥ 0 such that for any k, p ∈ N and for any x t 0 ∈ R d with ∥ x t 0 ∥ ≤ R

Therefore, combining this result and the Markov inequality, we get that for any x t 0 ∈ R d with ∥ x t 0 ∥ ≤ R we have

In addition, we have that for any x t 0 ∈ R d with R > 0

We also have that for any x t 0 ∈ R d with R > 0

Combining this result, (91), the fact that lim γ → 0 (4 /γ ) | 1 -e -γ | 2 = 0 and lim γ → 0 (2 /γ ) | e -γ -1 + γ | = 0 , we get that lim γ → 0 sup {∥ Σ γ ( x ) -Σ( x ) ∥ | x ∈ R d , ∥ x ∥ ≤ R } = 0 . Similarly, using (91), (94) and the fact that lim γ → 0 (4 /γ ) | 1 -e -γ | 2 = 0 , we get that lim γ → 0 sup {∥ b γ ( x ) -b ( x ) ∥ | x ∈ R d , ∥ x ∥ ≤ R } = 0 .

We can now conclude the proof of Theorem E.1.

Proof. We have that x ↦→ b ( x ) and x ↦→ Σ( x ) are continuous. Combining this result and Lemma E.2, we conclude the proof upon applying Stroock and Varadhan (2007, Theorem 11.2.3).

Theorem E.1 is a non-quantitative result which states what is the limit chain for the REPAINT procedure. Note that if we do not resample, we get that

Recalling (88), we get that (98) is an amortised version of b cond . Similar convergence results can be derived in this case. Note that it is also possible to obtain quantitative discretization bounds between ( X t ) t ≥ 0 and ( X 1 /n t ) t ≥ 0 under the ℓ 2 distance. These bounds are usually leveraged using the Girsanov theorem (Durmus and Moulines, 2017; Dalalyan, 2017). We leave the study of such bounds for future work.

We also remark that b ( x t 0 ) is not given by ∇ log p t 0 ( x t 0 | x c 0 ) . Denoting U t 0 such that for any x t 0 ∈ R d we have that ∇ U t 0 ( x t 0 ) = -b ( x t 0 ) , under mild integration assumptions. In addition, using Jensen's inequality, we have

Hence, π t 0 with density proportional to x ↦→ exp[ -U t 0 ( x )] defines a valid probability measure.

We make the following assumption which allows us to control the ergodicity of the process ( X t ) t ≥ 0 .

Assumption 2. There exist m > 0 and C ≥ 0 such that for any x t 0 ∈ R d and x c t 0 ∈ R p

The following proposition ensures the ergodicity of the chain ( X t ) t ≥ 0 . It is a direct application of Roberts and Tweedie (1996, Theorem 2.1).

Proposition E.3. Assume Assumption 1 and Assumption 2 . Then, π t 0 is the unique invariant probability measure of ( X t ) t ≥ 0 and lim t → 0 ∥L ( X t ) -π t 0 ∥ TV = 0 , where L ( X t ) is the distribution of X t .

Finally, for any t 0 > 0 , denoting π t 0 the probability measure with density U t 0 given for any x t 0 ∈ R d by

We show that the family of measures ( π t 0 ) t 0 > 0 approximates the posterior with density x 0 ↦→ p ( x 0 | x c 0 ) when t 0 is small enough.

Proposition E.4. Assume Assumption 1 . We have that lim t 0 → 0 π t 0 = π 0 where π 0 admits a density w.r.t. the Lebesgue measure given by x 0 ↦→ p ( x 0 | x c 0 ) .

Proof. This is a direct consequence of the fact that p t | 0 ( ·| x c 0 ) → δ x c 0 .

This last results shows that even though we do not target x t 0 ↦→ p t 0 | 0 ( x t 0 | x c 0 ) using this corrector term, we still target p ( x 0 | x c 0 ) as t 0 → 0 which corresponds to the desired output of the algorithm.

F Experimental details

Models, training and evaluation have been implemented in Jax (Bradbury et al., 2018). We used Python (Van Rossum and Drake Jr, 1995) for all programming, Hydra (Yadan, 2019), Numpy (Harris et al., 2020), Scipy (Virtanen et al., 2020), Matplotlib (Hunter, 2007), and Pandas (McKinney et al., 2010). The code is publicly available at https://github . com/cambridge-mlg/ neural_diffusion_processes .

F.1 Regression 1d

F.1.1 Data generation

We follow the same experimental setup as Bruinsma et al. (2020) to generate the 1d synthetic data. It consists of Gaussian (Squared Exponential (SE), MATÉRN ( 5 2 ) , WEAKLY PERIODIC) and nonGaussian (SAWTOOTH and MIXTURE) sample paths, where MIXTURE is a combination of the other four datasets with equal weight. Fig. 9 shows samples for each of these dataset. The Gaussian datasets are corrupted with observation noise with variance σ 2 = 0 . 05 2 . The left column of Fig. 9 shows example sample paths for each of the 5 datasets.

The training data consists of 2 14 sample paths while the test dataset has 2 1 2 paths. For each test path we sample the number of context points between 1 and 10 , the number of target points are fixed to 50 for the GP datasets and 100 for the non-Gaussian datasets. The input range for the training and interpolation datasets is [ -2 , 2] for both the context and target sets, while for the extrapolation task the context and target input points are drawn from [2 , 6] .

Architecture. For all datasets, except SAWTOOTH, we use 5 bi-dimensional attention layers (Dutordoir et al., 2022) with 64 hidden dimensions and 8 output heads. For SAWTOOTH, we obtained better performance with a wider and shallower model consisting of 2 bi-dimensional attention layers with a hidden dimensionality of 128 . In all experiment, we train the NDP-based models over 300 epochs using a batch size of 256 . Furthermore, we use the Adam optimiser for training with the following learning rate schedule: linear warm-up for 10 epochs followed by a cosine decay until the end of training.

F.1.2 Ablation Limiting Kernels

The test log-likelihoods (TLLs) reported in App. F.1.3 for the NDP models target a white limiting kernel and train to approximate the preconditioned score K ∇ log p t . Overall, we found this to be the best performing setting. App. F.1.3 shows an ablation study for different choices of limiting kernel and score parametrisation. We refer to Table 3 for a detailed derivation of the score parametrisations.

The dataset in the top row of the figure originates from a Squared Exponential (SE) GP with lengthscale ℓ = 0 . 25 . We compare the performance of three different limiting kernels: white (blue), a SE with a longer lengthscale ℓ = 1 (orange), and a SE with a shorter lengthscale ℓ = 0 . 1 (green). As the dataset is Gaussian, we have access to the true score. We observe that, across the different parameterisations, the white limiting kernel performance best. However, note that for the White kernel K = I and thus the different parameterisations become identical. For non-white limiting kernels we see a reduction in performance for both the approximate and exact score. We attribute this to the additional complexity of learning a non-diagonal covariance.

In the bottom row of App. F.1.3 we repeat the experiment for a dataset consisting of samples from the Periodic GP with lengthscale 0 . 5 . We draw similar conclusions: the best performing limiting kernel, across the different parametrisations, is the White noise kernel.

F.1.3 Ablation Conditional Sampling

Next, we focus on the empirical performance of the different noising schemes in the conditional sampling, as discussed in 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 1D squared exponential dataset (described above) as this gives us access to the true posterior. We use 2 12 samples to estimate the empirical mean and covariance, and fix the number of context points to 3.

In Fig. 11 we keep the total number of score evaluations fixed to 5000 and vary the number of steps in the inner ( L ) loop such that the number of outer steps is given by the ratio 5000 /L . From the figure, we observe that the particular choice of noising scheme is of less importance as long at least a couple ( ± 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.

-

2

.

5

0

.

0

2

.

5

SE

Data Prior

Model Posterior

True Posterior

-

2

.

5

0

.

0

2

.

5

Mat´ ern-

5

2

-

2

.

5

0

.

0

2

.

5

Weakly Periodic

-

0

.

5

0

.

0

0

.

5

Sawtooth

2

1

0

1

2

-

2

.

5

0

.

0

2

.

5

Mixture

2

1

0

1

2

2

1

0

1

2

-

-

-

White

SE

(

/lscript

= 1

.

0)

SE

(

/lscript

= 0

.

1)

exact score

approximate score

0

1

2

3

4

5

Precondition by

K

0

.

50

0

.

75

TLL

0

1

2

3

4

5

Precondition by

S

/latticetop

0

1

2

3

4

5

No preconditioning

0

1

2

3

4

5

Predict

y

0

(a) Squared Exponential dataset with lengthscale

ℓ

= 0

.

25

White

Periodic

(

/lscript

= 1

.

0)

Periodic

(

/lscript

= 0

.

1)

exact score

approximate score

0

1

2

3

4

5

Precondition by

K

0.00

1.00

TLL

0

1

2

3

4

5

Precondition by

S

/latticetop

0

1

2

3

4

5

No preconditioning

0

1

2

3

4

5

Predict

y

0

(b) Periodic dataset with lengthscale

ℓ

= 0

.

25

-

-

-

0

1

5

10

20

100

num_inner_steps

10 1

10 2

10 3

kl

method

independent_context_noise_without_resampling

sample_path_context_noise

independent_context_noise_with_resampling

raw_context

F.2 Gaussian process vector fields

Data We create synthetic datasets using samples from two-dimensional zero-mean GPs with the following E(2) -equivariant kernels: a diagonal Squared-Exponential (SE) kernel, a zero curl (CURLFREE) kernel and a zero divergence (DIV-FREE) kernel, as described in App. D.1. We set the variance to σ 2 = 1 and the lengthscale to ℓ = √ 5 . We evaluate these GPs on a disk grid, created via a 2D grid with 30 × 30 points regularly space on [ -10 , 10] 2 and keeping only the points inside the disk of radius 10 . We create a training dataset of size 80 × 10 3 . and a test dataset of size 10 × 10 3 .

Models We compare two flavours of our model GeomNDP. One with a non-equivariant attentionbased score network (Figure C.1, Dutordoir et al., 2022), referred as NDP * . Another one with a E(2) -equivariant score architecture, based on steerable CNNs (Thomas et al., 2018; Weiler et al., 2018). We rely on the e3nn library (Geiger and Smidt, 2022) for implementation. A knn graph E is built with k = 20 . The pairwise distances are first embed into µ ( r ab ) with a 'smooth_finite' basis of 50 elements via e3nn.soft_one_hot_linspace , and with a maximum radius of 2 . The time is mapped via a sinusoidal embedding ϕ ( t ) (Vaswani et al., 2017). Then edge features are obtained as e ab = Ψ ( e ) ( µ ( r ab ) || ϕ ( t )) ∀ ( a, b ) ∈ E k with Ψ ( e ) an MLP with 2 hidden layers of width 64 . We use 5 e3nn.FullyConnectedTensorProduct layers with update given by V k +1 a = ∑ b ∈N ( a, E k ) V k a ⊗ ( Ψ v ( e ab || V k a || V k b ) ) Y (ˆ r ab ) with Y spherical harmonics up to order 2 m Ψ v an

MLP with 2 hidden layers of width 64 acting on invariant features, and node features V k having irreps 12x0e + 12x0o + 4x1e + 4x1o . Each layer has a gate non-linearity (Weiler et al., 2018).

Wealso evaluate two neural processes, a translation-equivariant CONVCNP (Gordon et al., 2020) with decoder architecture based on 2D convolutional layers (LeCun et al., 1998) and a C4 ⋉R 2 ⊂ E(2) -equivariant STEERCNP (Holderrieth et al., 2021) with decoder architecture based on 2D steerable convolutions (Weiler and Cesa, 2021). Specific details can be found in the accompanying codebase https://github . com/PeterHolderrieth/Steerable_CNPs of Holderrieth et al. (2021).

Optimisation. Models are trained for 80 k iterations, via (Kingma and Ba, 2015) with a learning rate of 5 e -4 and a batch size of 32 . The neural diffusion processes are trained unconditionally, that is we feed GP samples evaluated on the full disk grid. Their weights are updated via with exponential moving average, with coefficient 0 . 99 . The diffusion coefficient is weighted by β : t ↦→ β min +( β max -β min ) · t , and β min = 1 e -4 , β max = 15 .

As standard, the neural processes are trained by splitting the training batches into a context and evaluation set, similar to when evaluating the models. Models have been trained on A100-SXM-80GB GPUs.

Evaluation. We measure the predictive log-likelihood of the data process samples under the model on a held-out test dataset. The context sets are of size 25 and uniformly sampled from a disk grid of size 648 , and the models are evaluated on the complementary of the grid. For neural diffusion processes, we estimate the likelihood by solving the associated probability flow ODE (53). The divergence is estimated with the Hutchinson estimator, with Rademacher noise, and 8 samples, whilst the ODE is solved with the 2nd order Heun solver, with 100 discretisation steps.

We also report the performance of the data-generating GP, and the same GP but with diagonal posterior covariance GP (DIAG.).

F.3 Tropical cyclone trajectory prediction

Data. The data is drawn from 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 at intervals of every 3 hours. From the dataset, we selected tracks of at least 50 time points long and clipped any longer to this length, resulting in 5224 cyclones. 90% was used fro training and 10% held out for evaluation. This split was changed across seeds. More interesting schemes of variable-length tracks or of interest, but not pursued here in this demonstrative experiment. Natively the track locations live in latitude-longitude coordinates, although it is processed into different forms for different models. The time stamps are processed into the number of days into the cyclone forming and this format is used commonly between all models.

Models.

Four models were evaluated.

The GP ( R → R 2 ) took the raw latitude-longitude data and normalised it. Using a 2-output RBF kernel with no covariance between the latitude and longitude and taking the cyclone time as input, placed a GP over the data. The hyperparameters of this kernel were optimised using a maximum likelihood grid search over the data. Note that this model places density outside the bounding box of [ -90 , 90] × [ -180 , 180] that defines the range of latitude and longitude, and so does not place a proper distribution on the space of paths on the sphere.

The STEREOGRAPHIC GP ( R → R 2 / { 0 } ) 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 a set of measure zero that does not correspond to the sphere, it does induce a proper distribution on the space of paths on the sphere.

The NDP ( R → R 2 ) uses the same preprocessing as GP ( R → R 2 ) but uses a Neural Diffusion Process from (Dutordoir et al., 2022) to model the data. This has the same shortcomings as the GP ( R → R 2 ) in not placing a proper density on the space of paths on the sphere. The network used for the score function and the optimisation procedure is detailed below. A linear beta schedule was used with β 0 = 1 e -4 and β 1 = 10 . The reverse model was integrated back to ϵ = 5 e -4 for numerical

stability. The reference measure was a white noise kernel with a variance 0 . 05 . ODEs and SDEs were discretised with 1000 steps.

The GEOMNDP ( R →S 2 ) works with the data projected into 3d space on the surface of the sphere. This projection makes no difference to the results of the model, but makes the computation of the manifold functions such as the exp map easier, and makes it easier to define a smooth score function on the sphere. This is done by outputting a vector for the score from the neural network in 3d space, and projecting it onto the tangent space of the sphere at the given point. For the necessity of this, see (De Bortoli et al., 2021). The network used for the score function and the optimisation procedure is detailed below. A linear beta schedule was used with β 0 = 1 e -4 and β 1 = 15 . The reverse model was integrated back to ϵ = 5 e -4 for numerical stability. The reference measure was a white noise kernel with a variance 0 . 05 . ODEs and SDEs were discretised with 1000 steps.

Neural network. The network used to learn the score function for both NDP ( R → R 2 ) and GEOMNDP ( R →S 2 ) is a bi-attention network from Dutordoir et al. (2022) with 5 layers, hidden size of 128 and 4 heads per layer. This results in 924k parameters.

Optimisation. NDP ( R → R 2 ) and GEOMNDP ( R →S 2 ) were both optimised using (correctly implemented) Adam for 250k steps using a batch size of 1024 and global norm clipping of 1. Batches were drawn from the shuffled data and refreshed each time the dataset was exhausted. A learning rate schedule was used with 1000 warmup steps linearly from 1e-5 to 1e-3, and from there a cosine schedule decaying from 1e-3 to 1e-5. With even probability either the whole cyclone track was used in the batch, or 20 random points were sub-sampled to train the model better for the conditional sampling task.

Conditional sampling. The GP models used closed-form conditional sampling as described. Both diffusion-based models used the Langevin sampling scheme described in this work. 1000 outer steps were used with 25 inner steps. We use a ψ = 1 . 0 and λ 0 = 2 . 5 . In addition at the end of the Langevin sampling, we run an additional 150 Langevin steps with t = ϵ as this visually improved performance.

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-likelihood of the whole trajectory and the log-likelihood of the context set only. The mean squared errors were computed using the geodesic distance between 10 conditionally sampled trajectories, described above.

Peer reviews:
comments: 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).

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).

QUESTIONS:
N/A

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nRegarding 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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$\n\nQUESTIONS:\n- 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\nLIMITATIONS:\nThe limitations are adequately addressed within the submission.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nI 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.\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nLimitations 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-19"}