대칭성을 활용한 확산 모델 재고찰: 표준화 방법을 통한 분자 그래프 생성 응용

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

화학 및 과학 분야의 많은 생성 작업은 그룹 대칭(예: 순열 및 회전)에 불변인 분포를 포함합니다. 일반적인 전략은 등방성(equivariant) 디노이징 모델 및 불변(invariant) 사전 분포과 같은 아키텍처 제약을 통해 불변성과 등방성을 강제합니다. 본 논문에서는 이러한 전통적인 접근 방식을 표준화 관점이라는 대안적인 방법으로 도전합니다. 먼저 각 샘플을 표준적인 자세 또는 순서를 갖는 궤적 대표값으로 매핑하고, 제약 없는(등방성이 아닌) 확산 또는 플로우 모델을 표준화된 데이터셋에 대해 학습시킨 후, 생성 시점에 임의의 대칭 변환을 샘플링하여 불변 분포를 복원합니다. 형식적인 몫 공간 관점에서, 본 연구는 표준화된 확산의 포괄적인 이론을 제공하며, 다음과 같은 내용을 증명합니다. (i) 표준화된 생성 모델이 불변 대상을 대상으로 할 때, 정확성, 보편성 및 우수한 표현력을 갖습니다. (ii) 표준화는 그룹 혼합으로 인한 확산 스코어 복잡성을 제거하고, 플로우 매칭에서의 조건부 분산을 줄여 학습 속도를 향상시킵니다. 또한, 정렬된 사전 분포와 최적 수송이 표준화와 상호 보완적으로 작용하여 학습 효율을 더욱 향상시킬 수 있음을 보여줍니다. 본 연구는 $S_n imes SE(3)$ 대칭 하에서 분자 그래프 생성을 위한 프레임워크를 구현합니다. 기하학적 스펙트럼 기반 표준화 및 가벼운 위치 인코딩을 활용하여, 표준화된 확산은 3차원 분자 생성 작업에서 등방성 기반 모델보다 훨씬 뛰어난 성능을 보이며, 유사하거나 더 적은 계산량으로 구현됩니다. 더욱이, 새로운 아키텍처인 Canon 및 CanonFlow는 어려운 GEOM-DRUG 데이터셋에서 최첨단 성능을 달성하며, 특히 적은 단계로 생성하는 경우에도 큰 이점을 제공합니다.

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.