Unified framework for diffusion generative models in SO(3): applications in computer vision and astrophysics

TL;DR

This paper extends diffusion generative models to the manifold of 3D rotations, SO(3), enabling applications in computer vision and astrophysics with state-of-the-art results and efficient training methods.

Contribution

It introduces a novel framework for diffusion models on SO(3), leveraging its tractable heat diffusion, and demonstrates their effectiveness in pose estimation and galaxy orientation prediction.

Findings

Achieved state-of-the-art results on synthetic SO(3) densities.

Successfully applied models to pose estimation tasks.

Predicted correlated galaxy orientations effectively.

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/cosmology 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. Additionally, we demonstrate the practicality of our model on pose estimation tasks and in predicting correlated galaxy orientations for astrophysics/cosmology.