대칭을 활용한 확산 모델 재구성: 분자 그래프 생성의 새로운 패러다임

초록

본 논문은 대칭 불변 분포를 다루는 기존의 equivariant 네트워크 방식 대신, 샘플을 정규화된 대표 형태(캐노니컬 포즈)로 변환한 뒤 비제한(non‑equivariant) 확산·플로우 모델을 학습하고, 생성 시 무작위 대칭 변환을 적용하는 “캐노니컬화” 접근법을 제안한다. 이론적으로 캐노니컬 모델이 불변 목표를 정확히 재현하고, 표현력·학습 효율이 우수함을 증명했으며, SE(3)·Sₙ 대칭을 갖는 3D 분자 그래프 생성 실험에서 기존 equivariant 베이스라인을 크게 능가한다. 특히 제안된 CanonFlow는 GEOM‑DRUG 데이터셋에서 최첨단 성능을 달성하고, 적은 샘플 단계에서도 강인한 결과를 보여준다.

상세 분석

이 논문은 대칭이 내재된 과학·화학 데이터(예: 분자 구조)의 생성 모델링에서 “대칭을 강제하는” 전통적 접근법을 근본적으로 재고한다. 기존 방법은 모델 자체를 equivariant 혹은 invariant하도록 설계해, 입력에 대한 회전·순열 변환에 대해 동일한 출력을 보장한다. 그러나 이러한 구조적 제약은 학습 시에 복잡한 그룹 혼합(score mixture) 문제를 야기하고, 조건부 분산의 분산을 증가시켜 수렴 속도를 저하시킨다. 저자들은 이를 “캐노니컬화”라는 관점으로 전환한다. 구체적으로, 각 데이터 샘플을 해당 대칭 궤도(orbit) 내에서 사전 정의된 기준 형태(예: 스펙트럼 기반 정렬)로 매핑하고, 이 “캐노니컬 슬라이스”에 대해 비제한 확산·플로우 모델을 학습한다. 생성 단계에서는 학습된 모델로부터 캐노니컬 샘플을 얻은 뒤, 무작위 대칭 변환을 적용해 원래의 불변 분포를 복원한다.

이 접근법의 핵심 이론적 기여는 세 가지이다. (i) 정확성: 캐노니컬 모델이 원래의 불변 목표 분포와 동일한 확률 측정을 제공함을 증명한다. (ii) 보편성·표현력: 캐노니컬 모델은 모든 equivariant 모델을 포함하는 상위 클래스이며, 특히 복잡한 대칭 그룹에 대해 더 풍부한 함수 공간을 탐색한다. (iii) 학습 효율: 캐노니컬화는 그룹 혼합에 의해 발생하는 스코어 복잡성을 제거하고, 플로우 매칭에서 조건부 분산을 감소시켜 학습 속도를 가속한다.

또한, 저자는 **정렬된 사전(prior)과 최적 수송(Optimal Transport)**을 결합해 캐노니컬화와 상호 보완적인 효과를 얻는다. 정렬된 사전은 캐노니컬 공간에서의 초기 분포를 보다 간단하게 만들고, 최적 수송은 학습 과정에서의 변환 경로를 최소화해 효율성을 높인다.

실험적으로는 $S_n \times SE(3)$ 대칭을 갖는 3차원 분자 그래프 생성에 초점을 맞춘다. 스펙트럼 기반 캐노니컬화와 가벼운 위치 인코딩을 사용해, 기존 equivariant denoiser 기반 확산 모델보다 높은 품질의 분자를 더 적은 연산량으로 생성한다. 특히, 제안된 CanonFlow 아키텍처는 GEOM‑DRUG 데이터셋에서 RMSD, Covariance, 그리고 유효성 평가 지표 모두에서 현존 최고 기록을 세우며, 특히 5~10 단계와 같은 적은 샘플 단계에서도 큰 성능 격차를 유지한다. 이는 캐노니컬화가 고차원 대칭 공간에서의 샘플링 효율을 크게 개선함을 실증한다.

전반적으로 이 논문은 “대칭을 강제하는” 대신 “대칭을 정규화하고 나중에 재삽입하는” 전략이 이론적 타당성과 실험적 우수성을 동시에 제공한다는 점에서, 향후 과학·화학 생성 모델링뿐 아니라 물리·그래프 기반 생성 전반에 새로운 설계 패러다임을 제시한다.