[논문 리뷰] Symphony: Symmetry-Equivariant Point-Centered Spherical Harmonics for 3D Molecule Generation

Ameya Daigavane, Song Eun Kim외 2명|arXiv (Cornell University)|2023. 11. 27.

Machine Learning in Materials Science인용 수 4

한 줄 요약

Symphony는 원자들의 반경 및 각도 분포를 동시에 모델링하기 위해 점 중심의 구면 조화 함수 투영을 사용하는 $E(3)$-동변성 자동회귀 생성 모델을 소개한다. QM9에서 이전 자동회귀 모델을 능가하며, 확산 모델과 경쟁 가능한 성능을 달성하면서도 더 빠른 추론 속도와 새로운 구조에 대한 강력한 일반화 능력을 유지한다.

ABSTRACT

We present Symphony, an $E(3)$-equivariant autoregressive generative model for 3D molecular geometries that iteratively builds a molecule from molecular fragments. Existing autoregressive models such as G-SchNet and G-SphereNet for molecules utilize rotationally invariant features to respect the 3D symmetries of molecules. In contrast, Symphony uses message-passing with higher-degree $E(3)$-equivariant features. This allows a novel representation of probability distributions via spherical harmonic signals to efficiently model the 3D geometry of molecules. We show that Symphony is able to accurately generate small molecules from the QM9 dataset, outperforming existing autoregressive models and approaching the performance of diffusion models.

연구 동기 및 목표

• 분자에 내재된 전역 및 국소 대칭성을 포함한 분자의 $E(3)$ 대칭성을 존중하는 생성 모델을 개발하기 위해.
• 구면 조화 함수를 통한 고차수 $E(3)$-동변성 특징을 활용하여 자동회귀 생성에서 3D 분자 기하 구조의 표현을 향상시키기 위해.
• 단일 초점 원자 기반 프레임워크에서 신규 원자의 반경 및 각도 분포를 동시에 정확하게 모델링할 수 있도록 하기 위해.
• 학습 데이터의 암기 외적 진정된 일반화 능력을 평가하기 위해 알려지지 않은 분자 조각에 대한 일반화 성능을 평가하기 위해.
• 확산 모델보다 더 빠른 추론 속도를 확보하면서도 높은 유효성과 기하 정확도를 유지하기 위해.

제안 방법

• Symphony는 각 새로운 원자를 단일 초점 원자로부터 예측하는 점 중심의 자동회귀 생성 과정을 사용한다. 이 초점 원자가 국소 기준점으로 기능한다.
• 후보 원자의 각도 및 반경 분포를 나타내기 위해 구면 조화 함수 $Y_{l,m}$에서 유도된 고차수 $E(3)$-동변성 특징을 활용한다.
• 모델은 구면 조화 신호를 사용하여 새로운 원자 위치의 조건부 확률 분포를 파arameterize하여 반경 및 각도 성분을 동시에 예측할 수 있도록 한다.
• E(3)-동변성 특징을 사용한 메시지 전달은 회전 및 이동에 대한 불변성을 보장하여 분자의 대칭성을 유지한다.
• 텐서 곱을 사용하여 특징을 조합하고 $SO(3)$의 고차수 기저 표현을 생성함으로써 더 rich한 기하 모델링이 가능해진다.
• 생성된 분자에서 국소 환경 일치도를 평가하기 위해 이분광 기반의 새로운 각도 정확도 지표를 도입한다.

실험 결과

연구 질문

• RQ1 구면 조화 함수 투영은 자동회귀 생성에서 3D 분자 기하 구조의 더 정확하고 효율적인 모델링을 가능하게 하는가?
• RQ2 고차수 특징을 갖는 $E(3)$-동변성은 회전 불변 모델에 비해 생성 품질을 어떻게 향상시키는가?
• RQ3 Symphony는 알려지지 않은 분자 조각의 완성에 얼마나 잘 일반화되는가? 이는 학습 데이터 암기 외의 진정한 일반화를 나타내는가?
• RQ4 더 높은 계산 비용에도 불구하고, Symphony의 추론 속도는 확산 모델 및 이전 자동회귀 모델과 비교해 어떻게 되는가?
• RQ5 구면 조화 함수의 사용은 룩업 테이블 기반 방법에 비해 결합 길이 분포의 이산화 오차를 줄이는 데 기여하는가?

주요 결과

• Symphony는 QM9의 모든 지표에서 G-SchNet 및 G-SphereNet를 능가하며, 원자 종류 수의 제닝-쇼너 분산(Jensen-Shannon Divergence, JSD)은 0.0003, 국소 환경 수의 JSD는 0.0039를 기록한다.
• 10개의 빈도가 높은 국소 환경에 대해 이분광 MMD(Mean Maximum Discrepancy)가 0.0039로, G-SchNet(0.0150) 및 G-SphereNet(0.1016)보다 유의미하게 낮다.
• Symphony는 알려지지 않은 테스트 조각에서 98.66%의 유효 완성률(VCR)을 달성하여 G-SchNet(98.15%)를 능가하며 강력한 일반화 능력을 보였다.
• 단일 A5000 GPU에서 분자당 0.293초의 추론 속도를 기록하여, 더 높은 계산 비용에도 불구하고 EDM(0.930 sec/mol)보다 약 3배 빠르다.
• 모델은 약 80 에포크 만에 훈련이 끝나 G-SchNet(330) 및 EDM(1100)보다 훨씬 적은 에포크를 요구하여 데이터 효율성이 향상됨을 시사한다.
• 강력한 성능에도 불구하고, 원자 간 거리 분포의 정확도는 낮은 반경 이산화로 인해 떨어져 있어 결합 차수 할당에 영향을 미친다.