Symmetry-Aware Generative Modeling through Learned Canonicalization

TL;DR

The paper addresses the challenge of modeling symmetric densities by arguing that learning the full invariant distribution $p_X$ is unnecessary. It introduces a learned canonicalization function $h$ that maps inputs to a canonical pose, defining the orbit representative map $c(x) = h(x)^{-1}x$ to model the distribution over orbits $p_{X/G}$ with a non-equivariant diffusion generator. Empirical results on RotMNIST and QM9 show that canon+GDM yields higher-quality samples and significantly faster inference than fully equivariant baselines, while a frozen-canonicalizer baseline still produces stable molecules. This symmetry-aware, architecture-agnostic approach preserves inductive biases while enabling flexible, efficient generative modeling in symmetry-rich domains, with broad potential applications in AI for science.

Abstract

Generative modeling of symmetric densities has a range of applications in AI for science, from drug discovery to physics simulations. The existing generative modeling paradigm for invariant densities combines an invariant prior with an equivariant generative process. However, we observe that this technique is not necessary and has several drawbacks resulting from the limitations of equivariant networks. Instead, we propose to model a learned slice of the density so that only one representative element per orbit is learned. To accomplish this, we learn a group-equivariant canonicalization network that maps training samples to a canonical pose and train a non-equivariant generative model over these canonicalized samples. We implement this idea in the context of diffusion models. Our preliminary experimental results on molecular modeling are promising, demonstrating improved sample quality and faster inference time.

Figures (2)

• Figure 1: Left: A distribution $p_X$ that's invariant to the group of rotations $G$. Right: a) A non-equivariant generative model must learn the whole distribution $p_X$. b) A fixed canonicalization method outputs group elements that map different samples of $p_X$ onto a single slice (red), which may be learned by a generative model. c) A learned canonicalizer can map samples onto a slice that yields a simpler distribution for the generative model to learn.
• Figure 2: a) Output of the learned canonicalizer on RotMNIST after training showing aligned digits. b) Select molecules generated from canon+GDM in a near-canonical pose.