DIFFUSION GENERATIVE MODELS ON SO(3)
Authors: Yesukhei Jagvaral
Francois Lanusse
Rachel Mandelbaum
Venue: Submitted to ICLR 2023 (2023)
Keywords: Deep Generative Models
Manifold Learning
SO(3)
Denoising Diffusion
Score-based models

Abstract: Diffusion-based generative models represent the current state-of-the-art for image generation. However, standard diffusion models are based on Euclidean geometry and do not translate directly to manifold-valued data. In this work, we develop extensions of both score-based generative models (SGMs) and Denoising Diffusion Probabilistic Models (DDPMs) to the Lie group of 3D rotations, SO(3). SO(3) is of particular interest in many disciplines such as robotics, biochemistry and astronomy/planetary science. Contrary to more general Riemannian manifolds, SO(3) admits a tractable solution to heat diffusion, and allows us to implement efficient training of Diffusion models. We apply both SO(3) DDPMs and SGMs to synthetic densities on SO(3) and demonstrate state-of-the-art results.

Under review as a conference paper at ICLR 2023
DIFFUSION GENERATIVE MODELS ON S0(3)
Anonymous authors
Paper under double-blee blind review
ABSTRACT
Diffusion-based generative models represent the current state-of-the-art for image generation. However, standard diffusion models are based on Euclidean geometry and do not translate directly to manifold- valued data. In this work, we develop ex- tensions of both score-based generative models (SGMs) and Denoising Diffusion Probabilistic Models (DDPMs) to the Lie group of 3D rotations, SO(3). SO(3) is of particular interest in many disciplines such as robotics, biochemistry and astrononyyplanetary science. Contrary to more general Riemannian manifolds, SO(3) admits tractable solution to heat diffusion, and allows us implement efficient training of diffusion models. We apply both SO(3) DDPMs and SGMs to synthetic densities on SO(3) and demonstrate state-of-ther -fettherrt results.
INTRODUCTION
Deep generative models (DGM) are trained to learn the underlying data distribution and then gener- new samples that match the empirical data. There are several classes of deep generative models, including Generative Adversarial Networks (Goodfellow et al., 2014), Variational Auto Encoders (Kingma & Welling 2013) and Normalizing Flows (Rezende & Mohamed, 2015). Recently, new class of DGMs based on Diffusion, such as Denoising Diffusion Probabilistic Models (DDPM) (Ho) ct al. 2020) and Score Matching with Langevin Dynamics (SMLD), subset of general score-based generative models (SGMs), (Song & Ermon, 2019). have achieved state-of-the-art quality in gen- crating images, molecules, audio and graph (Song et al., 2021). Unlike GANs, training diffusion models is usually very stable and straightforward, they do not suffer as much from mode collapse issues, and they can generate images of similar quality.
In parallel with the success of these diffusion models, [ong et al. (2021) demonstrated that both SGMs and DDPMs can mathematically be understood as variants of the same process. In both cases, the data distribution is progressively perturbed by noise diffusion process defined by specific Stochastic Differential Equation (SDE), which can then be time-reversed to generate realistic data samples from initial noise samples.
While the success of diffusion models has mainly been driven by data with Euclidean geometry (e.g., images), there is great interest in extending these methods manifold-valued data, which are ubiquitous in many scientific disciplines. Examples include high-energy physics (Brehmer & Cranmer 2020 Craven al. 2022). astrophysics (Hemmati al. 2019), geoscience (Gaddes al. 2019), and biochemistry (Zelesko et al. (202)). Very recently, pioneering work has started to develop generic frameworks for defining SGMs on arbitrary compact Riemannian manifolds (De Bortoli
et al. 2022), and non-compact Riemannian manifolds (Huang et al. 2022).
In this work, instead of considering generic Riemannian manifolds, we are specifically concerned with the Special Orthogonal group in 3 dimensions, SO(3), which corresponds to the Lie group of 3D rotations. Modeling 3D orientations is particularly high interest in many fields including for instance in robotics (estimating the pose of an object, [oqque et al. 202T); and in biochemistry (finding the conformation angle of molecules that minimizes the binding energy, Mansimov et al. 2019). Contrary more generic Riemannian manifolds. SO(3) benefits from specific properties, including a tractable heat kernel and efficient geometric ODE/SDE solvers, that will allow us to define very efficient diffusion models specifically for this manifold.
For comprehensive list. of articles score-based generative modeling, see scoreebaeddeeeeaat www.eeeeeeee.iiieeeeeeeeme vemode thub. io/
https: ///

---

Under review as a conference paper at ICLR 2023
t=T t= 3T/4 t T/2
t=0
Figure 1: Illustration of reversible diffusion of mixture of two Igso3 blobs SO(3). Samples from given base distribution (right most, denoted by circles) can evolved under the probability flow ODE (Eq. towards a noisy distribution (Ieft most), or vice-versa from noisy distribution back to the target distribution. Each point represents rotation matrix in SO(3) projected on the sphere according to its canonical axis, the color indicates the tilt around that axis (visualisation adopted from Murphy et al. 2021). An animation of this figure is available this link
The contributions of paper summarized as follows:
We reformulate Euclidean diffusion models the SO(3) manifold, and demonstrate how the tractable heat kernel solution on SO(3) can be used to recover simple and efficient algorithms on this manifold.
We provide concrete implementations of both Score-Based Generative Model and Denois- ing Diffusion Probabilistic Models specialized for SO(3).
We reach state-of-the-art in sample quality synthetic SO(3) distributions with proposed SO(3) Score-Basee Generative Model.
2 PRELIMINARIES AND NOTATIONS
In this work, we are exclusively considering the SO(3) manifold, corresponding to the Lie group of 3D rotation matrices. We will denote by exp so(3 -SO(3) and log SO(3) -> so(3) the exponential and logarithmic maps that connect SO(3) to its tangent space and Lie algebra so(3). so(3) corresponds to all skew-symmetric 3x3 matrices in R3, which can be parameterised in terms of a vector in g3, which corresponds in turn to the axis-angle representation of rotation matrices. SO(3), its axis-angle representation
For any rotation w = arccos (2- (tr(x) - 1)) E ( π, v. and v 2 (x X23, X13 X31, X21 X12) a unit vector of 3
R3 can be computed as wv with
We direct the interested reader more information representations of SO(3) in Appendix
DIFFUSION PROCESS ON SO(3)
Similarly to Euclidean diffusion models (Song et al. 2021 we begin by defining Brownian noising process that will be used perturb the data distribution. Let us assume Stochastic Differential Equation of the following form:
dx f(xx t) dt + g(t) dw,
where is a Brownian process on SO(3), f(-,t): SO(3) + SO(3() drift term, and g(-) : IR + IR is a diffusion term. If we sample initial conditions for this SDE at t = 0 from given data distribution x(0) Pdata- we will denote by Pt the marginal distribution of x(t) time Thus
Po = Paaaa and at final time T at which we stop the diffusion process will typically tend to known target distribution that will easy to sample from.
Just like in the Euclidean case, demonstrated in De Bortoli et al. (2022) under mild regularity conditions Equuation Jadmits reverse diffusion process on compact Riemannian manifolds such as SO(3), defined by the following reverse-time SDE:
dx [f(x, 1) = log (x))dd +g()dw
2

---

Under review as a conference paper at ICLR 2023
where w is reversed-time Brownian motion and the score function V ll E Tx3)(3) is the derivative of the log marginal density of the forward process at time t. Corresponding to this reverse-time SDE, one can also define a probability flow ODE (Song et al. 2021):
dx = [[(x. t) - ( log (x
(3)
This deterministic process is entirely defined once the score is known and maps pT to any interme- diate marginal distributions {pt hoster of the forward process, including Pu- particular, it can be seen the equivalent of Neural ODE-based Continuous Normalizing Flows (CNF, [Chnn al. 2018) with an explicit parameterization in terms of the score function. We illustrate this process in Figure T with samples from two Gaussian-like blobs SO(3) being transported reversibly through this ODE between = 0 and / T.
While these equations direct analog of the Euclidean SDEs and ODE described in Song et al. 2021 defining diffusion generative models on SO(3) will mainly differ on the two following points:
Defining the equivalent of the Gaussian heat kernel on SOG): this needed to easily sample from any intermediate without having to simulate an SDE.
Solving SDEs and ODEs on the manifold: contrary to the Euclidean case, the diffusion process must remained confined the SO(3) manifold, which requires specific solvers.
We address these two points below before moving on to defining our generative models on SO(3).
3.1 THE ISOTROPIC GAUSSIAN DISTRIBUTION ON SO(3)
In general, the main disadvantage of working on Riemannian manifolds compared to Euclidean space is that they lack closed form expression for the heat kernel, i.c., the solution of the diffusion process (which a Gaussian Euclidean space). For compact manifolds, the heat kernel in general only available (Nikolayev & Savyolov 1970):
infinite series, which in the case of SO(3), takes the following form
sin (e (in(((1/2)w)) 1/2)w)
fe(w) = - + 120 1) 8=
(4)
sin(w//)
where w = |wس| E (-T,, is the rotation angle of the axis-angle representation w of a given rotation matrix and is concentration parameter.
While for > 1 this series converges quickly (fmax 5 is sufficient to achieve sub-percent accu- racy), the convergence gets slower gets smaller, which makes impractical to model concen- trated distributions. Thankfully, this series has been thoroughly studied in the literature and Matthies al. (1988) shows that an excellent approximation of EEquatin can be achieved for using the following closed-fon form expression:
fe(w) = VTE 3/2 157--3/2 ww/2)3 (w - ((w 2 mw/e + (w + 2Tr)e 27))e=w///
(5)
2 2w/
Therefore, in practical applications, one can switch between using truncation of Equation for 2 1 and the approximation Equation for <1.
Because of the property of being a solution of diffusion process on SO(3), fo can be used define the manifold equivalent of the Euclidean isotropic Gaussian distribution, which we will refer to as IGS00), the Isotropic Gaussian on SO(3) (Leach et al. 2022, Ryu et al. 2022), also known in the literature as the normal distribution on SO(3) (Nikolayev & Savyolov 1970 Matthies et al., 1988). For given mean rotation f E SO(3) and scale e, the probability density of rotation x E (SO()) under IG is given by:
IG pu,)) fo (arccos [2 (tr(p1 Tx) 1)])
Sampling from Ig (14, E) is achieved in practice by inverse transform sampling. The cumulative distribution function over angles needed to sample with respect to the uniform distribution on SO(3) can be evaluated numerically given integrating 1-cos(w) (w) over (- T). To form a rotation ma- trix ~ [G ,4, e), one therefore first samples rotation angle by inverse transform sampling
3

---

Under review as a conference paper at ICLR 2023
Algorithm Geometric ODE solver on SO(3) (Heun's method) for dx = f(x. t) dt
Require: Step size h, initial condition Xo, time steps {4n (4n)=o number of steps
1: for {0. ..... .V 1 do
2: Y1 = h fl f(,,, tn)
3: Y2 f(exp( ( 2211) in hh)
4: Xn+11 = exp(y2)Xn
5: end for
6:
return )
given this CDF, then samples uniformly on S rotation axis v, yielding an axis axis-angle representa- tion of rotation matrix w wv, which is then shifted by the mean of the distribution according to l exp(w).
An important property of [Gsoo))() (H, c). which sets it apart from other distributions on SO(3) (e.g. Bingham, Matrix Fisher, Wrapped Normal, more on this in Appendix [. is that remains closed under convolution, as direct consequence of being the solution of diffusion process. The convo- lution of two centered Igso() distributions of scale parameter E₁ and eg is an IGsog) distribution of scale + 62-
We will also note two interesting asymptotic behaviors. For large e, it tends to Usoo)) the uniform distribution SO(3), while for small the distribution IGsoo) e) can locally be approximated in the anise representation of the tangent space by a normal distribution N N(0 oT) in IR3, with
3.2 SOLVING ORDINARY DIFFERENTIAL EQUATIONS ON SO(3)
Thanks to the existence of a tractable heat kernel on SO(3), the generative models we will define in the next section will not actually require us to solve the SDEs introduced at the beginning of this section, and we will only need to solve the probability flow ODE defined in Equation 3
Solving differential equations manifolds broadly be achieved using two distinct strategies, ei- ther projection methods using a Fuclidean solver followed by projection step onto the manifold, or intrinsic methods that rely on additional structure of the manifold define an iteration that remains by construction on the manifold. In this work, we are concerned with SO(3), which is not only compact Riemannian manifold, but also possesses a Lie group structure, which makes it amenable to efficient solvers. In particular, we will make use of the Runge- Kutta--utttaeehhee Kaas (RK-MK) class of algorithms and direct interested reader to Iserles et al. (2000) for a review Lie group integrators. We adopt in practice the Lie group equivalent of Heun's 's method, which is one variant of RK-MK integrators, and we provide the details of this integrator in Algorithm
While will not require it in practice, it is also possible to build SDE solvers SO(3) with a similar strategy, and we point the interested reader for instance to the Geodesic Random Walk algorithm described in De Bortoli et al. (2022).
4 DIFFUSION GENERATIVE MODELS ON SO(3)
The core idea of diffusion models is to perturb given empirical data distribution Pdsta by a noise process defined in terms of a stochastic differential equation of the form Eqqation I While sev- eral expressions can be proposed for this SDE. for simplicity we will consider here the case of the Sariance-Explodiig SDE (for our experiments we use a Variance- Pariance-Preserving SDE the DDPM
case), with f(x, t) = and g()) = v de(t)) for given choice of noise schedule e(l), which cor- responds to the canonical choice of the Euclidean Score-Matching Langevin Dynamics (Song & Ermon 2019):
dt
dx de(()) dw
di
For our fiducial model, and unless stated otherwise, we will further assume for simplicity the fol- lowing noise schedule: <(t) f. The main drawback of this SDE in Fuclidean geometry is that it
4

---

Under review as a conference paper at ICLR 2023
t=T
t t= /
t= T/2
T
t= 0
Sampling from a Diffusion Generative Model trained on a synthetic density on SO(3). Starting from Usorr), the uniform distribution on SO(3) at t ' (left), the sampling procedure (either based on SGMs or DDPMs) denoises this distribution back to the target density at f 0 (right). For visualization this density plot shows the distribution of canonical of sampled rotations projected the sphere; the tilt around that axis discarded.
Figure 2:
will tend to a Gaussian with infinitely large variance. However, on SO(3) this SDE will tend to the uniform distribution US0(3) which is a natural choice for the prior distribution large T.
Following from this choice SDE, we can define noise kernel pe (a = IGsoa) x; for æ, x ESO(3), such that the data distribution convolved by this noise kernel becomes
(x
Pinta ),, (ax ( dæ
S0(3))
and corresponds to pt, the marginal distribution of the diffusion process at time t: Pc(t) = Pt-
Having introduced specific choice of kernel and SDE well suited to the SO(3) manifold, we now move on to describing the two different approaches to build generative models: Score-Based Mod- els and Denoising Diffusion Probabilistic Models. They both will lead to sampling procedures illustrated in Figure
4.1 SCORE-BASED GENERATIVE MODEL
The first strategy directly extends Euclidean SGMs (Song & Ermon 2019 Song et al. 2021 and relies on the time-reversed diffusion process described in Equation 2 Samples from the learned distribution Po be sampled by first sampling XT Uso(3) and evolving these samples either through the reverse SDE (Equation 2} or probability flow ODE (Equation 3) back to t = 0. This process is entirely defined as soon as the score function of the marginal distribution at any interme- diate time t, V log Pc(t), is known. Therefore the first step is to establish a score- woreematching strategy on SO(3).
Let us consider {X i= orthonormal basis of the tangent space T.SO(3). The directional deriva- tive of the log density the noise kernel p. (x æ+) can be computed as:
d ds
Vx, (xlx)
log P. exx(exxpxx))))
which can be computed in practice by automatic differentiation given the explicit approximation formulae for the I9s03) distribution introduced in (sbssection 3.1 To match this derivative, we introduce neural score estimator so(x,c) SO(3) H + R. which can be trained directly under a conventional denoising score matching loss:
LPSM =
[[e] s0 (x( e) Vx (ogP)c(((x)) 1
(10)
where we sample at training time random noise scales ~ N(O0 of) similarly to Song & Ermon 2020 The minimum of this loss will be achieved for so(x, V Vlogre
Once the score function is estimated from data using this score matching loss, sampling from the generative model can be achieved by using the reverse SDE formula, using the ODE flow formula. In this work, we use the latter for its simplicity and speed, so that our specific fiducial sampling strategy becomes:
dx, de(!) dt so sx(, (f)) dt
XT
Us033
(11)
which we solve down to f = with the geometric ODE solver described in Algorithm I stochastic sampling strategies based simulating the reverse SDE, this approach has several
Compared
5

---

Under review as a conference paper at ICLR 2023
Algorithm 2 Sampling from Denoising Diffusion Probabilistic Model on SO(3)
Require: Trained neural networks ne <x (o() 1), number of steps N, time steps {4)]
1: x MsSK3)
2: for == {N,N N - -.......1 do
3: X;-1 pe (=(())))= [ (xx,,)
4: end for
5: return {xn) Jn
advantages. 1) It is much faster, and benefit from adaptive ODE solvers bringing down the number of score evaluations needed, 2) the same ODE can be used to evaluate the log likelihood of the model by applying the probability flow formula of CNFs.
4.2 DENOISING DIFFUSION PROBABILISTIC MODEL
As described in Song et al. 2021), when using a finite number of steps, the forward diffusion process defined by Equation x
(corresponding times {0 S t, S T)=)) can be interpreted
Markov process:
p(xo.N) = p(xo)pex (x1|x)) ... (x, (x-1) ... PeN (xN (xN-1)
(12)
with the transition kernel pe (x++1 (x:) (x)+ X;, 11+1), where 6,+1 e(1++1) el4))
The idea of DDPMs is to introduce reverse Markov process defined in terms variational transi- tion kernels po 1
po (x = (po 1 -oo Px
(13)
While one could choose any distribution on SO(3) to parameterize this inverse transition kernel (e.g., Matrix Fisher, Bingham), we adopt for convenience an Isotropic Gaussian on SO(3) and use the following expression:
po p0 (0 (0 lz); so (2);44))
(14)
where 6y SO(3) R+ SO(3) is neural network predicting the residual rotation to apply to æ; to obtain the mean of the reverse kernel and en SO(3) x IR+ - R+ is neural network predicting the variance of this reverse kernel. To parameterize the output of 5e adopt the 6D continuous rotation representation of (Zhou et al. 2019) and explore the impact of this choice in Apppeddix D)
If the reverse Markov process can be successfully trained match the forward process, provides direct sampling strategy to generate samples from Po by initializing the chain from pr and iteratively sampling from the reverse kernel po(x; (#;-1
In DDPMs, the training strategy is to write down the Evidence Lower Bound (ELBO), given this variational approximation for the reverse Markov process, in order train the individual transition kernels Pa (x:- {3) To reduce the variance of this loss over a naive evaluation of the ELBO, Sohl- Dickstein et al. 2015) and Ho et al. (2020) propose to use a closed form expression of the reverse kernel X: X,, Xo) when conditioned on Xo- This makes it possible to rewrite the ELBO in terms of analytic KL divergences between Gaussian transitions kernels. However, contrary to the Gaussian case of Euclidean DDPMs, for IGsoØ) we do not easily have access to closed form expression of the reverse kernel p(2;-1 |r, 2o) which is needed to derive the training loss used Ho et al. 2020 The same approach cannot be applied.
Instead, we consider the expression for the ELBO:
Po (x; 1xx;) E [- log po (xo)] S E. log (xNN Σ log p(x; (x;-))
=: LELBO
(15)
which will be optimized by maximizing the log likelihood of individual transition kernlls
log pe (/x) over samples x;-1, X; obtained through simulating the forward Markov diffusion
process over the training set. Our strategy on SO(3), is therefore to train each transition kernel by maximum likelihood using the following loss function:
LDDPM
Σ 120
E (xn) (x. Ep,, (x))))x) [- (x,
(16)

---

Under review as a conference paper at ICLR 2023
where the log probability of the IGso3) distribution used in our parameterised reverse kernel is defined in Equution 6) While this loss can indeed be used to train DDPM (as demonstrated in the next section), compared to the strategy oo al. (2020). we expect it to suffer from larger variance and is not explicitly parameterised in terms of the score function (Song et al. (021). Once trained, the sampling strategy described in Algorithm to draw from the generative model.
RELATED WORK
Most related to work is Song al. (2021) which introduces the diffusion framework we use in this paper, and served as a point of reference throughout. We survey below related works that have developed methodologies represent distributions SO(3).
Directional statistics The classical approach for modeling distributions on SO(3) relies on (mix- tures) of analytic distributions defined over the group of rotations. Common examples of using such distributions for modeling uncertainties over orientations include the Bingham distribution (Peretroukhin et al.) 2020: Srivatsan et al. 2018a; Gilitschenski et al. 2020) or the matrix Fisher distribution (Mohlin et al. 2020). The two main issues of these approaches are the lack of flexibil- ity/expressivity of these analytic distributions, and the general difficulty of computing their normal- ization constant, which is typically required to train these models by maximum likelihood.
Normalizing Flows A number of approach have been proposed to build density estimators on manifolds (which include SO(3)) based on Normalizing Flows. first class of methods proposes to use conventional Euclidean Normalizing Flow in R which is then mapped to the target manifold using an invertible mapping (Gemici et al.) 2016 Falorsi et al., 2019). This has some limitations however as the target manifold needs to be homeomorphic to Ik (which is the case for SO(3)), and this mapping can also present discontinuities. As an improvement over this approach, a second class of methods based on continuous normalizing flows (Chen et al. 2018) has emerged, defining directly flows on the manifold (Falorsi & Forré 2020 Mathieu & Nickel 2020). These approaches remain relatively costly training requires backpropagating through ODE solver. [Rzen et al. (202)) proposes to sidestep that issue by training the CNF through penalizing the divergence of the neural network. And finally, in recent work (Ben-Hamu et al., (022)) proposes train flow on manifolds by penalizing a Probability Path Divergence (PPD).
Diffusion models In concurrent work, Leach on SO(3) by direct analogy with Ho et al. (2020), based on the Isotropic Gaussian on SO(3) replacement for the Normal distribution the loss function used in Euclidan DDPMs does not directly translate to SO(3), which leads to imperfect density estimation as we will illustrate in our experiments. Finally, De Bortoli et al. 2022); ; Huang al. (2022); Thornton et al. (2022) introduce generic frameworks for diffusion models on Riemannian manifolds but only for Score-Based Generative Model (SGM). Their generic approach means they do benefit from the knowledge of solution the heat equation in SO(3), which we use extensively in our work to avoid the need to simulate SDEs and to efficiently generate samples from the forward diffusion process. In addition, we note that the method developed in Huang et al. 2022) is not particularly efficient on the orthogonal group as it requires projection operation, which involves a singular value decomposition.
2022) proposed an implementation of DDPMs
R" However, as mentioned the previous section,
Other approaches Murphy et al. 2021 develops parametrric representation of distributions on SO(3) by introducing a neural network to implicitly represent an unnormalized density on SO(3). Training this model by maximum likelihood requires computing the normalization constant of this implicit probability density function through brute-force evaluation on a tiling of SO(3), which is very costly in memory and limits the effective resolution of the learned densities.
EXPERIMENTS
We investigate the quality of the generative models described in the previous section series of synthetic test densities on SO(3). Details of the training procedures and architectural choices for all models can be found in Appendix A.
7

---

Under review as a conference paper at ICLR 2023
Model
Checkerboard
4- G-Gusssians
3-Stripes
SGM on SO(3) (ours)
0.50 0.01
0.50 0.01
0.51 0.01
DDPM on SO(3) (ours)
0.52 0.01
0.53# 0.01
0.52 0.01
RSGM (De Bortoli al.
2022
0.51 0.01
0.51
Moser Flow el al.
2021
0.56= 0.01
0.60= 0.02
0,53 0.02
DDPM (Leach et 2022)
0.7 1 0.04
0.90=
0.60=
Implicit-PDI (Murphy et al.
2021
0.59 0.04
0.8 + 0.09
0.63+ 0.04
Table I: Sample quality metric from the C2ST (lower better). the learned distribution is identical to
original one, the metric should be ~ 0.5; if is significantly different, the metric tends towards The errors on the metric were obtain from the standard deviation of the metric over k-fold cross validation samples for single training of the model. indicate failure evaluate the metric for particular model.
Model
Checkerboard
4-Gaussians
3-Stripes
True
SGM (Ours)
DDPM (Ours)
RSGM, DDe Bortoli et al. 2022)
Moser Flow,
Rozen et al. 2021
Leach et al. 2022
Implicit-PDD Murphy et al.
2021)
Figure 3: Density plot comparing samples from learned synthetic densities on SO(3). For visualization this density plot shows the distribution of canonical axes of sampled rotations projected on the sphere; the tilt around that axis is discarded.
8

---

Under review as a conference paper at ICLR 2023
Image
True
Scatter (Predicted)
Density (Predicted)
Figure 4: Example of pose estimation using an SO(3) SGM conditioned images from the SYMSOL dataset Murphy al. )
Test densities on SO(3) We adopt three different toy distributions SO(3): checkerboard pat- tern, a multi-modal distribution of 4 concentrated Gaussians and a stripe pattern that can be viewed as circles the sphere. We focus evaluating the generative models in terms of the quality of their
sample generation using the Classifier 2-Sample Tests (C2ST) metric (Lopez-Paz & Oquab] [217]]
Dalmasso et al. 2020 Lueckmann et al., (021). The C2ST metric has been used in particular in the context of simulation based inference to quantify the quality of inferred distributions. Concisely, the C2ST method uses neural network classifier to discriminate between true and the generated sam- ples, yielding value of 0.5 if the two distributions are perfectly indistinguishable to the classifier, up to value of if they are extremely different. contrast to the usual Negative Log Likelihood (NLL), C2ST can be consistently computed for all generative models we compare bellow.
We present in Figure 3 and TTbll the results of our (SGM, DDPM- VExp, DDPM-VPreS) com- parisons on these test densities against the implicit-pdf method of Murphy et al. (202T). the DDPM implementation of Leach et al.] (2022), Moser flow of Rozen et al. (2021). and the Riemannian Score-Based Generative Model (RSGM) of De Bortoli et al. (2022) (trained under their 4njo score matching loss). We find that in all cases our SGM implementatron on SO(3) yields the best C2ST metric, which is in line with the visual quality of distributions shown in Figure 3 Our DDPM im- plementation on SO(3) yields distributions that are comparatively less sharp, which we attribute to the larger variance of our training loss for that model. Compared to other models, our experiments illustrate a failure mode in the method Leach et al.] (2022) which we attribute to the fact that the usual DDPM loss function cannot be directly translated SO(3) (as discussed in (subsection 4.2). We also note that the Implicit-PDD. model, in comparison, is extremely limited in resolution because of the memory cost of evaluating the pdf on a tiling of SO(3), and thus yields much lower scores. The best results after our method are achieved by the RSGM model (De Bortoli et al. 2022). which is expected due to its similarity with our work, but is slower to train in the specific case SO(3). We find that the cost of simulating the forward SDE in the training phase leads to factor x8 in computation time per batch on a given GPU.
Pose estimation To test practical applications of our model, following (Murphy et al. 1021 we used a vision description obtained from a pre-trained ResNet architecture with ImageNet weights consisting of 2048 dimensional vector to condition an SO(3) SGM. Using images of symmetric solids from the SYMSOL dataset Murphy et al. (2021) we show that we correctly estimate poses of objects with degenerate symmetry, as shown in Fig. (and Appendix B).
7 CONCLUSIONS AND DISCUSSION
In this paper, we have presented a framework for score-based diffusion generative models on SO(3), as an extension of Euclidean SDE-based models (Song et al. 2021). Because it is developed specif- ically for the SO(3) manifold, our work proposes a simpler and more efficient alternative to other recent (and general) Riemannian diffusion models while reaching state-of-the-are quality on syn- thetic distributions on SO(3). One of the most promising applications of this work is in robotics and computer vision, for the general task outperforms current baselines (Murphy et al.) 2021). In the natural sciences, generative models on SO(3) are also of great interest and can be used for instance to find the angle of a molecule that mini- mizes the binding energy. Finally we note that this work, one could define a Schrödinger bridge approach (De Bortoli et al. 2021 Thorton et al.) 2022) specifically for SO(3), which would improve both sampling efficiency and sample quality.
cose-estimation, where our proposed model significantly
interesting extension of the models presented in
9

---

Under review as a conference paper at ICLR 2023
REFERENCES
Heli Ben-Hamu, Samuel Cohen, Joey Bose, Brandon Amos, Maximiliian Nickel, Aditya Grover, Ricky T. Q Chen, and Yaron Lipman. Matching normalizing flows and probability paths on manifolds. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 1749-1763, PMLR, 17-23 Jul 2022. URL httpss::: / hamu2 html
Johann Brehmer and Kyle Cranmer. Flows for simultaneous manifold learning and density esti- mation. In Proceedings of the 34th International Conference Neural Information Processing Systems, NIPS' 20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713829546.
Nicola De Cao and Wilker Aziz. The power spherical distribution. ArXiv, abs/2006.0443)) 2020.
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural Ordinary Differential Equations. In Advances in Neural Information Processing Systems, volume 31. Cur- Associates, Inc., 2018. URL / nwwrips....... cc/paper//018/ http:///aaa.aaa.rrooooooooo/r/pa
htmll
Sean Craven, Djuna Croon, Daniel Cutting, and Rachel Houtz. Machine learning manifold. Phys. Rev. D. 105:096030, May 2022. doi: 10.11003/hyyRevD..000 1006000000 URL
aps aps.org/doi 11 www.1///////......e 10w. 011.10000099939
Niccolo Dalmasso, Ann Lee, Rafael Izbicki, Taylor Pospisil, Ilmun Kim, and Chich-An Lin. Valida- tion of approximate likelihood and emulator models for computationally intensive simulations. In Silvia Chiappa and Roberto Calandra (eds.), Proceedings of the Twenty Third International Con- ference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning .mlr.
Research, pp. 3349-3361. PMLR, 26-28 Aug 2020. URL pttps://proceetee press/vi 08 da htt.://///aammasss2aammm
html
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion Schr "odinger Bridge with Applications to Score-Basee Generative Modeling, December 2021. URL http:
//arxiv. wwtp://////////eeeeee 01357 arXiv:2106.01357 [cs, math, stat].
Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornton, Yee Whye Teh, and Arnaud Doucet. Riemannian Score-Based Generative Modeling, May 2022. URL http:// 022202002773 02763 arXiv:2202.02763 [cs, math, stat].
arxiv. org/
Luca Falorsi and Patrick Forré. Neural Ordinary Differential Equations on Manifolds, June 2020.
URL htt.httpp///xiiv org/
06663 wxXiv22006006663 [cs, stat].
Luca Falorsi, Pim de Haan, Tim R. Davidson, and Patrick Forré. Reparameterizing Distri- butions on Lie Groups, March 2019. URL sttt://arxiy: org htrg/abs/1990..025 arXiv:1903.0200 (es, math, stat].
Wendelin Feiten, Muriel Lang, and Sandra Hirche. Rigid motion estimation using mixtures of projected gaussians. In Proceedings of the 16th International Conference on Information Fusion, pp. 1465-1472, 2013.
M. E. Gaddes, A. Hooper, and M. Bagnardi. Using machine learning to automatically detect volcanic unrest in time series of interferograms. Journal of Geophysical Research: Solid Earth, 124(11): 12304-12333 12322, 2019. doi: https///bb.orgg//..00002/0110000
URL https: /tt::tss::///atttttt htttps:/aquuubss. /awpppbg.......
onl ne Terary.rr... ww...wiii. www...lll..ooomm ctm./yccoooooooooooommmm.....l.. doi /tt/i//bs////// 0ew.11...2222000 9JR01 7519
Mevlana C. Gemici, Danilo Rezende, and Shakir Mohamed. Normalizing Flows on Rie- mannian Manifolds, November 2016. URL http://arxiv.orr www.v/oorg/as///61...22
1rXiv:1611.02250 [cs, math, stat].
Igor Gilitschenski, Roshni Sahoo, Wilko Schwarting, Alexander Amini, Sertac Karaman, and
Daniela Rus. Deep orientation uncertainty learning based bingham loss. In International
Conference on Learning Representations, 2020. Utthttps: URJAttpss//oeeeeeeeee.cc.
wwwp//peee.ew..ooeme wwtp:/wwwwww......eeeeeeeeeee forum??
@daryloog5Kke
10

---

Under review as a conference paper at ICLR 2023
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sher- jil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence. and K.Q. Weinberger (eds.), Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014. URL hnttps::/pproceedingg nee..e cc//paperr2014/111e 5c2aeeb222551110055ccchhaaaafff Paper. .df
A. Grigoryan. Heat Kernel and Analysis on Manifolds. AMS/PP studies in advanced mathemat- ics. American Mathematical Society, 2009. ISBN 9780821849354, URL
google. com/books 2aa.ookKs...--RRRRAAAAAA.
Richard I. Hartley, Jochen Trumpf, Yuchao Dai, and Hongdong Li. Rotation averaging. International
Journal of Computer Vision, 103(3):267-300 2013. URL httpp:://bo..orr/11..000
s s11263-B 1011663012-0000 012-0601-0]
Shoubaneh Hemmati, Peter Capak, Milad Pourrahmani. Hooshang Nayyeri, Daniel Stern, Bahram Mobasher, Behnam Darvish, lary Davidzon, Olivier Ilbert, Daniel Masters, and Abtin Shahidi. Bringing manifold learning and dimensionality reduction to SED fitters. The Astrophysical Jour- nal, (1) aug 2019. doi: 1444477/0441-8113/bb3444 3847/204 821 a/b3448
Jonathan Ho, Ajay Jain, and Picter Abbeel. Denoising diffusion probabilistic models. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 6840-6851. Curran Asso- ciates, Inc., 2020. URL https::/pproeeeiigg. new.oig..euuipssp. cc/pppper/222//ille #c5bcfec834lff0d5G. flbll11 abw1ab9ca11bbbbbb -aapeaaaaaaaaa....r.ee
Sabera Hoque, Md. Yasir Arafat, Shuxiang Xu, Ananda Maiti, and Yuchen Wei. A comprehensive review on 3d object detection and 6d pose estimation with deep learning. IEEE Access, 9:143746- 143770, 2021. doi: 10. 109//CCCESS.021..11
Chin-Wei Wei Huang, Milad Aghajohari, Avishek Joey Bose, Prakash Panangaden, and Aaron Courville. Riemannian Diffusion Models, August 2022. URL http::/aaxx.v.orr/bbb///222 07949 arXiv:2208.0799 [cs].
Arich Iserles, Hans Z. Munthe-Katt Syvert P. Nørsett, and Antonella Zanna. Lie-group methods. Acta Numerica, 9:215-366 2000. doi: 10.1017/5099222290000155
Diederik P Kingma and Max Welling. Auto-Encoding Variational Bayes. arXiv e-prints, art. arXiv:1312.6114, December 2013.
Adam Leach, Sebastian M Schmon, Matteo T Degiacomi, and Chris G Willcocks. Denoising diffu- sion probabilistic models on so (3) for rotational alignment. In ICLR 2022 Workshop on Geomet- rical and Topological Representation Learning, 2022.
David Lopez-Paz and Maxime Oquab. Revisiting classifier two-sample tests. In International Con- ference on Learning Representations, 2017. URL /tttps://ppenreviiee.. www.ennnttfoooteeeeee
id=SJkXfE5xx
Jan-Matthis Lueckmann, Jan Boelts, David S. Greenberg, Pedro J. Gonçalves, and Jakob H. Macke.
Benchmarking simulation-base inference. In AISTATS, pp. 343 343-351, 2021. URL http:// .mw..r.........................m /etee
Elman Mansimov, Omar Mahmood, Seokho Kang. and Kyunghyun Cho. Molecular geometry pre- diction using a deep generative graph neural network. Scientific Reports, 9, 2019.
Emile Mathieu and Maximilian Nickel. Riemannian Continuous Normalizing Flows, December 2020. TtRLbhtt:///axxxii org / abss 20066 10605 arXiv:2006.1000 10605 [cs, stat].
Siegfried Matthies, J. Muller, and G. W. Vinel. On the normal distribution in the orientation space. Textures and Microstructures, 10:77-96, 1988.

---

Under review as a conference paper at ICLR 2023
D. Mohlin, G. Bianchi, and J. Sullivan. Probabilistic orientation estimation with matrix Fisher distributions, June 2020. [cs].
arXiv:2006.09740
Kieran A Murphy, Carlos Esteves, Varun Jampani, Srikumar Ramalingam, and Ameesh Makadia. Implicit-pdf: Non-parametric representation of probability distributions on the rotation manifold. In International Conference on Machine Learning, pp. 7882-7893. PMLR, 2021.
Dmitry Nikolayev and Tatjana Savyolov. Normal distribution on the rotation group so (3). Tex- tures and Microstructures, 29, 1970.
Valentin Peretroukhin, Matthew Giamou, W. Nicholas Greene, David Rosen, Jonathan Kelly, and Nicholas Roy. A Smooth Representation of Belief over SO(3) for Deep Rotation Learning with Uncertainty. In Proceedings of Robotics: Science and Systems, Corvalis, Oregon, USA, July 2020. doi: 10. 1507//SSS 20000XX1.00000
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Francis Bach and David Blei (eds.), Proceedings the
the 32nd International Conference Machine Learn- ing. volume 37 of Proceedings of Machine Learning Research, pp. 1538, Lille, France, 07-09 Jul 2015. PMLR. [ttt::::////rreeenngg html
.om.smmmlr. wwwpinngs.mrr.....ee... vv7 r377reeeede1....
Noam Rozen, Aditya Grover, Maximilian Nickel, and Yaron Lipman. Moser Flow: Divergence- based Generative Modeling on Manifolds, November 2021. URL
2108. 08052 1XXi:22108080055 [cs, stat].
Hyunwoo Ryu, Jeong-Hoon Lee, Hong-in Lee, and Jongeun Choi. Equivariant Descriptor Fields: Equiia-quiaaaiaat Energy-Based Models for End-to-End Visual Robotic Manipulation Learning, June 2022. URL
08321 McXiv:2206.0832. [cs].
Jascha Sohl-Dickstein. Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Francis Bach and David Blei (eds.). Pro- ceedings of the 32nd International Conference on Machine Learning. volume 37 of Proceedings of Machine Learning Research, pp. 2256-2265, Lille, France, 07-09 Jul 2015. PMLR. URL bottps://procest
www.mpr...esssemmmmmmmm /vveess/v37// /shhl/ohl---- weallldiccsteennn html
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distri- bution. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019. URL 1001ef257407d5aa3 (117662227777771111666666228-3-3
pdf
Yang Song and Stefano Ermon. Tecoo for Training Score-Based Generative Models, October 2020.
09011 arXiv:2006.0900 00001 [cs, stat].
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- tional Conference Learning Representations, 2021. URL htttps::/cpenrevviee net./
Serrmmid=PxTIGII
Rangaprasad Arun Srivatsan, Mengyun Xu, Nicolas Zevallos, and Howie Choset. Probabilistic pose estimation using a bingham distribution-baun based linear filter. The International Journal of Robotics Research. 17(33-14)1610016331 2018a. doi: 10. 10.117770273366918877333.. doi /tt.oo///11..117/01110111
URL https://
Rangaprasad Arun Srivatsan, Mengyun Xu, Nicolas Zevallos, and Howie Choset. Probabilistic pose estimation using bingham distribution based linear filter. The International Journal of Robotics Research, 37(13-14:1610--631 2018b. doi: 10.1177007833:9118888333. URL https:// doi. 20o..org/10011777778888833333333
James Thomton, Michael Hutchinson, Emile Mathieu, Valentin De Bortoli, Yee Whye Tch, and Ar- naud Doucet. Riemannian Diffusion Schr', "'odinger Bridge, July 2022. URL :atpp:/ /arxiv.
org/ abs//ab//ii 2207. 03024 arXiv:2207.053: 003024 [es, stat].
12

---

Under review as a conference paper at ICLR 2023
Richard von Mises. Uber die "ganzzahligkeit' der atomgewicht und verwandte fragen. Physikal. Z. 19:490-500, 1918.
Nathan Zelesko, Amit Moscovich, Joe Kileel, and Amit Singer. Earthmover- manifold learning for analyzing molecular conformation spaces. In 2020 IEEF 17th International Symposium on Biomedical Imaging (ISBI), pp. 1715-1719, 2020. doi: 10.1 1.1109951145749 2020.9098723.
Yi Zhou, Connelly Barnes, Jingwan Lu, Jimei Yang, and Hao Li. On the continuity of rotation representations in neural networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 5745-5753 2019.
A IMPLEMENTATION AND TRAINING
We designed our neural networks with a size of {256.256,25..56.222 neurons each with leaky ReLU activation and with residual connection. The neural networks were implemented using the axis-angle representation of SO(3), i.e. the input and output elements were represented using axis- angle representation. Additionally, the neural networks were conditioned on the noise scheduler and the noise scales were also learned parameters. We trained our models using the Adam optimizer with learning rate of 10-4, exponential decay rates of 3 0.90 and B2 = 0.95, 400 000 iterations, and a batch size of 1024. NVIDIA Tesla V100 GPU was used as the hardware, with JAX and DeepMind-Hailee Python libraries as the software.
For the DDPM, adopt in practice the Variance Preserving SDE of (Ho et al. 2020 obtain better results empirically than with a Variance Exploding SDE.
B ADDITIONAL POSE ESTIMATION RESULTS
Here provide on Figure additional results on pose estimation.
C REPRESENTATIONS OF SO(3)
The special orthogonal group. SO(3), is the Lie group of all rotations about the origin in 3- dimensional space. There several ways to represent the elements of the group SO(3), each with its advantages and disadvantages:
Rotation Matrices E pax3 with determinant equal to 1. This representation has 9 parame- ters and can be subject to some numerical stabilities, such when computing the inverse or trigonometric functions.
Euler angles (also called yaw, pitch, and roll in robotics) three angles 3, so that can describe an orientation with respect to fixed coordinate system. This representation is subject to the infamous Gimbal lock, where one degree of freedom is lost when two axes of the gimbal become parallel.
Unit Quaternions defined as y = a bi cj- dk, where a, b, c, d are real number satisfy- ing Va2 12 + 2² + d = with i+ k denoting the vector (or imaginary) part of the unit quaternion. This representation has parameters and has elegant operations (Hamilton product) without trigonometric functions However, quaternions are antipodally symmet- rie which introduces some degeneracies.
Axis angle representation (normalized), Tangent space (unnormalized ) defined as fe (01, flg, 02) ((e,, C2, ez), where is the rotation angle and is the rotation axis. However, this representation does not have well defined operations to combine rotations, and is furthermore discontinuous 0 (Zhou et al.! 2019).
Therefore, in practice it is best to and convert back and forth among them. For comprehensive review on SO(3) representations and metrics, especially for computer scientists, please refer to (Hartley et al. 2013
combinations of the aforementioned representations
13

---

Under review as a conference paper at ICLR 2023
Image
True
Scatter (Predicted)
Density (Predicted)
2
m
2
S
Figure 5: Predicted poses for an image of solid with degenerate symmetry, here we only show it for the cone. The 1st column depicts the image of the symmetric solid. In column 2, each point represents a rotation matrix in SO(3) projected the sphere according to its canonical axis, the color indicates the tilt around that axis. For visualization density plot (column shows the distribution canonical axes of sampled rotations projected on sphere; tilt around that axis discarded.
14

---

Under review as a conference paper at ICLR 2023
(a) DDPM trained with axis angle inputs and out- puts.
(b) DDPM trained with 3x3 rotation matrices as inputs and using continuous 6D output rotation representation.
Figure 6: Comparison of distributions sampled from the DDPM under two different parameteriza- tions of both input and output rotations. Once the models fully trained illustrated here, the impact is small, but on a partially trained network discontinuities would be visible in the case of the axis-angle representation, mostly due the discontinuity in the input rotations.
D IMPACT OF ROTATION REPRESENTATIONS ON NEURAL DIFFUSION MODEL
As highlighted in Zhou et al. (2019). particular choice of rotation representation can affect the training and accuracy of neural networks which either take rotations an input or that output rota- tions. In particular, common representations such as axis-angle and quaternions are known to have discontinuities, which are needlessly difficult to capture for neural network. In that work, they propose in particular to use 5 or 6 dimensional representations which have the particularity of being continuous and demonstrate their benefit in neural network training.
In work, make the choice of providing as inputs of the neural networks directly the 3x3 rotation matrix. Only the network involved the DDPM needs to represent rotations as an output, and there adopt the 6D representation following Zhou al. (2019), which can be seen two 3D vectors, from which we can build a full orthogonal rotation matrix using cross- rross-products.
In comparing the impact of this choice against using only axis-angle representation as inputs and outputs, we observe the following points:
The choice of the output parameterization (in the case the DDPM) has no noticeable effect, which is expected the model outputs residual rotations, which remain small and thus away from the discontinuity in the axis-angle representation.
While the differences are small once the networks fully trained, illustrated Fig- ure 6 we notice for partially trained networks discontinuity the sampled distributions in the of an input axis- axis-angle representation. Therefore, directly feed the 3x3 rotation matrix as an input to our networks.
E QUATERNION OPERATIONS FOR SO(3)
Quaternions form a group under multiplication, defined by the Hamilton product. Given quternions Y1 and ?2, the Hamilton product defined by carrying out the ?1 %/2 (az + b,i + eij + d,k). (a2 + b,i + c2j + dzk) in distributive manner, keeping in mind the basis multiplication identities. This operation physically equivalent to rotating by 7 and then by 72. The identity of the group is the quaternion 7o 1 + 0i + Oj + Ok and the inverse of y* (also conjugate) is defined as dk.
** =
The reparametrization trick and variance preserving quaternions In Euclidean space for a Gaussian random variable from N(j II), the reparametrization trick is defined μ 2
15

---

Under review as a conference paper at ICLR 2023
where 8 ~ ~N(00;) I). We can define analogous operation in the quaternion group, such:
?
(17)
Here, the quternion raised to scalar power is defined the 7"t exp(In(7)a), in other words we take the quaternion to the tangent space from the manifold, perform the operation of multiplica- tion and bring back into the manifold by using the exponential map. and the variance preserving operation analogous Va v(1 a:) be defined as
x = x a 8V (I-a)
(18)
For the variance exploding case, we can directly sample from the heat kernel without resorting to the quaternion operations.
E.1 DISTRIBUTIONS ON SO(3)
In literature there are numerous ways to represent distributions on the hypersphere. Most of them in- volve taking standard distribution from the Euclidian space R" and then constraining projecting them on to the hypersphere S Some the popular distributions are:
Projected Gaussian(s) on the sphere where standard Gaussian(s) on the tangent space of the hypersphere are projected via central projection, as done in Feiten et al. (2013);
the von Mises-Fisher (yMF) distribution where an isotropic Gaussian on R² is restricted
the unit hypersphere von Mises 1918);
The recently developed Power Spherical distribution Cao & Aziz 20200, which addresses some of the challenges of the vMF distribution, such as numerical stability and scalability.
The antipodally symmetric Bingham distribution. The antipodal symmetry makes it suitable distribution to represent quaternions, since quaternions double cover the space of rotations on SO(3) Gilitschenski et al.] (2020); [Peretroukhin et al. (2020): (Srivatsan et al.) (2018b ab. However, the Bingham distribution is notorious for its normalization constant that is very hard to compute.
Unfortunately, these distributions are not closed under convolution (i.e. composition of their random variables), thus writing down closed form diffusion process akin the Euclidean Gaussian case intractable. One way to circumvent this problem is to use class of functions that are closed under convolutions the manifold. An obvious choice is the heat kemnel which is the canonical solution to the diffusion equation and is closed under convolutions by definition GGigoryan 2009).
16


Reviews:
[{"review_id": "jGo7xKFtJu", "reviewer": "Reviewer_aBeQ", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a generative diffusion model for SO(3) densities using the heat kernel on SO(3). The authors design both a score-based forward pipeline and well as a denoising pipeline for SO(3). Experiments are done on synthetic datasets of rotation samples. \n", "summary_of_the_review": "This paper would be an accept if it included experiments conditioned on input images or point clouds. It is the first diffusion on SO(3) as far as I know. Adopting the heat kernel is a great idea, and the implementation of the score-based model and the denoising model are carefully done. \n", "strength_and_weaknesses": "+: The main strength of the paper is the use of a heat kernel on SO(3).\n\n+: The authors carefully chose a tractable heat kernel and how to parametrize the isotropic Gaussian as well as explaining its advantages over Bingham and Fisher matrix distribution (being closed wrt convolutions).\n\n+: Thanks to the heat kernel there is no need to solve for the stochastic DE. \n\n+: Isotropic Gaussian enables a closed form in the KL-divergence term of the ELBO.\n\nAuthors could respond to the following weaknesses:\n\n-: The mapping from the axis-angle representation to SO(3) is not continuous, and it might affect the stability of the network components of the DE solutions. An example of a continuous mapping is the two first columns of the rotation matrix (a Stieffel manifold) as explained in Zhou et al.'s paper (On the continuity of rotation representations).\n\n-: The experimental evaluation is very limited regarding datasets and evaluation metrics. The significance of density estimators is in capturing multimodal distributions arising in real data. There is abundant real data in 6dof object pose in images and point clouds as well as in all joint orientations in human poses. Many insights can be gained by conditioning the diffusion on such inputs. State-of-the-art approaches like citations Murphy et al. or Mohlin et al. have used such data.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "RbcfIVGHJoZ", "reviewer": "Reviewer_XDGv", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper extends diffusion generative models to SO(3) manifold. It provides implementations for score-based generative models and denoising diffusion probabilistic models. The proposed methods are applied to synthetic densities on SO(3).", "summary_of_the_review": "The paper is heavy in mathematical formulation and theory, but is weak in experiments and applications.", "strength_and_weaknesses": "Strength:\n- Reformulate diffusion models on SO(3) and the loss functions to train the models.\n- The proposed method is better than others in the experiment of synthetic densities.\n\nWeakness\n- Experiments are weak, with only three synthetic densities. The paper claims the proposed method has efficient training, but it’s not evaluated in the experiments.\n- Applications are also weak. The paper mentioned the proposed method can be used for pose estimation, and claims it’s better than previous work. However, it doesn’t explain how the diffusion generative models can be applied to pose estimation, and there’s no experiment to support the claim.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "NfmoMEK2vGD", "reviewer": "Reviewer_aEa8", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper builds on top of recent, good papers. Diffusion-based generative models have been recently lifted from Euclidean spaces to compact, connected Riemannian spaces (De Bortoli et al, 2022), where given a family of \"normal\" distributions, forward diffusion processes are proved to admit backward diffusion processes similarly to Euclidean spaces. Normal distributions on those spaces are required to be closed under convolution so that the backward processes can be tractable However, such distributions are typically very cumbersome to deal with as they often exist as infinite serieses. Many approximations have to be introduced. The paper specialises to SO(3), and uses a classical, but less well-known family of \"normal\" distributions on SO(3) called Isotropic Gaussian distributions $\\mathcal{IG}_{SO(3)}$ (Nikolayev & Savyolov, 1970), which is not only closed under convolution but also much easier to work with from the calculus point of view.\n\nThe paper goes on and derives an equivalent of score-based generative models (SGMs) on SO(3), and an equivalent of (the spirit of) denoising diffusion probabilistic models (DDPMs) on SO(3). It is shown theoretically that SGM on SO(3) enjoys a lot of similar properties as the Euclidean SGM, as the cumbersome normalizing factor of $\\mathcal{IG}_{SO(3)}$ distributions can be mostly avoided.\n\nIn contrast, in the DDPM on SO(3) case, because there is no closed form for the distance between two distributions in $\\mathcal{IG}_{SO(3)}$, the ELBO cannot be reduced to an analytical form during training.\n\nExperiments on synthetic data comparing the proposed variants showing results favouring SGM on SO(3). However, there is no results on real-world data.", "summary_of_the_review": "I thank the authors for having submitted this paper. I have had a good read.\n\nThe paper has the right motivation. Its theoretical side is very good. There are minor points of improvements for DDPM on SO(3). However, experiments are restricted to synthetic data only, and there is no result related to training time or inference time compared to the existing best method RSGM, which is a bit unfortunate because the paper promotes efficient computations on SO(3).\n\nNonetheless, I find the pros outweigh the cons.", "strength_and_weaknesses": "## Strength\n\nThe paper is very well-written. It has an excellent review of existing works.\n\nThe idea of using Isotropic Gaussian distributions for diffusion generative models on SO(3) has been recently introduced in (Leach et al 2022). However, their approach is restricted to DDPM, while this paper is a more complete treatment of the same topic.\n\n## Weaknesses\n\nWhile SGM on SO(3) enjoys many properties similar to Euclidean SGM, I think the treatment of DDPM on SO(3) can be further improved. In the paper, the authors mentioned that they did not try out any other distribution families like matrix Fisher distributions or Bingham distributions for the reverse diffusion process but instead used Isotropic Gaussians for convenience. This part can be elaborated further. What convenience do you get by staying with Isotropic Gaussians? Why not try other families? In my view, once you have set up a reverse Markov process as machine-learnable, you are free to choose the distributions that work best for your data. As mentioned in the paper, by staying with Isotropic Gaussians you have to accept an inconvenience in that there is no analytical form for the KL divergence between two distributions of such kind, leading to an extra forward Markov process during training. Why not use a distribution family where the KL divergence has an analytical form?\n\nI think it makes sense to report some comparisons in training times and inference times between SGM on SO(3) and its closest rival RSGM, since RSGM is very close to SGM on SO(3) in spirit but SGM on SO(3) enjoys a much more efficient implementation thanks to the Isotropic Gaussians. Speed-related results are somewhat lacking in the paper.\n\nEquation (14) is somewhat puzzled to me. In the current draft, it looks like the network predicts the mean rotation using the axis-angle representation. However, the axis-angle representation does not work very well when the mean rotation angle approaches $\\pi$. Perhaps the authors meant that the function $\\mu$ represents the relative rotation (in axis-angle form) from $\\mathcal{x}_i$ to the predicted mean rotation, thereby actually having $\\mathcal{x}_i$ next to $\\exp(\\cdot)$ in the equation?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject