arXiv:2502.09890cs.LG2025-02NeurIPS被引 3

提出低方差梯度估计器,提升对称性建模的训练稳定性与效率。

Rao-Blackwell Gradient Estimators for Equivariant Denoising Diffusion

  • 将数据增强视为梯度的蒙特卡洛估计,用Rao-Blackwell化降方差
  • 单样本仅需一次前向后向传播,实现更快收敛与更稳定优化
  • 适用于分子、蛋白质、晶体生成,适合对称性敏感的生成任务

在分子与蛋白质生成等领域,物理系统具有内在对称性,对建模至关重要。现有方法主要分为两类:设计等变网络架构或使用数据增强近似等变性。等变架构虽能保证对称性,但复杂度高且优化困难;数据增强灵活但难以完全捕捉对称性。本文框架通过将数据增强视为训练梯度的蒙特卡洛估计,并应用Rao-Blackwell化,显著降低训练方差,提供可证明更低方差的梯度估计器。该方法只需每样本一次前向与后向传播,实现更稳定的优化、更快的收敛与更低的方差。我们提出了名为Orbit Diffusion的实用实现,结合损失函数与采样过程。理论上保证损失函数存在等变最小值点。实验表明,Orbit Diffusion在GEOM-QM9分子构象生成上达到当前最优性能,提升晶体结构预测能力,在Perov-5与MP-20文本引导晶体生成任务中表现优异,并增强了蛋白质结构生成的设计性。代码已开源。

原文摘要 · Abstract (English)

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.

扩散模型对称性建模生成模型分子生成

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。