Quotient-Space Diffusion Models

TL;DR

This paper introduces quotient-space diffusion models that leverage symmetry in generative tasks, simplifying learning and improving performance in molecular structure generation over previous equivariant and alignment-based methods.

Contribution

The work develops a novel quotient-space diffusion framework that fully exploits symmetry, outperforming existing equivariant and alignment-based approaches in generative modeling.

Findings

Improves molecular structure generation performance.

Outperforms equivariant diffusion models.

Universal superiority over alignment-based methods.

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.