# Rao-Blackwell Gradient Estimators for Equivariant Denoising Diffusion

arXiv: https://arxiv.org/abs/2502.09890 · HTML: https://mlsanity.com/paper/2502.09890/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 53 of its 73 references that are in the map.
- First posted on arXiv 14 Feb 2025.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 30 of 53)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [Variational Diffusion Models](https://mlsanity.com/paper/2107.00630/) (2021)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [Equivariant Diffusion for Molecule Generation in 3D](https://mlsanity.com/paper/2203.17003/) (2022)
- [SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks](https://mlsanity.com/paper/2006.10503/) (2020)
- [DiGress: Discrete Denoising diffusion for graph generation](https://mlsanity.com/paper/2209.14734/) (2022)
- [GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation](https://mlsanity.com/paper/2203.02923/) (2022)
- [MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields](https://mlsanity.com/paper/2206.07697/) (2022)
- [Vector Neurons: A General Framework for SO(3)-Equivariant Networks](https://mlsanity.com/paper/2104.12229/) (2021)
- [Harmonic Networks: Deep Translation and Rotation Equivariance](https://mlsanity.com/paper/1612.04642/) (2016)
- [Crystal Diffusion Variational Autoencoder for Periodic Material Generation](https://mlsanity.com/paper/2110.06197/) (2021)
- [Torsional Diffusion for Molecular Conformer Generation](https://mlsanity.com/paper/2206.01729/) (2022)
- [Equivariant Flows: Exact Likelihood Generative Learning for Symmetric Densities](https://mlsanity.com/paper/2006.02425/) (2020)
- [GEOM: Energy-annotated molecular conformations for property prediction and molecular generation](https://mlsanity.com/paper/2006.05531/) (2020)
- [SE(3) diffusion model with application to protein backbone generation](https://mlsanity.com/paper/2302.02277/) (2023)
- [Structure-based Drug Design with Equivariant Diffusion Models](https://mlsanity.com/paper/2210.13695/) (2022)
- [3D Equivariant Diffusion for Target-Aware Molecule Generation and Affinity Prediction](https://mlsanity.com/paper/2303.03543/) (2023)
- [SE(3)-Stochastic Flow Matching for Protein Backbone Generation](https://mlsanity.com/paper/2310.02391/) (2023)
- [Crystal Structure Prediction by Joint Equivariant Diffusion](https://mlsanity.com/paper/2309.04475/) (2023)
- [Equivariant flow matching](https://mlsanity.com/paper/2306.15030/) (2023)
- [Deep neural network solution of the electronic Schr"odinger equation](https://mlsanity.com/paper/1909.08423/) (2019)
- [Equivariant 3D-Conditional Diffusion Models for Molecular Linker Design](https://mlsanity.com/paper/2210.05274/) (2022)
- [Protein structure generation via folding diffusion](https://mlsanity.com/paper/2209.15611/) (2022)

## Built on by (most cited first, 5 of 5)

- [Mofasa: A Step Change in Metal-Organic Framework Generation](https://mlsanity.com/paper/2512.01756/) (2025)
- [Scaling Laws and Symmetry, Evidence from Neural Force Fields](https://mlsanity.com/paper/2510.09768/) (2025)
- [SymDrift: One-Shot Generative Modeling under Symmetries](https://mlsanity.com/paper/2605.06140/) (2026)
- [Budgeted Quotient-Residual Guidance for Frozen Pocket-Conditioned Molecular Diffusion](https://mlsanity.com/paper/2609.31222/) (2026)
- [Characterizing and Mitigating Distributional Bias in Physics-Informed Generative Models](https://mlsanity.com/paper/2610.16582/) (2026)

## Abstract (arXiv)

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.