SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Emile Mathieu*
University Of Cambridge
Vincent Dutordoir*
Secondmind Labs
Michael J. Hutchinson*
University Of Oxford
Valentin De Bortoli
Center for Science of Data, ENS Ulm
Yee Whye Teh
Richard E. Turner
University Of Cambridge,
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.
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: \mathbb{R}^n \to \mathbb{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: \mathcal{M} \to 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: \mathbb{R}^3 \to \mathbb{R}$ , representing wind over the Earth's surface as  $f: S^2 \to T S^2$ , where  $T S^2$  is the tangent-space of the sphere, or modelling the stress in a deformed object as  $f: \mathrm{Object} \to \mathrm{TR}^3 \otimes \mathrm{TR}^3$ , where  $\otimes$  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.
Secondly, we look at fields with manifold codomain, and in particular, at functions of the form  $f: \mathbb{R} \to \mathcal{M}$ . The challenge in dealing with manifold-valued output, arises from the lack of vector-space structure. In applications, these functions typically appear when modelling processes indexed by time that take values on a manifold. Examples include tracking the eye of 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  $\mathrm{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:\mathcal{X}\to \mathcal{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.
Denoising diffusion models. We briefly recall here the key concepts behind diffusion models on  $\mathbb{R}^d$  and refer the readers to (Song et al., 2021) for a more detailed introduction. We consider a forward noising process  $(\mathbf{Y}_t)_{t\geq 0}$  defined by the following Stochastic Differential Equation (SDE)
$$
\mathrm {d} \mathbf {Y} _ {t} = - \frac {1}{2} \mathbf {Y} _ {t} \mathrm {d} t + \mathrm {d} \mathbf {B} _ {t}, \quad \mathbf {Y} _ {0} \sim p _ {0}, \tag {1}
$$
where  $(\mathbf{B}_t)_{t\geq 0}$  is a  $d$  -dimensional Brownian motion and  $p_0$  is the data distribution. The process  $(\mathbf{Y}_t)_{t\geq 0}$  is simply an Ornstein-Ulthenbeck (OU) process which converges with geometric rate to  $\mathrm{N}(0,\mathrm{Id})$  (Durmus and Moulines, 2017). Under mild conditions on  $p_0$ , the time-reversed process  $(\bar{\mathbf{Y}}_t)_{t\geq 0} = (\mathbf{Y}_{T - t})_{t\in [0,T]}$  also satisfies an SDE (Cattiaux et al., 2021) given by
$$
\mathrm {d} \bar {\mathbf {Y}} _ {t} = \left\{\frac {1}{2} \bar {\mathbf {Y}} _ {t} + \nabla \log p _ {T - t} (\bar {\mathbf {Y}} _ {t}) \right\} \mathrm {d} t + \mathrm {d} \mathbf {B} _ {t}, \quad \bar {\mathbf {Y}} _ {0} \sim p _ {T}, \tag {2}
$$
where  $p_t$  denotes the density of  $\mathbf{Y}_t$ . Unfortunately we cannot sample exactly from (2) as  $p_T$  and the scores  $(\nabla \log p_t)_{t \in [0, T]}$  are unavailable. First,  $p_T$  is substituted with the limiting distribution  $\mathrm{N}(0, \mathrm{Id})$  as it converges towards it. Second, one can easily show (Kallenberg, 2021, Thm 5.1) that the score  $\nabla \log p_t$  is the minimiser of  $\ell_t(\mathbf{s}) = \mathbb{E}[\| \mathbf{s}(\mathbf{Y}_t) - \nabla_{y_t} \log p_{t|0}(\mathbf{Y}_t | \mathbf{Y}_0) \|^2]$  over functions  $\mathbf{s} : [0, T] \times \mathbb{R}^d \to \mathbb{R}^d$  where the expectation is over the joint distribution of  $\mathbf{Y}_0, \mathbf{Y}_t$ , and as such can readily be approximated by a neural network  $\mathbf{s}_{\theta}(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 = \operatorname{E}(d)$ , which elements  $g$  admits a unique decomposition  $g = uh$  where  $h \in \mathrm{O}(d)$  is a  $d \times d$  orthogonal matrix and  $u \in \mathrm{T}(d)$  is a translation which can be identified as an element of  $\mathbb{R}^d$ ; for a vector  $x \in \mathbb{R}^d$ ,  $g \cdot 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)$
$\rho (h)f(g^{-1}x)$
Figure 1: Illustration of a vector field  $f:\mathbb{R}^2\to \mathbb{R}^2$  with representation  $\rho (h) = h$  being steered by a group element  $h = 90^{\circ}\in \mathrm{O}(2)\subset \mathrm{E}(2)$ .
We are interested in learning a probabilistic model over functions of the form  $f: \mathcal{X} \to \mathbb{R}^d$  such that a group  $G$  acts on  $\mathcal{X}$  and  $\mathbb{R}^d$ . We call a feature field a tuple  $(f, \rho)$  with  $f: \mathcal{X} \to \mathbb{R}^d$  a mapping between input  $x \in \mathcal{X}$  to some feature  $f(x)$  with associated representation  $\rho: G \to \mathrm{GL}(\mathbb{R}^d)$  (Scott and Serre, 1996). This feature field is said to be  $G$ -steerable if it is transformed for all  $x \in \mathcal{X}, g \in G$  as  $g \cdot f(x) = \rho(g)f(g^{-1} \cdot x)$ . In this setting, the action of  $\operatorname{E}(d) = \operatorname{T}(d) \rtimes \operatorname{O}(d)$  on the feature field  $f$  yields  $g \cdot f(x) = (uh) \cdot f(x) \triangleq \rho(h)f\left(h^{-1}(x - u)\right)$ . Typical examples of feature fields include scalar fields with  $\rho_{\mathrm{triv}}(g) \triangleq 1$  transforming as  $g \cdot f(x) = f\left(g^{-1}x\right)$  such as temperature fields, and vectors or potential fields with  $\rho_{\mathrm{Id}}(g) \triangleq h$  transforming as  $g \cdot f(x) = hf\left(g^{-1}x\right)$  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  $\mu(g \cdot f) = \mu(f), \forall g \in 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: \mathcal{X} \to \mathcal{Y}$ , with  $\mathcal{Y} = \mathbb{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  $(\mathbf{Y}_0(x))_{x\in \mathcal{X}}$ . Given any  $x = (x^{1},\ldots ,x^{n})\in \mathcal{X}^{n}$ , we consider the following forward noisiging process  $(\mathbf{Y}_t(x))_{t\geq 0}\triangleq (\mathbf{Y}_t(x^1),\dots ,\mathbf{Y}_t(x^n))_{t\geq 0} = (\mathbf{Y}_t^1,\dots ,\mathbf{Y}_t^n)_{t\geq 0}\in \mathcal{Y}^n$  defined by the following SDE
$$
\mathrm {d} \mathbf {Y} _ {t} (x) = \frac {1}{2} \left\{m (x) - \mathbf {Y} _ {t} (x) \right\} \beta_ {t} \mathrm {d} t + \beta_ {t} ^ {1 / 2} \mathrm {K} (x, x) ^ {1 / 2} \mathrm {d} \mathbf {B} _ {t}, \tag {3}
$$
where  $\mathrm{K}(x,x)_{i,j} = k(x^i,x^j)$  with  $k:\mathcal{X}\times \mathcal{X}\to \mathbb{R}$  a kernel and  $m:\mathcal{X}\rightarrow \mathcal{Y}$ . The process  $(\mathbf{Y}_t(x))_{t\geq 0}$  is a multivariate Ornstein-Uhlenbeck process—with drift  $b(t,x,\mathbf{Y}_t(x)) = m(x) - \mathbf{Y}_t(x)$  and diffusion coefficient  $\sigma (t,x,\mathbf{Y}_t(x)) = \mathrm{K}(x,x)$ —which converges with geometric rate to  $\mathrm{N}(m(x),\mathrm{K}(x,x))$ . Using Phillips et al. (2022), it can be shown that this convergence extends to the process  $(\mathbf{Y}_t)_{t\geq 0}$  which converges to the Gaussian Process with mean  $m$  and kernel  $k$ , denoted  $\mathbf{Y}_{\infty}$ .
In the specific instance where  $k(x,x^{\prime}) = \delta_{x}(x^{\prime})$ , then the limiting process  $\mathbf{Y}_{\infty}$  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  $\mathbf{Y}_{\infty}$ . 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  $\mathbf{Y}_0^2$ , the time-reversal process  $(\bar{\mathbf{Y}}_t(x))_{t\geq 0}$  satisfies
$$
\mathrm {d} \bar {\mathbf {Y}} _ {t} (x) = \left\{- \frac {1}{2} \left(m (x) - \bar {\mathbf {Y}} _ {t} (x)\right) + \mathrm {K} (x, x) \nabla \log p _ {T - t} (\bar {\mathbf {Y}} _ {t} (x)) \right\} \beta_ {T - t} \mathrm {d} t + \beta_ {T - t} ^ {1 / 2} \mathrm {K} (x, x) ^ {1 / 2} \mathrm {d} \mathbf {B} _ {t}, \tag {4}
$$
with  $\bar{\mathbf{Y}}_0\sim \mathrm{GP}(m,k)$  and  $p_t$  the density of  $\mathbf{Y}_t(x)$  w.r.t. Lebesgue. In practice, the  $\log p_{T - t}$  term-known 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  $\bar{\mathbf{Y}}_0\sim \mathrm{GP}(m,k)$  and then simulating the reverse diffusion (4) (e.g. via Euler-Maruyama discretisation).
$^{2}$ Haussmann and Pardoux (1986) assumes Lipschitz continuity of the drift and volatility matrix and Cattiaux et al. (2021, Thm 4.9) which only assume a finite entropy condition on the space of processes.
Manifold valued outputs. So far we have defined our generative model with  $\mathcal{V} = \mathbb{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  $\mathcal{V}$  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  $\mathrm{K}(x,x)\nabla \log p_t$ , we directly approximate it with a neural network  $(\mathrm{Ks})_{\theta}:[0,T]\times \mathcal{X}^n\times \mathcal{Y}^n\to \mathrm{T}\mathcal{Y}^n$ , where  $\mathrm{T}\mathcal{V}$  is the tangent bundle of  $\mathcal{V}$ , see App. C. The conditional score of the noising process (3) is given by
$$
\nabla_ {\mathbf {Y} _ {t}} \log p _ {t} (\mathbf {Y} _ {t} (x) | \mathbf {Y} _ {0} (x)) = - \Sigma_ {t | 0} ^ {- 1} (\mathbf {Y} _ {t} (x) - m _ {t | 0}) = - \sigma_ {t | 0} ^ {- 1} \mathrm {K} (x, x) ^ {- 1 / 2} \varepsilon , \tag {5}
$$
since  $\mathbf{Y}_t = m_{t|0} + \Sigma_{t|0}^{1/2}\varepsilon$  with  $\varepsilon \sim \mathrm{N}(0, \mathrm{Id})$ , and  $\Sigma_{t|0} = \sigma_{t|0}^2\mathrm{K}$  with  $\sigma_{t|0} = (1 - \mathrm{e}^{-\int_0^t\beta(s)\mathrm{d}s})^{1/2}$ , see App. B.1. We learn the preconditioned score  $(\mathrm{Ks})_\theta$  by minimising the following denoising score matching (DSM) loss (Vincent et al., 2010) weighted by  $\Lambda(t) = \sigma_{t|0}^2\mathrm{K}^\top\mathrm{K}$
$$
\mathcal {L} (\theta ; \Lambda (t)) = \mathbb {E} [ \| \mathrm {s} _ {\theta} (t, \mathbf {Y} _ {t}) - \nabla \log p _ {t} (\mathbf {Y} _ {t} | \mathbf {Y} _ {0}) \| _ {\Lambda (t)} ^ {2} ] = \mathbb {E} [ \| \sigma_ {t | 0} \cdot (\operatorname {K s}) _ {\theta} (t, \mathbf {Y} _ {t}) + \mathrm {K} ^ {1 / 2} \varepsilon \| _ {2} ^ {2} ], \tag {6}
$$
where  $\| x\|_{\Lambda}^{2} = x^{\top}\Lambda x$ . Note that when targeting a unit-variance white noise, then  $\mathrm{K} = \mathrm{Id}$  and the loss (6) reverts to the DSM loss with weighting  $\lambda (t) = 1 / \sigma_t^2 |_0$  (Song et al., 2021). In App. B.2, we explore several preconditioning terms and associated weighting  $\Lambda (t)$ . Overall, we found the preconditioned score  $\mathrm{K}\nabla \log p_t$  parameterisation, in combination with the  $\ell_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 \sim \mu$  is said to be  $G$ -invariant if  $\mu(g \cdot \mathsf{A}) = \mu(\mathsf{A})$  for any  $g \in G$ , with  $\mu \in \mathcal{P}(\mathrm{C}(\mathcal{X},\mathcal{Y}))$ , where  $\mathcal{P}$  is the space of probability measure on the space of continuous functions and  $\mathsf{A} \subset \mathrm{C}(\mathcal{X},\mathcal{Y})$  measurable. From a sample perspective, this means that with input-output pairs  $\mathcal{C} = \{(x^i,y^i)\}_{i=1}^n$ , and denoting the action of  $G$  on this set as  $g \cdot \mathcal{C} \triangleq \{(g \cdot x^i,\rho(g)y^i)\}_{i=1}^n$ ,  $f \sim \mu$  is  $G$ -invariant if and only if  $g \cdot \mathcal{C}$  has the same distribution as  $\mathcal{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  $\mathrm{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^{\prime}\in \mathcal{X},g\in G$
$$
m (g \cdot x) = \rho (g) m (x) a n d k (g \cdot x, g \cdot x ^ {\prime}) = \rho (g) k (x, x ^ {\prime}) \rho (g) ^ {\top}. \tag {7}
$$
Trivial examples of  $\mathrm{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  $(\bar{\mathbf{Y}}_t(x))_{x\in \mathcal{X},t\in [0,T]}$  the process induced by the time-reversal SDE (4) where the score is approximated by a score network  $\mathbf{s}_{\theta}:[0,T]\times \mathcal{X}^n\times \mathcal{Y}^n\to \mathrm{T}\mathcal{Y}^n$  , and the limiting process is given by  $\mathcal{L}(\bar{\mathbf{Y}}_0) =$ $\mathrm{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.  $\mathbf{s}_{\theta}(t,g\cdot x,\rho (g)y) = \rho (g)\mathbf{s}_{\theta}(t,x,y)$  for all  $x\in \mathcal{X},g\in G$  , then for any  $t\in [0,T]$  the process  $(\bar{\mathbf{Y}}_t(x))_{x\in \mathcal{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  $(\mathbf{Y}_0(x))_{x\in \mathcal{X}}$  such as a temperature field, we need the score network to admit the invariance constraint  $\mathbf{s}_{\theta}(t,g\cdot x,y) = \mathbf{s}_{\theta}(t,x,y)$ .
Equivariant conditional process. Often precedence is given to modelling the predictive process
Figure 2: Samples from equivariant neural diffusion processes conditioned on context set  $\mathcal{C}$  (in red) and evaluated on a regular grid  $x^{*}$  for scalar (Left) and 2D vector (Right) fields. Same model is then conditioned on transformed context  $g\cdot \mathcal{C}$ , with group element  $g$  being a translation of length 2 (Left) or a  $90^{\circ}$  rotation (Right).
given a set of observations  $\mathcal{C} = \{(x^c, y^c)\}_{c \in \mathcal{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  $\mu$  given a context  $\mathcal{C}$  is said to be conditionally  $G$ -equivariant if the conditional satisfies  $\mu(\mathsf{A}|g \cdot \mathcal{C}) = \mu(g \cdot \mathsf{A}|\mathcal{C})$  for any  $g \in G$  and  $\mathsf{A} \in \mathrm{C}(\mathcal{X}, \mathcal{Y})$  measurable, as illustrated in Fig. 2.
Proposition 3.3. Equivariant conditional process. Assume a stochastic process  $f \sim \mu$  is  $G$ -invariant. Then the conditional process  $f|_{\mathcal{C}}$  given a set of observations  $\mathcal{C}$  is  $G$ -equivariant.
Originally stated in Holderrieth et al. (2021) in the case where the process is over functions of the form  $f: \mathbb{R}^n \to \mathbb{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  $\nabla_{x}\log p(x|y) = \nabla_{x}\log p(x,y) - \nabla_{x}\log p(y) = \nabla_{x}\log p(x,y)$  we can therefore, for any point in the diffusion time, conditionally sample using Langevin dynamics, following the SDE  $\mathrm{dY}_s = \frac{1}{2}\mathrm{K}\nabla \log p_{T - t}(\mathbf{Y}_s)\mathrm{d}s +$
Figure 3: Illustration of Langevin corrected conditional sampling. The black line represents the noising process dynamics  $(p_{t})_{t\in [0,T]}$ . The time reversal (i.e. predictor) step, is combined with a Langevin corrector step projecting back onto the dynamics.
$\sqrt{\mathrm{K}} \mathrm{d}\mathbf{B}_s$ , by only applying the diffusion to the variables of interest and holding the others fixed. While we could still be difficult in practice. Similar to the motivation of Song and R. diffusion, taking a number of conditional Langevin steps forward noisig SDE to the conditioning points at each step sampling set in a region that the score function will be well illustrated in Alg. 1. In App. E.1 we draw links with REP
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\in C})$  given a context set  $\{x^i,y^i\}_{i\in C}$ . A simple approach is to simply rely on the conditional probability rule evaluate  $p(y^{*}|x^{*},\{x^{i},y^{i}\}_{i\in C}) = p(y^{*},\{y^{i}\}_{i\in C}|x^{*},\{x^{i}\}_{i\in C}) / p(\{y^{i}\}_{i\in C}|\{x^{i}\}_{i\in 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.