DIFFUSION GENERATIVE MODELS ON SO(3)

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.

1. INTRODUCTION

Deep generative models (DGM) are trained to learn the underlying data distribution and then generate 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, a new class of DGMs based on Diffusion, such as Denoising Diffusion Probabilistic Models (DDPM) (Ho et al., 2020) and Score Matching with Langevin Dynamics (SMLD) , a subset of general score-based generative models (SGMs), (Song & Ermon, 2019) , have achieved state-of-the-art quality in generating images, molecules, audio and graphs 1 (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, Song 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 a noise diffusion process defined by a 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 to manifold-valued data, which are ubiquitous in many scientific disciplines. Examples include high-energy physics (Brehmer & Cranmer, 2020; Craven et al., 2022) , astrophysics (Hemmati et al., 2019) , geoscience (Gaddes et al., 2019) , and biochemistry (Zelesko et al., 2020) . 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 of particularly high interest in many fields including for instance in robotics (estimating the pose of an object, Hoque et al. 2021); and in biochemistry (finding the conformation angle of molecules that minimizes the binding energy, Mansimov et al. 2019) . Contrary to 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. given base distribution (right most, denoted by circles) can be evolved under the probability flow ODE (Eq. 3) towards a noisy distribution (left most), or vice-versa from the noisy distribution back to the target distribution. Each point represents a 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 at this link.



For a comprehensive list of articles on score-based generative modeling, see https:// scorebasedgenerativemodeling.github.io/



Figure 1: Illustration of reversible diffusion of a mixture of two IG SO(3) blobs on SO(3). Samples from a

annex

The contributions of our paper are summarized as follows:• We reformulate Euclidean diffusion models on 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 Denoising Diffusion Probabilistic Models specialized for SO(3).• We reach a new state-of-the-art in sample quality on synthetic SO(3) distributions with our proposed SO(3) Score-Based 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 R 3 , which can be parameterised in terms of a vector in R 3 , which corresponds in turn to the axis-angle representation of rotation matrices.For any rotation x ∈ SO(3), its axis-angle representation ω ∈ R 3 can be computed as ω = ωv with ω = arccos 2 -1 (tr(x) -1) ∈ (-π, π] and v = 1 2 sin ω (x 32 -x 23 , x 13 -x 31 , x 21 -x 12 ) a unit vector of R 3 .We direct the interested reader to more information on representations of SO(3) in Appendix C.

3. DIFFUSION PROCESS ON SO(3)

Similarly to Euclidean diffusion models (Song et al., 2021) , we begin by defining a Brownian noising process that will be used to perturb the data distribution. Let us assume a Stochastic Differential Equation of the following form:where w is a Brownian process on SO(3), f (• , t) : SO(3) → T x SO(3) is a drift term, and g(•) : R → R is a diffusion term. If we sample initial conditions for this SDE at t = 0 from a given data distribution x(0) ∼ p data , we will denote by p t the marginal distribution of x(t) at time t > 0. Thus p 0 = p data , and at final time T at which we stop the diffusion process p T will typically tend to a known target distribution that will be easy to sample from. 

