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

Apr 23, 2026

Quotient-Space Diffusion Models

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.

Yixian Xu, Yusong Wang, Shengjie Luo +4

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

Ask PDF

Mar 24, 2026

Permutation-Symmetrized Diffusion for Unconditional Molecular Generation

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.

Gyeonghoon Ko, Juho Lee

Korea Advanced Institute of Science and Technology

Molecular Generation Diffusion Models 3D Molecular Generation Permutation-Invariant Learning

Ask PDF

May 7, 2026

SymDrift: One-Shot Generative Modeling under Symmetries

Generative modeling of physical systems, such as molecules, requires learning distributions that are invariant under global symmetries, such as rotations in three-dimensional space. Equivariant diffusion and flow matching models can incorporate such invariances effectively, even when trained on a non-invariant empirical distribution, but they typically rely on costly multi-step sampling. Recently, drifting models have emerged as an efficient alternative, enabling single-step generation and achieving state-of-the-art performance in generative modeling tasks. However, we show that drifting models face a symmetry-specific challenge, since an equivariant generator does not generally produce the same drifting field as the one obtained from the symmetrized target distribution. Addressing this issue would require expensive symmetrization of the empirical distribution. To avoid this cost, we propose SymDrift, a framework that makes the drifting field itself symmetry-aware. We introduce two complementary strategies: (i) a symmetrized drift in coordinate space based on optimal alignment, and (ii) a

G

-invariant embedding that removes symmetry ambiguity by construction. Empirically, SymDrift outperforms existing one-shot methods on standard benchmarks for conformer and transition state generation, while remaining competitive with significantly more expensive multi-step approaches. By enabling one-shot inference, SymDrift reduces computational overhead by up to 40

\times

compared to existing baselines, making it promising for high-throughput applications such as virtual drug screening and large-scale reaction network exploration.

Samir Darouich, Vinh Tong, Lluís Pastor-Pérez +3

Institute of Artificial Intelligence Institute of Theoretical Chemistry University of Stuttgart

One-Step Generative Modeling Generative Modeling Drifting Models

Ask PDF

Abstract

Generative Modeling

Explore similar work

Quotient-Space Diffusion Models

Permutation-Symmetrized Diffusion for Unconditional Molecular Generation

SymDrift: One-Shot Generative Modeling under Symmetries