Permutation-Symmetrized Diffusion for Unconditional Molecular Generation

TL;DR

This paper introduces a novel diffusion model directly on the permutation-invariant quotient manifold for molecular generation, improving efficiency and performance over traditional permutation-equivariant methods.

Contribution

It models diffusion directly on the quotient manifold, deriving an explicit heat kernel expression and a permutation-symmetrized score approximation, advancing molecular generation techniques.

Findings

Competitive generation quality on QM9 dataset

Improved efficiency over existing methods

Practical permutation symmetrization approach

Abstract

Permutation invariance is fundamental in molecular point-cloud generation, yet most diffusion models enforce it indirectly via permutation-equivariant networks on an ordered space. We propose to model diffusion directly on the quotient manifold $\tilde{\calX}=\sR^{d\times N}/S_N$, where all atom permutations are identified. We show that the heat kernel on $\tilde{\calX}$ admits an explicit expression as a sum of Euclidean heat kernels over permutations, which clarifies how diffusion on the quotient differs from ordered-particle diffusion. Training requires a permutation-symmetrized score involving an intractable sum over $S_N$; we derive an expectation form over a posterior on permutations and approximate it using MCMC in permutation space. We evaluate on unconditional 3D molecule generation on QM9 under the EQGAT-Diff protocol, using SemlaFlow-style backbone and treating all variables continuously. The results demonstrate that quotient-based permutation symmetrization is practical and yields competitive generation quality with improved efficiency.