Rao-Blackwell Gradient Estimators for Equivariant Denoising Diffusion

TL;DR

This paper introduces a Rao-Blackwell gradient estimator for equivariant denoising diffusion models, reducing variance and improving stability in training for molecular and protein generation tasks.

Contribution

It proposes a novel variance reduction technique using Rao-Blackwellization for equivariant diffusion models, with a practical implementation called Orbit Diffusion.

Findings

Achieves state-of-the-art results on molecular conformation generation

Improves crystal structure prediction accuracy

Enhances protein structure generation capabilities

Abstract

In domains such as molecular and protein generation, physical systems exhibit inherent symmetries that are critical to model. Two main strategies have emerged for learning invariant distributions: designing equivariant network architectures and using data augmentation to approximate equivariance. While equivariant architectures preserve symmetry by design, they often involve greater complexity and pose optimization challenges. Data augmentation, on the other hand, offers flexibility but may fall short in fully capturing symmetries. Our framework enhances both approaches by reducing training variance and providing a provably lower-variance gradient estimator. We achieve this by interpreting data augmentation as a Monte Carlo estimator of the training gradient and applying Rao-Blackwellization. This leads to more stable optimization, faster convergence, and reduced variance, all while requiring only a single forward and backward pass per sample. We also present a practical implementation of this estimator incorporating the loss and sampling procedure through a method we call Orbit Diffusion. Theoretically, we guarantee that our loss admits equivariant minimizers. Empirically, Orbit Diffusion achieves state-of-the-art results on GEOM-QM9 for molecular conformation generation, improves crystal structure prediction, and advances text-guided crystal generation on the Perov-5 and MP-20 benchmarks. Additionally, it enhances protein designability in protein structure generation. Code is available at: https://github.com/vinhsuhi/Orbit-Diffusion.git.