[2602.15022] Rethinking Diffusion Models with Symmetries through Canonicalization with Applications to Molecular Graph Generation
Summary
This paper explores a novel approach to diffusion models by emphasizing canonicalization to enhance molecular graph generation, demonstrating superior performance over traditional methods.
Why It Matters
The research addresses a significant gap in generative modeling by challenging existing architectural constraints. By introducing canonicalization, it offers a more efficient and effective method for generating molecular graphs, which is crucial for advancements in chemistry and materials science.
Key Takeaways
- Canonicalization improves training efficiency in diffusion models.
- The proposed method outperforms equivariant baselines in 3D molecular generation tasks.
- Aligned priors and optimal transport enhance the canonicalization framework.
Computer Science > Machine Learning arXiv:2602.15022 (cs) [Submitted on 16 Feb 2026] Title:Rethinking Diffusion Models with Symmetries through Canonicalization with Applications to Molecular Graph Generation Authors:Cai Zhou, Zijie Chen, Zian Li, Jike Wang, Kaiyi Jiang, Pan Li, Rose Yu, Muhan Zhang, Stephen Bates, Tommi Jaakkola View a PDF of the paper titled Rethinking Diffusion Models with Symmetries through Canonicalization with Applications to Molecular Graph Generation, by Cai Zhou and 9 other authors View PDF HTML (experimental) Abstract:Many generative tasks in chemistry and science involve distributions invariant to group symmetries (e.g., permutation and rotation). A common strategy enforces invariance and equivariance through architectural constraints such as equivariant denoisers and invariant priors. In this paper, we challenge this tradition through the alternative canonicalization perspective: first map each sample to an orbit representative with a canonical pose or order, train an unconstrained (non-equivariant) diffusion or flow model on the canonical slice, and finally recover the invariant distribution by sampling a random symmetry transform at generation time. Building on a formal quotient-space perspective, our work provides a comprehensive theory of canonical diffusion by proving: (i) the correctness, universality and superior expressivity of canonical generative models over invariant targets; (ii) canonicalization accelerates training by removing diff...