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.
翻译:化学与科学领域的许多生成任务涉及对群对称性(如置换与旋转)不变的分布。传统方法通常通过架构约束(如等变去噪器与不变先验)来强制不变性与等变性。本文中,我们通过替代性的规范化视角挑战这一传统:首先将每个样本映射至具有规范姿态或顺序的轨道代表元,在规范切片上训练无约束(非等变)的扩散或流模型,最终在生成时通过随机对称变换恢复不变分布。基于形式化的商空间视角,本研究通过证明以下两点建立了规范化扩散的完备理论:(i)规范化生成模型对不变目标的正确性、普适性及更优的表达能力;(ii)规范化通过消除群混合引入的扩散得分复杂性并降低流匹配中的条件方差,从而加速训练。我们进一步证明对齐先验与最优传输可与规范化互补并进一步提升训练效率。我们在 $S_n \times SE(3)$ 对称性下的分子图生成任务中实例化了该框架。通过利用基于几何谱的规范化与温和的位置编码,规范化扩散在三维分子生成任务中显著优于等变基线模型,且计算量相当甚至更少。此外,借助新颖的 Canon 架构,CanonFlow 在极具挑战性的 GEOM-DRUG 数据集上取得了最先进的性能,且在少步生成中仍保持显著优势。