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

Motivation

• Known: 기존 분자 생성 모델들은 SN × SE(3) 대칭을 만족시키기 위해 equivariant denoiser와 invariant prior 등 아키텍처적 제약을 두어 score/velocity field가 군 작용을 respect하도록 설계해왔다(e.g., EDM, GeoLDM 계열).
• Gap: Equivariant 아키텍처는 원칙적이지만 tensor algebra 등으로 인한 상당한 계산·구조적 오버헤드를 유발하며, 대칭성으로 인한 latent "gauge" 모호성이 intermediate noisy state에서 다수의 동등한 group-transformed configuration으로 나타나 trajectory crossing과 상충하는 gradient를 발생시키는 문제(mismatch loss)가 충분히 다뤄지지 않았다.
• Why: Canonicalization을 통해 non-equivariant한 범용 Transformer/GNN 백본을 그대로 사용할 수 있게 되면 아키텍처 설계 부담과 계산 비용을 크게 줄이면서도 diffusion score의 복잡도를 낮추고 flow matching의 conditional variance를 줄여 훈련 효율과 few-step 생성 품질을 동시에 개선할 수 있어, 분자 생성뿐 아니라 대칭성이 관여하는 다양한 과학적 생성 과제에 폭넓게 적용 가능한 대안적 패러다임을 제공한다.
• Approach: 각 샘플을 canonical pose/order를 갖는 orbit representative로 매핑하는 canonicalization map Ψ를 정의하고, quotient-space 관점에서 canonical slice 위에 제약 없는 diffusion/flow model을 학습한 뒤 생성 시 Haar measure에서 무작위 대칭 변환을 샘플링하여 invariant 분포를 복원하는 이론적 프레임워크를 제시하고, 이를 SN × SO(3) 대칭 하의 분자 그래프 생성에 구체적으로 구현한다.

Achievement

Figure 5. OT and canonicalization act complementarily. OT can reduce conditional variances with or without canonicalizat

1. 이론적 정합성 및 표현력 증명: canonical generative model이 invariant target에 대해 correctness, universality, 그리고 equivariant model보다 우월한 expressivity를 가짐을 증명.
2. 훈련 가속 메커니즘 규명: flow-matching identity를 통해 conditional variance가 within-slice term과 symmetry-ambiguity term으로 분해됨을 보이고, canonicalization이 group mixture로 인한 diffusion score 복잡도와 symmetry-ambiguity term을 제거함을 증명.
3. OT와의 상호보완성 입증: aligned prior와 optimal transport가 canonicalization과 상호보완적으로 작용하여 within-slice term을 추가로 줄이고 훈련 효율을 더욱 향상시킴을 보임.
4. 실용적 파이프라인 및 신규 아키텍처 제안: geometric spectra 기반 canonicalization(Fiedler-vector rank, rank-anchored frame)과 mild positional encoding을 결합한 파이프라인, 그리고 canonical 정보를 활용하는 신규 아키텍처 Canon을 설계.
5. State-of-the-art 성능 달성: QM9, GEOM-DRUG 등 3D 분자 생성 벤치마크에서 canonicalized 모델(예: canonicalized SemlaFlow)이 equivariant baseline 대비 유사하거나 더 적은 계산량으로 크게 우수한 성능을 보였고, CanonFlow는 GEOM-DRUG에서 SOTA를 달성하며 few-step 생성에서도 우위를 유지함.

How

Figure 2. Overview of our canonicalized generation pipeline. (a) Canonicalization: map a molecule Z to a slice represent

• 분자 상태를 Z=(X,H,A) (좌표, 원자 특성, 결합/엣지 특성)로 정의하고, G=SO(3)×SN 작용에 대한 orbit, quotient space, canonicalization map Ψ 및 canonical slice S를 형식화.
• Haar measure λ를 이용해 invariant measure를 canonical slice 분포와 Haar randomization의 factorization으로 표현하는 정리를 증명.
• score-based diffusion(VP-SDE)과 conditional flow matching 프레임워크 하에서, irreducible regression error가 conditional variance로 표현됨을 이용해 group-aligned lifting construction으로 within-slice term과 symmetry-ambiguity term으로 분해.
• SN에 대해서는 Fiedler-vector 기반 rank로, SO(3)에 대해서는 rank-anchored frame으로 permutation과 rotation gauge를 동시에 고정하는 geometric spectra 기반 canonicalizer를 설계.
• gap-free prior와 canonical condition을 결합해 non-equivariant Transformer/GNN 백본(diffusion/flow model)을 canonical space에서 직접 학습.
• 원자 hidden state에 canonical 정보를 명시적으로 통합해 canonicality-aware denoising을 수행하는 신규 아키텍처 Canon을 설계하고 CanonFlow에 적용.
• SemlaFlow 등 기존 baseline을 canonicalized 버전으로 변형해 QM9와 GEOM-DRUG 데이터셋에서 비교 실험 수행.

Originality

• 기존의 equivariant 아키텍처 제약을 통한 대칭 처리 방식 대신, quotient-space 상의 canonicalization을 통해 대칭성을 명시적으로 깨고 non-equivariant 모델을 사용하는 대안적 패러다임을 제시.
• 대칭으로 인한 latent gauge 모호성이 diffusion score/flow의 conditional variance를 증가시킨다는 것을 conditional variance 분해(within-slice term vs symmetry-ambiguity term)로 정량화하고 이론적으로 규명.
• Optimal transport와 canonicalization이 서로 다른 term(각각 within-slice, symmetry-ambiguity)을 타겟팅하여 상호보완적으로 작동함을 처음으로 명확히 규명.
• geometric spectra(Fiedler vector) 기반의 permutation-rotation joint canonicalization과 이를 활용하는 새로운 아키텍처 Canon을 제안.

Limitation & Further Study

• Assumption 2.5(free action a.s.)에 의존하는데, 실제 분자에 대칭적 하위구조(e.g., 대칭 분자, graph automorphism)가 존재할 경우 canonical representative의 불연속성/다중값 문제(Remark 2.4)가 발생할 수 있어 이에 대한 완화책(weighted/probabilistic frame 등)의 실제 효과에 대한 심층 검증이 제한적일 수 있음.
• Fiedler-vector 기반 rank나 rank-anchored frame과 같은 특정 canonicalization 선택이 다른 분자 유형이나 더 큰 그래프, 다른 도메인(단백질, 소재 등)에도 안정적으로 일반화되는지에 대한 추가 검증이 필요함.
• 이론은 conditional variance 감소를 통한 효율성 개선을 다루지만, 실제 생성 품질(다양성, 물리적 타당성 등)에 미치는 영향에 대한 폭넓은 ablation이나 실패 사례 분석이 부족할 수 있음.
• 후속 연구로 다른 대칭군(예: 주기적 대칭을 갖는 결정 구조, 다른 물리적 불변성)에 대한 canonicalization 이론과 아키텍처의 확장이 기대됨.

Evaluation

총평: 이론과 실용성을 모두 갖춘 논문으로, equivariant 아키텍처가 사실상 필수라는 통념에 도전하며 canonicalization이라는 간결한 대안이 이론적으로도 탄탄하고 실험적으로도 SOTA 성능과 계산 효율을 동시에 달성함을 보여준 인상적인 연구이다.