Research map / Probabilistic and scientific modeling / Molecules and proteins
Rao-Blackwell Gradient Estimators for Equivariant Denoising Diffusion
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Orbit Diffusion (OrbDiff), a Rao-Blackwellized gradient estimator for training symmetric (equivariant or augmented) diffusion models that replaces the noisy augmented target with an orbit-weighted conditional mean, reducing variance at no extra forward passes.
Problem
Training diffusion models to learn G-invariant distributions, either with equivariant architectures (which are hard to optimize) or with data augmentation (which yields high-variance gradients and may not capture symmetry fully).
Approach
Interpret data augmentation as a Monte Carlo estimator of the symmetrized diffusion loss gradient and apply Rao-Blackwellization, replacing the sampled x0 target with an orbit-weighted conditional expectation E[x0|xt] estimated via self-normalized importance sampling over the group orbit of the training example, with a single forward/backward pass per sample.
Key results
- On Perov-5 text-guided CSP (TGDMat short), OrbDiff_WN raises match rate from 59.39% to 65.57%.
- On MP-20 with DiffCSP, OrbDiff_WN improves match rate from 51.49% to 55.70%, with RMSE 0.053 vs 0.063.
- Finetuning ProteiNA with OrbDiff raises designability to 95.6% and lowers scRMSD to 0.93, versus naive finetuning.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add OrbDiff_WN to DiffCSP (ablation) | DiffCSP · Perov-5 | Match: 52.02 -> 52.39 | Helped |
| Add OrbDiff_WN to DiffCSP (ablation) | DiffCSP · Perov-5 | RMSE: 0.076 -> 0.069 | Helped |
| Add OrbDiff_WN to DiffCSP (ablation) | DiffCSP · MP-20 | Match: 51.49 -> 55.70 | Helped |
| Add OrbDiff_WN to DiffCSP (ablation) | DiffCSP · MP-20 | RMSE: 0.063 -> 0.053 | Helped |
| Use OrbDiff_U instead of the base objective (ablation) | DiffCSP · Perov-5 | RMSE: 0.076 -> 0.078 | Hurt |
| Use OrbDiff_U instead of the base objective (ablation) | DiffCSP · MP-20 | Match: 51.49 -> 54.47 | Helped |
| Add OrbDiff_WN (ablation) | TGDMat (S) · Perov-5, short descriptions | Match: 59.39 -> 65.57 | Helped |
| Add OrbDiff_WN (ablation) | TGDMat (S) · MP-20, short descriptions | Match: 59.90 -> 61.29 | Helped |
| Use OrbDiff_U instead of the base objective (ablation) | TGDMat (S) · MP-20, short descriptions | Match: 59.90 -> 56.50 | Hurt |
| Add OrbDiff_WN (ablation) | TGDMat (L) · MP-20, long descriptions | Match: 61.91 -> 66.74 | Helped |
| Add Orbit Diffusion to EquiNet (ablation) | EquiNet · Synthetic reflection data | RMSD: roughly 25× lower RMSD | Helped |
| Add data augmentation to an equivariant denoiser (ablation) | EquiNet · Synthetic reflection data | training loss fluctuations | No clear effect |
Surprising findings
- Standard data augmentation gives no variance benefit for already-equivariant denoisers, as the augmented and non-augmented losses are equivalent.
- The authors could not reproduce the reported ETFlow performance and report all results under their own evaluation protocol.
Limitations
- Orbit sampling constrains the choice of proposal distributions and may limit flexibility for diverse or adaptive noise structures.
- The number of orbit samples trades off computational cost against estimation accuracy.
- The OrbDiff target differs from the exact conditional expectation, so the method relies on biased-proposal design choices.
Experimental setup
Datasets: GEOM-QM9 Perov-5 MP-20 ProteiNA protein structure data
Benchmarks: GEOM-QM9 Perov-5 MP-20 protein designability benchmark
Metrics: Cov@0.1 Recall AMR Precision Match rate RMSE designability scRMSD diversity novelty RMSD
Baselines: GeoMol Torsional Diffusion MCF ETFlow DiffCSP P-cG-SchNet CDVAE TGDMat FoldFlow (OT) ProteiNA
Builds on: denoising diffusion ETFlow DiffCSP ProteiNA Rao-Blackwell theorem
Closest papers by meaning
- Scalable Non-Equivariant 3D Molecule Generation via Rotational Alignment (ICML 2025 · similarity 0.90)
- Equivariant Diffusion for Crystal Structure Prediction (ICML 2024 · similarity 0.88)
- Equivariant Blurring Diffusion for Hierarchical Molecular Conformer Generation (NeurIPS 2024 · similarity 0.88)
- SymDrift: One-Shot Generative Modeling under Symmetries (NeurIPS 2026 · similarity 0.87)
- Equivariant Denoisers Cannot Copy Graphs: Align Your Graph Diffusion Models (ICLR 2025 · similarity 0.87)
- SymDiff: Equivariant Diffusion via Stochastic Symmetrisation (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.