Rethinking Diffusion Models with Symmetries through Canonicalization for Molecular Graph Generation

Paper: 2602.15022 Authors: Cai Zhou, Zijie Chen, Zian Li, Jike Wang, Kaiyi Jiang, Pan Li, Rose Yu, Muhan Zhang, Stephen Bates, Tommi Jaakkola Categories: cs.LG, cs.AI, math.GR, q-bio.BM

Abstract

This paper introduces a paradigm shift in handling symmetries for generative models in scientific domains. Rather than enforcing equivariance through architectural constraints, the authors propose canonicalization: mapping samples to canonical representatives, training unconstrained models on this canonical slice, and recovering the full invariant distribution by applying random symmetry transforms during generation. The approach is grounded in quotient-space theory and demonstrates both theoretical superiority and practical advantages. Applied to molecular graph generation under $S_n \times SE(3)$ symmetries, canonical diffusion significantly outperforms equivariant baselines while using similar or less computation, achieving state-of-the-art results on the GEOM-DRUG dataset.

Key Contributions

• Introduction of the canonicalization framework as an alternative to equivariant architectures for symmetry-invariant generative modeling
• Comprehensive theoretical analysis proving correctness, universality, and superior expressivity of canonical models over invariant targets
• Demonstration that canonicalization accelerates training by removing score complexity from group mixtures and reducing conditional variance in flow matching
• Novel Canon architecture combining geometric spectra-based canonicalization with mild positional encodings
• State-of-the-art performance on 3D molecular generation, particularly on the challenging GEOM-DRUG dataset with significant advantages in few-step generation

Theoretical Framework and Methodology

The core insight of this work lies in decomposing the symmetry-invariant generation problem through the lens of quotient spaces. Given a distribution $p(x)$ invariant under group $G$, traditional approaches constrain the model architecture to respect $G$-equivariance. The canonical approach instead:

1. Defines a canonicalization function $\phi: \mathcal{X} \rightarrow \mathcal{X}/G$ that maps each sample to a unique representative in its orbit
2. Trains an unconstrained diffusion or flow model $q_\theta$ on the canonical slice $\mathcal{X}/G$
3. Recovers the full distribution via $p_\theta(x) = \int_G q_\theta(\phi(g \cdot x)) dg$

The authors prove that this approach is not only correct but also more expressive than equivariant models. Specifically, they show that canonical models can represent any invariant distribution, while equivariant architectures may be fundamentally limited in their expressivity. The key theoretical results establish:

• The canonical slice $\mathcal{X}/G$ has lower intrinsic dimensionality than the full space, making it easier to model
• Training on canonical data removes the “group mixture” complexity that equivariant models must handle in their score functions
• For flow matching, canonicalization reduces the conditional variance of the velocity field, leading to more stable training

The framework is instantiated for molecular graphs with $S_n$ (permutation) and $SE(3)$ (rotation and translation) symmetries. The canonicalization uses geometric spectra—eigenvalues and eigenvectors of graph Laplacians and distance matrices—to establish canonical orderings and orientations that are provably unique up to discrete ambiguities.

Experimental Results and Performance

The empirical validation focuses on 3D molecular generation tasks, comparing canonical diffusion against strong equivariant baselines including EDM, GeoLDM, and other state-of-the-art methods. Key findings include:

On the QM9 dataset, canonical diffusion achieves comparable or superior performance to equivariant models across all molecular property metrics (validity, uniqueness, novelty) while using similar computational resources. The model demonstrates particular strength in generating chemically valid structures with correct bond lengths and angles.

On the more challenging GEOM-DRUG dataset containing larger, more complex drug-like molecules, the Canon architecture with CanonFlow achieves state-of-the-art results. The performance gap widens significantly in few-step generation scenarios, where canonical models maintain high quality while equivariant baselines degrade. This suggests that the simplified score landscape from canonicalization enables more efficient sampling trajectories.

Ablation studies reveal that the benefits of canonicalization are complementary to other techniques like aligned priors and optimal transport. Combining these approaches yields the best results, with each component contributing measurably to final performance.

Computational analysis shows that despite the overhead of canonicalization operations, the overall training time is competitive or even faster than equivariant models due to the simplified learning problem on the canonical slice. The inference cost is comparable, as the canonicalization and symmetry recovery steps are relatively lightweight compared to the denoising network evaluations.

Implications and Future Directions

This work challenges the prevailing wisdom that symmetries must be baked into model architectures. The canonicalization perspective offers several advantages: simpler architectures without equivariance constraints, better theoretical properties, and improved training dynamics. The success on molecular generation suggests broad applicability to other scientific domains with symmetries.

The theoretical framework opens questions about the fundamental relationship between symmetry, expressivity, and learning efficiency in generative models. The finding that removing symmetry from the learning problem (via canonicalization) can improve results contradicts intuitions about inductive biases, suggesting that the complexity cost of equivariant architectures may outweigh their benefits in some settings.

For molecular design and drug discovery, the practical improvements are significant. Faster generation with fewer steps enables more efficient exploration of chemical space, while the improved quality of generated molecules increases the likelihood of discovering viable candidates. The framework’s flexibility also allows easy integration with other advances in diffusion models and flow matching.

Future work could explore canonicalization for other symmetry groups, investigate learned canonicalization functions that adapt to data, and extend the approach to conditional generation tasks where symmetries interact with conditioning information.

Takeaways

1. Canonicalization provides a theoretically grounded and practically effective alternative to equivariant architectures for symmetry-invariant generation
2. Training on canonical representations simplifies the learning problem by removing group mixture complexity and reducing score function variance
3. The approach achieves state-of-the-art results on molecular graph generation, particularly excelling in few-step generation scenarios
4. Canonical models are provably more expressive than equivariant models for invariant distributions, challenging conventional wisdom about symmetry inductive biases
5. The framework is complementary to other advances like aligned priors and optimal transport, enabling further performance improvements through combination