Geometric Neural Diffusion Processes

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

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes

Summary

In this paper, the authors propose a generative model over vector fields. The proposed approach is a combination of ideas from neural diffusion processes [1], equivariant learning of stochastic fields [2] and equivariant scores [3]. The authors represent vector fields as vectors in finite sets of points and use Euclidean groups to incorporate desired symmetries on vector fields. The proposed diffusion model defined via Ornstein–Uhlenbeck process converges to Gaussian Process with equivariant mean and kernel. The model is trained on the denoising score matching loss with score parametrized by an equivariant neural network. This methodology can be extended to manifold-valued functions using Riemannian diffusion models. Further, the authors explore conditional sampling from the proposed model. [1] Dutordoir, V., Saul, A., Ghahramani, Z., and Simpson, F. (2022). Neural Diffusion Processes. [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. [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 arXiv:2302.02277.

Strengths

The paper is well written and provides all necessary information for understanding. Main distinguishing features of proposed model are well demonstrated in experiments. The proposed model needs much less training examples to reach high log-likelihood. In several small-scale experiments the proposed model outperforms existing approaches.

Weaknesses

The experiments are small-dimensional. It would be interesting to see how the proposed approach scales to higher dimensionality, where equivariant models could become even more beneficial.

Questions

Regarding Figure 5: The default NDP fails to achieve the same level of log-likelihood even with large amount of training data. Would it reach the same log-likelihood as E(2)-GeomNDP when the number of training examples tends to “infinity” (> 10^5)? Does it need more parameters? Or is there something else that prevents it from adequately fitting the data?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes

Summary

This paper proposes GeomNDP, a continuous-time functional diffusion model which can incorporate symmetries of the problem of interest. In particular, this model is constructed via a collection of finite-dimensional Ornstein-Uhlenbeck processes (one for each finite-dimensional marginal), which target a given prior Gaussian process. In addition, the authors study (Prop. 3.2, 3.3) conditions under which the resulting process is invariant to a given group. A conditional sampling method is also proposed (Section 3.3) which generalizes the RePAINT method (Lugmayr et al.) to the functional setting.

Strengths

- The paper is well-written, clear, and generally has a high-quality presentation. - 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. - The proposed model shows competitive performance against several baselines across multiple tasks. In general, the experimental evaluation is clear and convincing. - 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. - The work is well-contextualized within the relevant literature.

Weaknesses

- The exposition of the conditional sampler (Section 3.3) was somewhat unclear (see questions), and required a careful read of the pseudocode in order to understand what exactly was being proposed. - 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. ### Minor comments - The submission was not properly anonymized (see the Appendix in the main submission) - Line 85: it was somewhat unclear what exact assumptions are made on the domain $X$

Questions

- Could you elaborate on the differences between the proposed conditional sampler (Sect. 3.3) and the "replacement" method? E.g. previously proposed conditional functional diffusion models [1-2] work in a manner that, from the description starting on Line 199, in a very similar manner. - 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? [1] Neural Diffusion Processes. Dutordoir, Saul, Ghahramani, Simpson, 2023. [2] Diffusion Generative Models in Infinite Dimensions. Kerrigan, Ley, Smyth, 2023.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

The limitations are adequately addressed within the submission.

Summary

The present thesis deals with the problem of learning geometric-sensitive Denoising Diffusion Models (DDMs) for modeling equivariant conditional processes in (non)-Euclidean observation spaces. In doing so, the authors extend existing work that defines diffusions over function spaces with infinite dimensions using a (non-)Euclidean score-based modeling approach. They incorporate a series of geometric priors by constructing a noise process that allows a geometric Gaussian process as the limiting distribution transforming under the symmetry group of interest, and approximate the score as part of the denoising process by a neural network that is equivariant with respect to this group. They show that the generative function model admits the same symmetry under these conditions. The scalability and capacity of the proposed model is evaluated in the context of a novel Langevin-based conditional sampler based on synthetic and real data in the Euclidean and non-Euclidean scenario.

Strengths

I greatly appreciate any attempt to incorporate geometrically sensitive analysis into machine learning. From my perspective it is advantageous, to assume that data incorporates structural information and to give models the ability to address them. The proposed method builds on a novel combination of well-known techniques, and is original to the best of my knowledge. Further, the paper is quite well positioned to existing work, the content is well organized and, with a few exceptions, has no spelling or grammatical flaws. Empirical evidence is accessed through experiments on toy and real world data and the authors make effort to extend the small-scale experimental setting.

Weaknesses

I would rate the overall added value of the contribution as rather moderate. The biggest shortcomings relate to the clarity of the submission and the technical soundness which I will address with examples in the following. - Some general points: - 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. - 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. - 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. - 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. - 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. - Last but not least this setting even succeeds in case of non-Euclidean codomain such as compact manifolds. - 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. - Background - l. 63 Please provide a convergence result reference. - l. 69 Please provide a Can-easily-be-shown result reference. - 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$ - l. 83 Please check wording. - l. 73 Can you provide a motivation for considering steerable (feature) fields? - l. 89 What exactly is the feature field in (2)? - l. 90 What is $p$? - l. 93 Can you specify $\rho(h)$ in Fig. 1? - l. 94f. I don't understand this comment, can you rephrase it? - Section 3.1 - 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 - How is a data process defined; can you define the observation model in more detail? - 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)? - 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. - l. 120 Do you mean __Gaussian__ white noise? - l. 125 Please add a reference for what kind of mild conditions. - l. 128 Please add a reference for Stein score. - l. 136ff This part (i.e. Training) including the conditional score is unclear to me! Why does the noising process include a score? - What numerical solver is used in the manifold case for simulating the reverse diffusion? - Section 3.2 - l. 150 Sec. 2 include the definition of a feature field; how is this connected to a tensor field - l. 158 (Prop. 3.1) what about the stationary conditional of Theorem 1 (Holderrieth et al.)? - 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? - l. 181 Please define or give reference for a definition of a conditional process. - l. 184 Please define SP. - Section 3.3 - 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? - l. 218 Please check wording? - Section 4 - 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? - Section 5.1 - l. 271 Please define GEOMNDP and you are using GeomNDP a few times instead. - What about comparing models from Holderrieth et al. and Phillips et al. which are closer conceptional? - l. 292 Please define T(1)-GEOMNDP. - l. 295 Please define $\mathbb{I}(x=x’)$ or provide reference. - 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? Important side note: Including supplementary material with unfinished notes is not of advantage! However, 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.

Questions

- Regarding the usage of numerical solvers: From my point of view, numerical (SDE) solvers play a key role in the proposed approach, as it involves the simulation of (reverse) diffusions defined via SDEs. The authors use an Euler-Maruyama discretisation scheme, which is a very naive solver. Despite the fact that more advanced systems would drastically reduce computational efficiency, it would be important to address their benefits. Last but not least, the vanilla Euler-Maruyama scheme is based on a vector space structure. Therefore, in the case of manifold valued SDEs, there are more advanced options that take these intrinsic structures into account, e.g., a geometric Euler-Maruyama scheme. How could this improve the simulation process? - l. 66 & 72 (4) -> (2) ? - l. 173 admits -> admit - l. 179 $A \subset C$ and how is $x^\ast$ defined in Fig. 2?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

Limitations are addressed (e.g., cf. Chapter 6) and relate to a computational overhead and modeling weaknesses. No evidence was found that the authors addressed potential negative societal impact. Please make the effort to include this mandatory information in your work.