Quotient-Space Diffusion Models

Yixian Xu, Yusong Wang, Shengjie Luo, Kaiyuan Gao, Tianyu He, Di He, Chang Liu

Keywords: Diffusion Models, Generative Modeling, Geometric Deep Learning, Molecular Structure Generation

TL;DR: We propose a principled way to leverage group symmetry of the target distribution by defining a diffusion model on the quotient space, which achieves both easier learning and correct sampling for the first time.

Abstract:

Diffusion-based generative models have reformed generative AI, and also enabled new capabilities in the science domain, _e.g._, fast generation of 3D structures of molecules. In such tasks, there is often a _symmetry_ in the system, identifying elements that can be converted by certain transformations as equivalent. Equivariant diffusion models guarantee a symmetric distribution, but miss the opportunity to make learning easier, while alignment-based simplification attempts fail to preserve the target distribution.
In this work, we develop _quotient-space diffusion models_, a principled generative framework to fully handle and leverage symmetry.
By viewing the intrinsic generation process on the quotient space, the exact construction that removes symmetry redundancy, the framework simplifies learning by allowing model output to have an arbitrary intra-equivalence-class movement, while generating the correct symmetric target distribution with guarantee.
We instantiate the framework for molecular structure generation which follows $\\mathrm{SE}(3)$ (rigid-body movement) symmetry.
It improves the performance over equivariant diffusion models and outperforms alignment-based methods universally for small molecules and proteins, representing a new framework that surpasses previous symmetry treatments in generative models.

Primary Area: generative models

Submission Number: 7627