Rao-Blackwell Gradient Estimators for Equivariant Denoising Diffusion

• URL: http://arxiv.org/abs/2502.09890v3
• Date: Mon, 16 Jun 2025 18:20:12 GMT
• Title: Rao-Blackwell Gradient Estimators for Equivariant Denoising Diffusion
• Authors: Vinh Tong, Trung-Dung Hoang, Anji Liu, Guy Van den Broeck, Mathias Niepert,
• Abstract summary: In domains such as molecular and protein generation, physical systems exhibit inherent symmetries that are critical to model.<n>We present a framework that reduces training variance and provides a provably lower-variance gradient estimator.<n>We also present a practical implementation of this estimator incorporating the loss and sampling procedure through a method we call Orbit Diffusion.
• Score: 41.50816120270017
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Zero-Variance Gradients for Variational Autoencoders [32.818968022327866]
  Training deep generative models like Variational Autoencoders (VAEs) is often hindered by the need to backpropagate gradients through sampling of their latent variables.<n>In this paper, we propose a new perspective that sidesteps this problem, which we call Silent Gradients.<n>Instead of improving estimators, we leverage specific decoder architectures analytically to compute the expected ELBO, yielding a gradient with zero variance.
  arXiv  Detail & Related papers  (2025-08-05T15:54:21Z)
• Learning (Approximately) Equivariant Networks via Constrained Optimization [25.51476313302483]
  Equivariant neural networks are designed to respect symmetries through their architecture.<n>Real-world data often departs from perfect symmetry because of noise, structural variation, measurement bias, or other symmetry-breaking effects.<n>We introduce Adaptive Constrained Equivariance (ACE), a constrained optimization approach that starts with a flexible, non-equivariant model.
  arXiv  Detail & Related papers  (2025-05-19T18:08:09Z)
• Energy-Based Coarse-Graining in Molecular Dynamics: A Flow-Based Framework Without Data [0.0]
  We introduce a data-free generative framework for coarse-graining that directly targets the all-atom Boltzmann distribution.<n>A potentially learnable, bijective map from the full latent space to the all-atom configuration space enables automatic and accurate reconstruction of molecular structures.
  arXiv  Detail & Related papers  (2025-04-29T17:05:27Z)
• Preconditioned Inexact Stochastic ADMM for Deep Model [35.37705488695026]
  This paper develops an algorithm, PISA, which enables scalable parallel computing and supports various second-moment schemes. Grounded in rigorous theoretical guarantees, the algorithm converges under the sole assumption of Lipschitz of the gradient. Comprehensive experimental evaluations for or fine-tuning diverse FMs, including vision models, large language models, reinforcement learning models, generative adversarial networks, and recurrent neural networks, demonstrate its superior numerical performance compared to various state-of-the-art Directions.
  arXiv  Detail & Related papers  (2025-02-15T12:28:51Z)
• Spatially-Aware Diffusion Models with Cross-Attention for Global Field Reconstruction with Sparse Observations [1.371691382573869]
  We develop and enhance score-based diffusion models in field reconstruction tasks. We introduce a condition encoding approach to construct a tractable mapping mapping between observed and unobserved regions. We demonstrate the ability of the model to capture possible reconstructions and improve the accuracy of fused results.
  arXiv  Detail & Related papers  (2024-08-30T19:46:23Z)
• Improving Equivariant Model Training via Constraint Relaxation [31.507956579770088]
  We propose a novel framework for improving the optimization of such models by relaxing the hard equivariance constraint during training.<n>We provide experimental results on different state-of-the-art network architectures, demonstrating how this training framework can result in equivariant models with improved generalization performance.
  arXiv  Detail & Related papers  (2024-08-23T17:35:08Z)
• Uniform Transformation: Refining Latent Representation in Variational Autoencoders [7.4316292428754105]
  We introduce a novel adaptable three-stage Uniform Transformation (UT) module to address irregular latent distributions. By reconfiguring irregular distributions into a uniform distribution in the latent space, our approach significantly enhances the disentanglement and interpretability of latent representations. Empirical evaluations demonstrated the efficacy of our proposed UT module in improving disentanglement metrics across benchmark datasets.
  arXiv  Detail & Related papers  (2024-07-02T21:46:23Z)
• On the Trajectory Regularity of ODE-based Diffusion Sampling [79.17334230868693]
  Diffusion-based generative models use differential equations to establish a smooth connection between a complex data distribution and a tractable prior distribution. In this paper, we identify several intriguing trajectory properties in the ODE-based sampling process of diffusion models.
  arXiv  Detail & Related papers  (2024-05-18T15:59:41Z)
• Scaling and renormalization in high-dimensional regression [72.59731158970894]
  We present a unifying perspective on recent results on ridge regression.<n>We use the basic tools of random matrix theory and free probability, aimed at readers with backgrounds in physics and deep learning.<n>Our results extend and provide a unifying perspective on earlier models of scaling laws.
  arXiv  Detail & Related papers  (2024-05-01T15:59:00Z)
• Online Variational Sequential Monte Carlo [49.97673761305336]
  We build upon the variational sequential Monte Carlo (VSMC) method, which provides computationally efficient and accurate model parameter estimation and Bayesian latent-state inference. Online VSMC is capable of performing efficiently, entirely on-the-fly, both parameter estimation and particle proposal adaptation.
  arXiv  Detail & Related papers  (2023-12-19T21:45:38Z)
• Ensemble Kalman Filtering Meets Gaussian Process SSM for Non-Mean-Field and Online Inference [47.460898983429374]
  We introduce an ensemble Kalman filter (EnKF) into the non-mean-field (NMF) variational inference framework to approximate the posterior distribution of the latent states. This novel marriage between EnKF and GPSSM not only eliminates the need for extensive parameterization in learning variational distributions, but also enables an interpretable, closed-form approximation of the evidence lower bound (ELBO) We demonstrate that the resulting EnKF-aided online algorithm embodies a principled objective function by ensuring data-fitting accuracy while incorporating model regularizations to mitigate overfitting.
  arXiv  Detail & Related papers  (2023-12-10T15:22:30Z)
• Multi-Response Heteroscedastic Gaussian Process Models and Their Inference [1.52292571922932]
  We propose a novel framework for the modeling of heteroscedastic covariance functions. We employ variational inference to approximate the posterior and facilitate posterior predictive modeling. We show that our proposed framework offers a robust and versatile tool for a wide array of applications.
  arXiv  Detail & Related papers  (2023-08-29T15:06:47Z)
• Efficient Training of Energy-Based Models Using Jarzynski Equality [13.636994997309307]
  Energy-based models (EBMs) are generative models inspired by statistical physics. The computation of its gradient with respect to the model parameters requires sampling the model distribution. Here we show how results for nonequilibrium thermodynamics based on Jarzynski equality can be used to perform this computation efficiently.
  arXiv  Detail & Related papers  (2023-05-30T21:07:52Z)
• Reflected Diffusion Models [93.26107023470979]
  We present Reflected Diffusion Models, which reverse a reflected differential equation evolving on the support of the data. Our approach learns the score function through a generalized score matching loss and extends key components of standard diffusion models.
  arXiv  Detail & Related papers  (2023-04-10T17:54:38Z)
• ClusterQ: Semantic Feature Distribution Alignment for Data-Free Quantization [111.12063632743013]
  We propose a new and effective data-free quantization method termed ClusterQ. To obtain high inter-class separability of semantic features, we cluster and align the feature distribution statistics. We also incorporate the intra-class variance to solve class-wise mode collapse.
  arXiv  Detail & Related papers  (2022-04-30T06:58:56Z)
• Equivariant vector field network for many-body system modeling [65.22203086172019]
  Equivariant Vector Field Network (EVFN) is built on a novel equivariant basis and the associated scalarization and vectorization layers. We evaluate our method on predicting trajectories of simulated Newton mechanics systems with both full and partially observed data.
  arXiv  Detail & Related papers  (2021-10-26T14:26:25Z)
• A data-driven peridynamic continuum model for upscaling molecular dynamics [3.1196544696082613]
  We propose a learning framework to extract, from molecular dynamics data, an optimal Linear Peridynamic Solid model. We provide sufficient well-posedness conditions for discretized LPS models with sign-changing influence functions. This framework guarantees that the resulting model is mathematically well-posed, physically consistent, and that it generalizes well to settings that are different from the ones used during training.
  arXiv  Detail & Related papers  (2021-08-04T07:07:47Z)
• EBM-Fold: Fully-Differentiable Protein Folding Powered by Energy-based Models [53.17320541056843]
  We propose a fully-differentiable approach for protein structure optimization, guided by a data-driven generative network. Our EBM-Fold approach can efficiently produce high-quality decoys, compared against traditional Rosetta-based structure optimization routines.
  arXiv  Detail & Related papers  (2021-05-11T03:40:29Z)
• Autoregressive Score Matching [113.4502004812927]
  We propose autoregressive conditional score models (AR-CSM) where we parameterize the joint distribution in terms of the derivatives of univariable log-conditionals (scores) For AR-CSM models, this divergence between data and model distributions can be computed and optimized efficiently, requiring no expensive sampling or adversarial training. We show with extensive experimental results that it can be applied to density estimation on synthetic data, image generation, image denoising, and training latent variable models with implicit encoders.
  arXiv  Detail & Related papers  (2020-10-24T07:01:24Z)
• Learning Invariances in Neural Networks [51.20867785006147]
  We show how to parameterize a distribution over augmentations and optimize the training loss simultaneously with respect to the network parameters and augmentation parameters. We can recover the correct set and extent of invariances on image classification, regression, segmentation, and molecular property prediction from a large space of augmentations.
  arXiv  Detail & Related papers  (2020-10-22T17:18:48Z)
• Neural Control Variates [71.42768823631918]
  We show that a set of neural networks can face the challenge of finding a good approximation of the integrand. We derive a theoretically optimal, variance-minimizing loss function, and propose an alternative, composite loss for stable online training in practice. Specifically, we show that the learned light-field approximation is of sufficient quality for high-order bounces, allowing us to omit the error correction and thereby dramatically reduce the noise at the cost of negligible visible bias.
  arXiv  Detail & Related papers  (2020-06-02T11:17:55Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.