Rethinking Diffusion Models with Symmetries through Canonicalization with Applications to Molecular Graph Generation

Cai Zhou; Zijie Chen; Zian Li; Jike Wang; Kaiyi Jiang; Pan Li; Rose Yu; Muhan Zhang

cs.LGFeb 16, 2026

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,

Organizations: Massachusetts Institute of Technology · Zhejiang University · Peking University · Princeton University ·

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 diffusion score complexity induced by group mixtures and reducing conditional variance in flow matching. We then show that aligned priors and optimal transport act complementarily with canonicalization and further improves training efficiency. We instantiate the framework for molecular graph generation under $S_n \times SE(3)$ symmetries. By leveraging geometric spectra-based canonicalization and mild positional encodings, canonical diffusion significantly outperforms equivariant baselines in 3D molecule generation tasks, with similar or even less computation. Moreover, with a novel architecture Canon, CanonFlow achieves state-of-the-art performance on the challenging GEOM-DRUG dataset, and the advantage remains large in few-step generation.

Figures & tables

Figure 1: Motivation of Canonical Diffusion: Position-wise supervision induces mismatch loss under symmetries for equivalent molecules in diffusion models.

Figure 2: Overview of our canonicalized generation pipeline. (a) Canonicalization: map a molecule $Z$ to a slice representative $\tilde{Z}=\Psi_{\phi}(Z)$ under $\mathcal{G}=SO(3)\times S_{N}$ , inducing $q_{0}=(\Psi_{\phi})_{\#}p_{0}$ . (b) Training: learn a diffusion/flow model on $\tilde{Z}_{0}\sim q_{0}$ with slice prior $q_{1}$ (optionally using an OT coupling for aligned pairings). (c) Sampling: generate on the slice (optionally with projected canonical sampling) and then apply $g\sim\mathrm{Haar}(\mathcal{G})$ to obtain an invariant sample $Z=g\to\hat{Z}_{0}$ .

Figure 3: Training trajectories of our canonicalized model (visualization of the learned transport dynamics).

Methods	Atom Stab $\uparrow$	Mol Stab $\uparrow$	Valid $\uparrow$	Opt-RMSD $\downarrow$	NFE $\downarrow$
EDM	98.7	82.0	91.9	–	1000
GCDM	98.7	85.7	94.8	–	1000
MUDiff	98.3	89.9	95.3	–	1000
FlowMol	99.7	96.2	97.3	–	100
MiDi	99.8	97.5	97.9	–	500
EQGAT-diff	99.9 ${}_{\pm\text{0.0}}$	98.7 ${}_{\pm\text{0.18}}$	99.0 ${}_{\pm\text{0.16}}$	–	500

Table 1: Performance of canonicalized SemlaFlow on QM9 dataset. Atom stability, molecule stability, and validity are reported as percentages; Opt-RMSD is reported in Å.

Methods	Atom Stab $\uparrow$	Mol Stab $\uparrow$	Valid $\uparrow$	Unique $\uparrow$	Novel $\uparrow$	NFE $\downarrow$
FlowMol	99.0	67.5	51.2	-	-	100
MiDi	99.8	91.6	77.8	100.0	100.0	500
EQGAT-diff	99.8 ${}_{\pm\text{0.0}}$	93.4 ${}_{\pm\text{0.21}}$	94.6 ${}_{\pm\text{0.24}}$	100.0 ${}_{\pm\text{0.0}}$	99.9 ${}_{\pm\text{0.07}}$	500
SemlaFlow w/ OT	99.8 ${}_{\pm\text{0.0}}$	97.3 ${}_{\pm\text{0.08}}$	93.9 ${}_{\pm\text{0.19}}$	100.0 ${}_{\pm\text{0.0}}$	99.6 ${}_{\pm\text{0.03}}$	100
Canon. SemlaFlow w/ OT anneal (ours)	99.8 ${}_{\pm\text{0.0}}$	98.1 ${}_{\pm\text{0.03}}$	95.0 ${}_{\pm\text{0.20}}$	100.0 ${}_{\pm\text{0.0}}$	99.7 ${}_{\pm\text{0.02}}$	100
CanonFlow w/ OT anneal (ours)	99.9 ${}_{\pm\text{0.0}}$	98.4 ${}_{\pm\text{0.02}}$	95.9 ${}_{\pm\text{0.08}}$	100.0 ${}_{\pm\text{0.0}}$	99.7 ${}_{\pm\text{0.01}}$	100

Table 2: Performance of canonicalized SemlaFlow and CanonFlow on GEOM-DRUG dataset. Atom stability, molecule stability, validity, as well as uniqueness and novelty are reported as percentages.

Model	Mol Stab $\uparrow$	Valid $\uparrow$	Sample Time $\downarrow$	NFE $\downarrow$
EQGAT-diff	93.4 ${}_{\pm\text{0.21}}$	94.6 ${}_{\pm\text{0.24}}$	2293.0	500
SemlaFlow 50 w/ OT	97.0 ${}_{\pm\text{0.21}}$	93.9 ${}_{\pm\text{0.12}}$	49.8	50
Canon. SemlaFlow 50	97.5 ${}_{\pm\text{0.14}}$	94.9 ${}_{\pm\text{0.07}}$	50.1	50
SemlaFlow 100 w/ OT	97.3 ${}_{\pm\text{0.08}}$	93.9 ${}_{\pm\text{0.19}}$	99.3	100
Canon. SemlaFlow 100	97.9 ${}_{\pm\text{0.09}}$	94.9 ${}_{\pm\text{0.18}}$	99.7	100
Data	100.0	100.0	–	–

Table 3: Comparison of few-step generation on GEOM-DRUG (w/o CFG in training and sampling). Sample time is measured by the average number of seconds to generate 1000 molecules. Canonical SemlaFlow significantly outperforms baselines w/ fewer sampling steps, with negligible increased sample time.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: OT and canonicalization act complementarily. OT can reduce conditional variances with or without canonicalization, but OT solutions are generally non-unique in the presence of symmetry (if

\gamma

is optimal then

(g,g)_{\#}\gamma

is also optimal); averaging yields a diagonal-invariant optimum but may destroy Monge-ness. Canonicalization fixes a gauge (collapsing each orbit to a unique slice representative), eliminating group ambiguity and stabilizing OT on the slice.

a=\operatorname{argmax}_{k\in[\lfloor N/3\rfloor,\lfloor 2N/3\rfloor)}\|(\mathbf{x}_{k}-\mathbf{x}_{h})\times(\mathbf{x}_{t}-\mathbf{x}_{h})\|

Appendix

Algorithm 2 Spectral

SO(3)

Canonicalization

Panel of Figure 5 — Figure 5: Spectral canonicalization via the (signed) Fiedler vector produces stable, locality-preserving atom orderings for both highly symmetric and weakly symmetric molecules (blue: start, red: end), demonstrating robust performance across diverse symmetry regimes.

Figure 6: Overview of the Canon architecture. Three parallel streams—coordinates (

\mathbf{CS}

), node features (

\mathcal{H}

), and canonical rank (

\mathbf{R}

)—are updated through

L

message-passing layers consisting of feed-forward blocks, pairwise edge messages, and attention-based aggregation. Gray blocks denote components inherited from Semla; blue blocks denote novel canonical-rank components introduced by Canon.

Model	Mol Stab $\uparrow$	Valid $\uparrow$	NFE
EQGAT-diff	93.4 ${}_{\pm\text{0.21}}$	94.6 ${}_{\pm\text{0.24}}$	500
SemlaFlow 50 w/ OT	97.0 ${}_{\pm\text{0.21}}$	93.9 ${}_{\pm\text{0.12}}$	50
CanonFlow 50	98.0 ${}_{\pm\text{0.08}}$	95.4 ${}_{\pm\text{0.09}}$	50
SemlaFlow 100 w/ OT	97.3 ${}_{\pm\text{0.08}}$	93.9 ${}_{\pm\text{0.19}}$	100
CanonFlow 100	98.4 ${}_{\pm\text{0.02}}$	95.9 ${}_{\pm\text{0.08}}$	100
Data	100.0	100.0	–

Appendix

Table 4: CanonFlow w/ our novel Canon architecture on GEOM-DRUG. Canonicalization is conducted on

S_{N}

Model	CFG scale	Mol Stab $\uparrow$	Valid $\uparrow$	NFE
EQGAT-diff	–	93.4 ${}_{\pm\text{0.21}}$	94.6 ${}_{\pm\text{0.24}}$	500
SemlaFlow 50	–	97.0 ${}_{\pm\text{0.21}}$	93.9 ${}_{\pm\text{0.12}}$	50
Canon. SemlaFlow 50	1.0	97.6 ${}_{\pm\text{0.13}}$	94.6 ${}_{\pm\text{0.32}}$	50
Canon. SemlaFlow 50	1.5	97.4 ${}_{\pm\text{0.07}}$	94.7 ${}_{\pm\text{0.12}}$	50
Canon. SemlaFlow 50	2.0	96.7 ${}_{\pm\text{0.04}}$	94.5 ${}_{\pm\text{0.03}}$	50
SemlaFlow 100	–	97.3 ${}_{\pm\text{0.08}}$	93.9 ${}_{\pm\text{0.19}}$	100

Appendix

Table 5: Canonical SemlaFlow on GEOM-DRUG dataset w/ classifier-free guidance (CFG). Canonicalization is conducted on

S_{N}

Model	CFG scale	Mol Stab $\uparrow$	Valid $\uparrow$	NFE
EQGAT-diff	–	93.4 ${}_{\pm\text{0.21}}$	94.6 ${}_{\pm\text{0.24}}$	500
SemlaFlow 100	–	97.3 ${}_{\pm\text{0.08}}$	93.9 ${}_{\pm\text{0.19}}$	100
Canon. SemlaFlow 100	1.0	97.9 ${}_{\pm\text{0.09}}$	93.8 ${}_{\pm\text{0.30}}$	100
Canon. SemlaFlow 100	1.5	97.7 ${}_{\pm\text{0.06}}$	94.2 ${}_{\pm\text{0.23}}$	100
Canon. SemlaFlow 100	2.0	97.6 ${}_{\pm\text{0.09}}$	94.4 ${}_{\pm\text{0.05}}$	100
Data	–	100.0	100.0	–

Appendix

Table 6: Canonical SemlaFlow on GEOM-DRUG dataset w/ classifier-free guidance (CFG). Canonicalization is conducted on

S_{N}\times SO(3)

Group	Model	PE	OT	Mol Stab $\uparrow$	Valid $\uparrow$	NFE
	EQGAT-diff	-	-	93.4 ${}_{\pm\text{0.21}}$	94.6 ${}_{\pm\text{0.24}}$	500
	SemlaFlow 20	-	equivariant	95.3 ${}_{\pm\text{0.14}}$	93.0 ${}_{\pm\text{0.10}}$	20
	SemlaFlow 50	-	equivariant	97.0 ${}_{\pm\text{0.21}}$	93.9 ${}_{\pm\text{0.12}}$	50
	SemlaFlow 100	-	equivariant	97.3 ${}_{\pm\text{0.08}}$	93.9 ${}_{\pm\text{0.19}}$	100
	Canon. SemlaFlow 20	-	equivariant	95.1 ${}_{\pm\text{0.09}}$	93.5 ${}_{\pm\text{0.17}}$	20
$\quad S_{N}$	Canon. SemlaFlow 50	-	equivariant	97.3 ${}_{\pm\text{0.08}}$	94.4 ${}_{\pm\text{0.02}}$	50

Appendix

Table 7: Canonical SemlaFlow with aligned prior on GEOM-DRUG. Canonicalization is conducted on

S_{N}

, and we study both equivariant and non-equivariant backbones w/ or w/o OT. Note that the prior is computed consistently with different group

G

Figure 7: Random samples from Canonicalized Semlaflow model trained on GEOM Drugs. These samples were generated using 100 ODE integration steps.

Generative Modeling

Generative modeling learns data distributions to create new samples, from images and text to molecules, simulations, and structured records.

See all 493 papers

Quotient-Space Diffusion Models
Apr 23, 2026Yixian Xu, Yusong Wang, Shengjie Luo +4Protein Structure Generation Generative Modeling
State Key Laboratory of General Artificial Intelligence, Peking University, Beijing, China · State Key Laboratory of Human-Machine Hybrid Augmented Intelligence, Institute of Artificial Intelligence and Robotics, Xi’an Jiaotong University · Zhongguancun Academy, Beijing, China +2
Permutation-Symmetrized Diffusion for Unconditional Molecular Generation
Mar 24, 2026Gyeonghoon Ko, Juho LeeMolecular Generation Diffusion Models
Korea Advanced Institute of Science and Technology
SymDrift: One-Shot Generative Modeling under Symmetries
May 7, 2026Samir Darouich, Vinh Tong, Lluís Pastor-Pérez +3One-Step Generative Modeling Generative Modeling
Institute of Artificial Intelligence Institute of Theoretical Chemistry University of Stuttgart

Apr 23, 2026J/K move · Enter open · S save

Quotient-Space Diffusion Models

Yixian Xu, Yusong Wang, Shengjie Luo, Kaiyuan Gao+3

State Key Laboratory of General Artificial Intelligence, Peking University, Beijing, China · State Key Laboratory of Human-Machine Hybrid Augmented Intelligence, Institute of Artificial Intelligence and Robotics, Xi’an Jiaotong University · Zhongguancun Academy, Beijing, China+2

Protein Structure Generation Generative Modeling 3D Molecular Generation Molecular Generation Diffusion Models

Open paper AskPDF arXiv

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.

Abstract

Generative Modeling

Explore similar work

Quotient-Space Diffusion Models

Permutation-Symmetrized Diffusion for Unconditional Molecular Generation

SymDrift: One-Shot Generative Modeling under Symmetries