Rethinking Diffusion Models with Symmetries through Canonicalization with Applications to Molecular Graph Generation
이 논문은 분자 그래프 생성과 같은 과학적 생성 작업에서 대칭성을 처리하기 위해 기존에 널리 사용되던 등변성 (equivariance) 제약 대신, 샘플을 표준형 (canonical) 으로 매핑하여 제약 없는 모델을 학습한 후 생성 시 대칭 변환을 적용하는 '정규화 (canonicalization)' 접근법을 제안하고, 이를 통해 GEOM-DRUG 데이터셋에서 최첨단 성능을 달성함과 동시에 학습 효율성을 크게 향상시켰음을 보여줍니다.
원저자:Cai Zhou, Zijie Chen, Zian Li, Jike Wang, Kaiyi Jiang, Pan Li, Rose Yu, Muhan Zhang, Stephen Bates, Tommi Jaakkola
분자 생성 AI 를 공부할 때 가장 큰 난관은 **'대칭성 (Symmetry)'**입니다.
비유: imagine you are trying to teach a child to draw a specific chair.
만약 아이가 의자를 앞에서 그렸든, 옆에서 그렸든, 거꾸로 그렸든, 결국 그 의자는 같은 의자입니다.
하지만 AI 에게는 이 세 가지 그림이 완전히 다른 그림으로 보입니다.
기존 AI 는 "어떤 방향이 정답일까?"라고 고민하며, 모든 방향을 다 고려해야 하므로 미로 속에서 헤매는 것처럼 학습이 느리고 헷갈립니다. (이걸 '대칭성 혼란'이라고 합니다.)
기존 연구들은 AI 의 구조 자체를 아주 복잡하게 만들어서 "어떤 방향이든 똑같이 인식하게" 만들었습니다. 하지만 이는 마치 미로를 헤매지 않게 하기 위해 미로 전체를 벽으로 막아버리는 것과 같아, AI 가 배울 수 있는 것이 줄어들고 계산도 매우 무거워집니다.
2. 해결책: "표준화 (Canonicalization)"라는 나침반
이 논문은 **"대칭성을 깨뜨리는 것"**이 오히려 더 빠르고 똑똑한 길이라고 말합니다.
핵심 아이디어: AI 가 학습할 때, 모든 분자를 **하나의 '표준 자세 (Canonical Pose)'**로 맞춰버리는 것입니다.
비유:
의자 그림을 가르칠 때, AI 에게 "의자는 무조건 정면에서, 다리가 바닥에 닿은 상태로만 그려라"라고 규칙을 정해버리는 것입니다.
AI 는 이제 "어떤 방향일까?"라고 고민할 필요가 없습니다. 오직 하나뿐인 정답만 학습하면 되므로, 미로가 사라지고 직선으로 빠르게 달릴 수 있게 됩니다.
학습이 끝난 후, 실제로 분자를 만들어낼 때는 AI 가 만든 '표준 의자'를 임의의 방향으로 돌려서 (랜덤하게 회전시켜서) 내보냅니다. 이렇게 하면 AI 는 여전히 모든 방향의 분자를 만들 수 있지만, 학습 과정은 훨씬 쉬워진 것입니다.
3. 왜 이것이 더 좋은가요? (세 가지 장점)
학습 속도가 빨라집니다 (직선 항해):
기존 방식은 AI 가 여러 개의 정답 (회전된 버전) 사이를 오가며 헷갈려서 진로를 자주 바꿨습니다.
이 방식은 **하나의 정해진 길 (표준화 된 공간)**만 따라가게 하므로, AI 가 훨씬 더 직관적이고 빠르게 학습합니다. 마치 미로를 헤매지 않고 직선 도로를 달리는 것과 같습니다.
더 복잡한 분자를 만들 수 있습니다 (표현력):
기존 방식은 AI 가 "어떤 방향이든 똑같이 인식해야 한다"는 복잡한 규칙을 따르느라, 뇌 (네트워크) 를 제한했습니다.
이 방식은 규칙을 깨뜨리고 자유로운 AI를 쓰게 하므로, 훨씬 더 정교하고 복잡한 분자 구조를 상상해 낼 수 있습니다.
적은 노력으로 더 좋은 결과 (Few-step Generation):
보통 AI 가 분자를 만들려면 많은 단계 (예: 100 번의 계산) 가 필요합니다.
이 방식은 학습이 깔끔하므로 단계가 훨씬 적어도 (예: 50 번) 매우 정확한 분자를 만들어냅니다. 마치 숙달된 요리사가 레시피를 줄여도 맛있는 요리를 만드는 것과 같습니다.
4. 실제 성과: "CanonFlow"라는 새로운 요리사
저자들은 이 아이디어를 적용하여 CanonFlow라는 새로운 AI 모델을 만들었습니다.
실험 결과: 유명한 약물 데이터셋 (GEOM-DRUG) 에서 기존에 가장 잘하던 모델들보다 더 안정적이고, 더 유효한 분자를 만들어냈습니다.
특이점: 계산 비용은 비슷하거나 오히려 더 적게 들면서, 학습 시간과 생성 속도 면에서 압도적인 우위를 보였습니다.
요약
이 논문은 **"분자 만들기를 가르칠 때, 모든 방향을 고려하게 하느라 헤매게 하지 말고, 일단 하나의 '표준 자세'로 맞춰서 가르친 뒤, 나중에 자유롭게 돌려주자"**는 발상의 전환을 제시합니다.
이는 미로에서 헤매는 아이를 구출하여 직선 도로로 안내하는 것과 같아, AI 가 분자 세계를 더 빠르고 정확하게 탐험할 수 있게 해줍니다.
1. 문제 정의 (Problem)
대칭성의 이중성: 화학 및 과학 데이터 (특히 3D 분자) 는 원자 순서 (SN) 와 회전/이동 ($SE(3)$) 에 대해 불변 (Invariant) 입니다. 즉, 분자의 인덱스를 바꾸거나 회전시켜도 물리적으로 동일한 분자입니다.
기존 접근법의 한계:
기존 모델은 대칭성을 유지하기 위해 **공변성 (Equivariance)**을 갖는 아키텍처 (Equivariant Layers) 나 불변 Prior 를 사용했습니다.
그러나 이는 계산 비용이 크고, 아키텍처의 표현력 (Expressivity) 을 제한합니다.
더 중요한 문제는 잠재적 '게이지 (Gauge)' 모호성입니다. 노이즈 상태 (noisy state) 에서 대칭적으로 변환된 여러 구성이 모두 유효할 수 있어, 모델이 올바른 방향을 결정하기 어렵습니다. 이로 인해 학습 중 '궤적 교차 (trajectory crossing)'와 상충되는 기울기가 발생하여 학습이 비효율적이고, 소수 단계 (few-step) 생성 시 성능이 저하됩니다.
이미지 생성과의 차이: 이미지 생성은 회전이나 반전에 대해 완전히 불변하지 않아 (위아래가 뒤집힌 풍경은 의미가 없음) 고정된 그리드나 위치 인코딩 (PE) 을 사용하여 대칭성을 깨뜨립니다. 반면 분자 생성은 대칭성이 물리적으로 필수적이므로 이를 어떻게 다룰지 고민해야 합니다.
2. 방법론 (Methodology)
저자들은 정규화 (Canonicalization) 관점을 도입하여 대칭성을 학습 과정에서 명시적으로 깨뜨리는 새로운 프레임워크를 제안합니다.
핵심 아이디어: 궤적 대표자 (Orbit Representative) 매핑
정규화 (Canonicalization): 각 분자 샘플을 군 (Group) 궤도 (Orbit) 의 유일한 대표자 (Canonical Pose/Order) 로 매핑합니다.
순열 (SN) 정규화: 그래프 라플라시안의 Fiedler 벡터 (두 번째 고유벡터) 를 사용하여 원자의 고유한 순서를 결정합니다. 이는 분자의 기하학적 구조를 반영한 자연스러운 순서화입니다.
회전 ($SO(3)$) 정규화: 정규화된 순서를 기반으로 고정된 기준 좌표계 (Frame) 를 설정하여 회전 불변성을 깨뜨립니다.
학습 (Training):
정규화된 공간 (Canonical Slice) 에서 **비공변적 (Non-equivariant)**인 확산 또는 플로우 매칭 모델을 학습합니다.
이 과정에서 군 혼합 (Group Mixture) 으로 인한 복잡성이 제거되어, 학습할 스코어 (Score) 장이 매끄러워지고 조건부 분산 (Conditional Variance) 이 감소합니다.
정렬된 Prior 와 최적 수송 (OT): 정규화된 공간 내에서 데이터와 Prior 가 잘 정렬되도록 하고, 최적 수송 (Optimal Transport) 을 사용하여 궤적을 직선화합니다.
샘플링 및 복원 (Sampling & Restoration):
학습된 모델을 통해 정규화된 공간에서 샘플을 생성합니다.
생성된 샘플에 무작위 대칭 변환 (Haar Randomization) 을 적용하여 최종적으로 대칭성을 갖춘 불변 분포를 복원합니다.
아키텍처: Canon
기존 SemlaFlow 와 같은 베이스 모델에 Canonical Rank Embedding을 추가하거나, 원자의 정규화된 순서 정보를 숨은 상태 (Hidden State) 로 명시적으로 유지하는 새로운 아키텍처인 Canon을 설계했습니다.
이를 통해 모델은 대칭성을 깨뜨린 상태에서도 표현력을 극대화할 수 있습니다.
3. 주요 기여 및 이론적 분석 (Key Contributions & Theory)
정규화 확산의 이론적 근거:
보편성 (Universality): 정규화된 슬라이스 (Slice) 에서 학습하고 Haar 무작위화를 통해 샘플링하는 방식이 임의의 불변 타겟 분포를 표현하기에 충분함을 증명했습니다.
표현력 우위: 비공변적 백본을 정규화 공간에서 사용하면, 공변적 아키텍처보다 더 강력한 표현력을 가지면서도 계산 비용을 줄일 수 있음을 보였습니다.
학습 가속화 메커니즘:
조건부 분산 감소: 대칭성 모호성 (Symmetry Ambiguity) 으로 인한 분산 항을 제거하여, 플로우 매칭 (Flow Matching) 의 학습 오차 하한을 낮춥니다.
스코어 복잡성 제거: 군 혼합으로 인한 스코어 필드의 급격한 변화를 제거하여 학습 안정성을 높입니다.
상호 보완적 관계:
정규화 (대칭성 제거) 와 최적 수송 (OT, 궤적 직선화) 이 서로 보완적으로 작용하여 학습 효율성을 극대화함을 이론적으로 증명했습니다.
4. 실험 결과 (Results)
분자 생성 벤치마크 (QM9, GEOM-DRUG) 에서 실험을 수행했습니다.
성능 향상:
GEOM-DRUG: 제안한 CanonFlow는 기존 최상위 모델 (SOTA) 인 SemlaFlow, EQGAT-diff 등을 크게 앞지르며, 분자 안정성 (Molecule Stability) 과 유효성 (Validity) 에서 새로운 SOTA 를 달성했습니다.
QM9: 기존 모델 대비 분자 안정성과 유효성이 개선되었으며, 최적화된 RMSD (Opt-RMSD) 가 26% 감소하여 에너지 최소화 구조에 더 가까운 분자를 생성함을 보였습니다.
소수 단계 생성 (Few-step Generation):
학습된 모델은 50100 단계의 샘플링으로도 기존 모델 (5001000 단계) 과同等하거나 더 좋은 성능을 보였습니다. 이는 정규화가 학습된 동역학을 직선화하여 빠른 수렴을 가능하게 했음을 시사합니다.
계산 효율성:
정규화 과정은 학습/추론 시 추가적인 계산 오버헤드를 거의 발생시키지 않으며, 오히려 학습 수렴 속도를 높였습니다.
5. 의의 및 결론 (Significance)
패러다임 전환: 분자 생성 분야에서 "대칭성을 아키텍처에 강제로 주입하는 것"이 아니라, "학습 공간을 정규화하여 대칭성을 깨뜨린 후 복원하는 것"이 더 효율적이고 강력함을 입증했습니다.
실용성: 복잡한 공변적 레이어 없이도 표준 Transformer 나 GNN 을 사용하여 고품질 3D 분자를 생성할 수 있는 길을 열었습니다.
확장성: 이 프레임워크는 3D 분자 생성뿐만 아니라, 대칭성을 가진 다른 과학적 데이터 생성 작업에도 적용 가능한 일반적인 원리를 제공합니다.
요약하자면, 이 논문은 **정규화 (Canonicalization)**를 통해 대칭성으로 인한 학습의 불확실성을 제거하고, 비공변적 모델의 강력한 표현력을 활용하여 더 빠르고 정확한 3D 분자 생성을 가능하게 하는 이론적, 실증적 성과를 제시했습니다.