Scaling Riemannian Diffusion Models

Riemannian diffusion models draw inspiration from standard Euclidean space diffusion models to learn distributions on general manifolds. Unfortunately, the additional geometric complexity renders the diffusion transition term inexpressible in closed form, so prior methods resort to imprecise approximations of the score matching training objective that degrade performance and preclude applications in high dimensions. In this work, we reexamine these approximations and propose several practical improvements. Our key observation is that most relevant manifolds are symmetric spaces, which are much more amenable to computation. By leveraging and combining various ansätze, we can quickly compute relevant quantities to high precision. On low dimensional datasets, our correction produces a noticeable improvement, allowing diffusion to compete with other methods. Additionally, we show that our method enables us to scale to high dimensional tasks on nontrivial manifolds. In particular, we model QCD densities on $SU(n)$ lattices and contrastively learned embeddings on high dimensional hyperspheres.

Paper

References (52)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer VM6C5/10 · confidence 3/52023-06-24

Summary

This paper considers the diffusion model on Riemann space. They propose several strategies for numerically computing the heat kernel on Riemannian symmetric spaces. They also propose two techniques to sample from their proposed heat kernel approximations, including a rejection sampling method and a probability flow method.

Strengths

They compare several techniques to approximate the heat kernel in Riemannian symmetric space. They propose two techniques to obtain samples from the approximated heat kernel. They scale the models to the task of 128 dimension, which seems to be promising.

Weaknesses

The proposed several techniques in section 3 seem quite standard, and do not involve much novelty. From Figure 1, these several heat kernel estimators could be very sensitive, so the hyper-parameter $\tau$ in Algorithm 1 needs very careful tuning. The experimental parts consider only few baselines. For example, in figure 3 and table 1, can you compare with more previous Riemannian diffusion model works?

Questions

Rejection sampling is known to be non-scalable (low acceptance rate.) in high dimension. Do you have this problem when you apply Cheap Rejection Sampling Methods in section 4? Can you compare the two techniques in section 4? Which one is more convenient to use in practice?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

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

2 fair

Contribution

2 fair

Limitations

.

Reviewer 611J7/10 · confidence 3/52023-06-25

Summary

Developing diffusion model on Riemannian manifolds has received attention recently. However, the current proposals did not scale well to high dimension as they used sliced-score matching or denoising score matching with bad approximations of the heat kernel. This work focuses on using better approximations of the heat kernel on (compact) symmetric Riemannian manifolds. Then, they show the benefit of their proposals on several manifolds such as the torus, $S^2$, $S^{127}$ or $SU(n)$, and apply it to different problems such as density estimation of earth data or on the torus, or for contrastively learned hypersperical embeddings.

Strengths

Developing efficient generative models on Riemannian manifolds is of much interest, and this paper proposes to use efficient approximations of the transition kernel, which is pretty appealing. - Different kernel approximations with an analysis on the setting in which they work or not - Very competitive results compared to previous methods, especially on high-dimensional manifolds - Good Applications

Weaknesses

- I find the introduction of the transition kernel a bit quick and it is thus not very clear how to relate it to the equation (7,8) in the Euclidean diffusion. Maybe it would be more clear rewrite the equations with the kernel. - The experiment in Section 5.2 is not very clear.

Questions

The results of the experiment in Section 5.2 are not clear to me as it only provides an ESS of 0.62, but without any comparison or value of any other method. Thus, I do not understand what it brings. Not very related to the point of the paper, but it is stated that for non compact Riemannian manifolds, one can "map points to surfaces and learn the diffusion on $\mathbb{R}^d$. Is it not a bit naive? I would expect such method to be less performant than by directly applying the diffusion on the manifold. Typos: - Equation 10: lack closing parenthesis - Line 109: "prior to work to" - The notation "cot" in equation (15) is not defined. - Line 288: the abreviation ESS is not defined

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

4 excellent

Presentation

3 good

Contribution

4 excellent

Limitations

Yes

Reviewer 1Bu56/10 · confidence 4/52023-07-02

Summary

The theory of stochastic differential equations on vector spaces forms the basis for recent diffusion generative models, and recent efforts have attempted to extend this framework to stochastic differential equations on Riemannian manifolds (Riemannian diffusion models). From a practical computational perspective this extension is less than straightforward as pointed out in this paper, e.g., computing things like "straight lines" (minimal geodesics) become very difficult two-point boundary value problems; quantities that admitted closed-form analytic formulas in the vector space setting do not translate to closed-form formulas in the Riemannian setting. Past attempts at practical work-arounds and computational approximations are limited to low-dimensional cases; these methods do not scale well to higher-dimensional spaces. By restricting the focus to Riemannian manifolds that are also symmetric spaces, the authors present a collection of computation and approximation techniques that, taken together, demonstrate both computational tractability and scalability. Experiments validating these claims are presented.

Strengths

The paper is technically sound and well-written, and the contributions are laid out in a clear and concise manner without claiming too much. Rather than making claims about the significance of a particular result, I appreciate that the authors claim that it is the sum of a set of smaller contributions that, taken together, lead to a computationally practical, scalable set of techniques to make Riemannian diffusion models workable in practice. The focus on Riemannian symmetric spaces seems reasonable, and the detailed computation formulas and approximations are laid out in enough detail that they can be implemented by a reader familiar with stochastic differential equations and some Riemannian geometry. The examples and experiments appear to support the claims of the authors about computational tractability and scalability to higher-dimensional problems.

Weaknesses

-The following is a general comment applicable not only to this paper, but to all papers that present a “geometric” version of an existing algorithm or method in machine learning. In the past, developing a geometric, manifold version of a vector space algorithm could be regarded as a meaningful contribution in itself, and I would argue that the Riemannian diffusion model is somewhat in this spirit (equivariant models I would argue are not in this spirit however). Of course, it’s hard to argue against the claim that the Riemannian diffusion model is needed for problems in which the underlying data are manifold-valued; in that case I would be much more convinced by examples and case studies drawn from mainstream applications rather than narrow ones. As geometric methods have become more mainstream, the threshold for what constitutes a meaningful contribution is justifiably higher: given the much more difficult computations involved in computing, e.g., derivatives, gradients, Laplacians, minimal geodesics, kernels, etc., the considerable increase in computation needs to be justified by results. For this paper, the authors recognize and point out the difficulty of computing the geometric quantities and propose more efficient approximations, which is worthwhile (but whether they deserve to be published in NeurIPS is another matter). It would be helpful if the authors can address this question more directly – are the added computational difficulties justified by a commensurate improvement in the results? The experiments do not strike me as mainstream problems, for one thing. (Note: This comment could also have been placed in the "questions" section, but I place it here as it could be a potential weakness). -The definition of symmetric space could be sharpened, as this is an important underlying assumption in the paper, e.g., clarify the isometry requirement (isometry between what spaces?), make the distinction between local vs global symmetric space, provide brief intuition (reflexive symmetry about a point), and most importantly, provide examples of spaces arising in ML applications that are symmetric (the authors list some -- a few more would be helpful, e.g., the space of symmetric positive-definite matrices, which is nonimpact but a space that arises constantly in ML) and also examples of non-symmetric spaces (e.g., hyperbolic manifolds). -The assumption of the Riemannian manifold being embedded in a higher-dimensional ambient Euclidean space in itself is not restrictive thanks to Whitney and Nash, and most practical problems that I've encountered usually admit some natural embedding. The compactness assumption, however, could be somewhat restrictive for certain problems in which the underlying data manifold is potentially unbounded. The authors brush aside this case by asserting that noncompact Riemannian manifolds are diffeomorphic to Euclidean space, and therefore the standard way of using the embedding in R^n to model diffusion, then mapping back to the surface, is sufficient. One could then ask why this approach doesn't work for any manifold embedded in Euclidean space; it would be helpful if the authors could clarify this point.

Questions

-Explaining in more detail the Riemannian exponential map would be helpful, so that the reader has a better idea of what the computation entails. -The examples are for SU(3) and the hypersphere, which are rather specific manifolds. Particularly in the case of the hypersphere, I suspect that the heavy Riemannian diffusion machinery may not be needed to arrive at the result, since Brownian motion on the sphere is well-characterized and projections to the hypersphere are trivial.

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

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

No potential negative societal impact beyond that of other typical NeurIPS submissions that I can detect.

Reviewer NaGf6/10 · confidence 4/52023-07-03

Summary

In this work the authors leverage properties of symmetric spaces to get better approximation of the heat kernel and therefore of the denoising score matching loss for training diffusion models on compact manifolds. In particular, using the concept of maximal torus and roots of the symmetric space, the Laplacian can be written in terms of the radial Laplacian which itself depends on these roots and their multiplicity. Most interestingly, for (compact) Lie groups, the heat kernel can also be written as a summation over all possible paths between the origin and evaluation points, similarly to the heat kernel expression on the torus. Additionally, for small time scale, the authors propose to use the Schwinger-Dewitt approximation of the heat kernel which can be seen (I think) as a higher order Taylor approximation than the Varadhan one, and therefore is accurate for much higher values of $t$ while staying computationally fast. The authors then show that these approximations are crucial for training diffusion models on 'high' dimensional manifolds, but still improves performance on smaller ones such as $\mathcal{S}^2$ and $\mathrm{SU}(3)$.

Strengths

- I enjoyed reading this paper, as it is well organised, sounds and generally well written. - I was not aware of the result from Equation 21 'Sum Over Paths', and have been wondering in the past why does the approximation from a truncation of Eq 13 was getting worse as $t$ gets smaller, while it is the opposite for the heat kernel of the torus. This is really interesting, and would like to know more about the generalisation mentioned at line 219. Figure 1 is quite illustrative in what the Schwinger-Dewitt approximation brings compared the Varadhan! - The exact sampling from the (marginal) Brownian motion is also neat.

Weaknesses

- Section 3.1: Would be useful to give some intuition to the Definition 3.1 of symmetric manifolds. Generally, the concept of maximal torus and Equation 15 deserve a better presentation. The 'positive roots of our symmetric space' are not defined. Would be useful to say how the 'radial' Laplacian relates to the Laplacian of this symmetric space, and similarly for its heat kernel to the heat kernel of the symmetric space. - Table 1: Would be worth including [22] as they overperform [3] on several datasets. - Flow matching: Why not including the results from Flow Matching [9]? They seem to report better results on the Volcano and Fire datasets. Why does flow matching fails on the chequerboard on the torus but not on the sphere? Isn't the vector field also discontinuous there?

Questions

- Equation 11: As stated in the paragraph below this noising process is indeed for compact manifolds, yet [3] more generally proposed an 'Ornstein-Uhlenbeck' like process for non-compact manifolds, which simplifies to Eq 11 when compact. Would be worth having this more general setting in Section 2.2? - line 105: Perhaps slightly misleading to say that the heat kernel has 'has no closed form' since it is given by the Sturm–Liouville decomposition of Equation 13, although it is a converging series, and has such not 'computable' (since there is generally no finite time algorithm). - Line 181: It is indeed good to stress that one can bypass directly computing the eigenfunctions, although I don't think that in the past people have done this (naive approach)? E.g. equation 18 can be found in appendix F of [3]. - Section 3.2.2: Worth linking this with the $\exp(-\lambda_i t)$ term in Eq. 13, so for small values of $t$, terms in the series would tend to contribute more and more equally. - Eq 20: What is $v$? Where does this come from? - Figure 1: Would you know why the Schwinger-Dewitt approximation has this very fast change in accuracy? - 'RFM introduces several geodesic-based vector fields, but these break the smoothness assumption and the theoretical framework'. Could you please expand on that? This is using which (conditional) vector field? - Section 3.3: This is a really neat idea. Can the ratio be trivially bounded? Worth adding this in the paper! Does this help during training (for the Monte Carlo estimator of the denoising score matching loss)?

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

- This paper 'do not consider the case of noncompact Riemannian symmetric spaces, as these are diffeomorphic to Euclidean space'. I agree that since the topology is Euclidean it is perhaps less crucial to parameterise a probabilistic model directly on that space, yet the geometry is still different and only locally linear, so the geometry is more and more distorted further away, e.g. at boundary for the Poincare model of hyperbolic geometry. Is there some equivalent notion of the maximal Torus for non compact manifolds? - Would be worth expanding the evaluations in Section 5.2 which is a really interesting setting.

Reviewer VM6C2023-08-14

Reply

Thank you for your reply! I have a follow-up question: Have you considered comparing with non-diffusion based methods, such as the one blow? Cohen S, Amos B, Lipman Y. Riemannian convex potential maps. In International Conference on Machine Learning 2021 Jul 1 (pp. 2028-2038). PMLR.

Reviewer 1Bu52023-08-15

Just a short note to confirm that I've read your responses to my queries, thank you. Re your point about CLIP and the hypersphere, I suppose I was looking for other more generic examples other than the hypersphere (all of the Riemannian geometric machinery developed for diffusions on general manifolds typically reduces to very simple notions on the hypersphere, and often can be derived by appealing to more intuitive concepts).

Reviewer NaGf2023-08-18

response

Thanks for your detailed response! > we have found that RFM tends to quickly degrade when the task becomes harder (through dimension or data), which hampers applications in real world examples. Re the checkerboard from Figure 3: I am still struggling to understand how come the (optimal) flow suffers from this 'lack of smoothness' more than the Stein score. Would you please be able to expand on this? > Section 3.3: The ratio can be trivially bounded for most manifolds, and we found that the main benefit for our simulation strategies was a noticeable speed increase when compared with the many SDE Riemannian exponential map simulation steps that previous methods have to use to get good results [3, 22], which matters when working with many dimensions. That's cool! > On section 3.1: we agree with the reviewer that our presentation of these topics was a bit rushed. We will update our final paper to better present these concepts. That is really important indeed.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC