[논문 리뷰] E(n) Equivariant Normalizing Flows for Molecule Generation in 3D

Victor García Satorras, Emiel Hoogeboom외 3명|arXiv (Cornell University)|2021. 05. 19.

Machine Learning in Materials Science인용 수 17

한 줄 요약

이 논문은 E(n) 등변 정규화 플로우(E-NFs)를 제안하며, 이는 연속 시간 정규화 플로우에 E(n) 등변 그래프 신경망을 통합함으로써 3D 분자의 생성을 학습하는 가능도 기반 생성 모델이다. 이 방법은 QM9의 3D 입자 시스템과 분자에서 최신 기준 성능을 달성하며, 정확한 E(n) 등변성을 갖춘 최초의 가능도 기반 3D 분자 생성 모델임을 나타낸다.

ABSTRACT

This paper introduces a generative model equivariant to Euclidean symmetries: E(n) Equivariant Normalizing Flows (E-NFs). To construct E-NFs, we take the discriminative E(n) graph neural networks and integrate them as a differential equation to obtain an invertible equivariant function: a continuous-time normalizing flow. We demonstrate that E-NFs considerably outperform baselines and existing methods from the literature on particle systems such as DW4 and LJ13, and on molecules from QM9 in terms of log-likelihood. To the best of our knowledge, this is the first likelihood-based deep generative model that generates molecules in 3D.

연구 동기 및 목표

• 유럽 기하학적 변환(회전 및 이동)에 대한 정확한 불변성을 갖춘 3D 분자 구조를 생성하는 딥 생성 모델을 개발하는 것.
• 역행 가능하고 등변적인 생성을 위해 연속 시간 정규화 플로우 프레임워크에 E(n) 등변 메시지 전달 신경망을 통합하는 것.
• QM9 및 DW4, LJ13와 같은 벤치마크 데이터셋을 포함한 3D 분자 및 입자 시스템에서 최신 기준 가능도 성능를 달성하는 것.
• 문헌에서 3D 분자 생성에 대한 관심이 증가하고 있음에도 불구하고, 3D 분자에 대한 가능도 기반 생성 모델의 격차를 메우는 것.

제안 방법

• 모델은 연속 정규화 플로우의 신경망 구성 요소로 E(n) 등변 그래프 신경망을 사용하여, 회전 및 이동에 대한 등변성을 보장한다.
• 플로우는 GNN에서 유도된 학습된 벡터장으로 정의되며, ODE 솔버를 통해 통합되어 역행 가능한 변환을 생성한다.
• 정규화 플로우는 가능도 최대화를 통해 엔드 투 엔드로 훈련되며, 로그 밀도의 변화는 자코비안의 추적값을 통해 계산된다.
• 유지되는 정확한 E(n) 등변성은 원자 위치와 특징을 처리하는 등변 GNN에서 유도된 플로우 벡터장을 구성함으로써 달성된다.
• 모델은 3D 입자 시스템과 분자에서 훈련되며, 생성을 위해 잠재공간 샘플링을 사용한다.
• 이 방법은 생성과 가능도 평가를 모두 가능하게 하여 다른 가능도 기반 모델과 직접 비교할 수 있다.

실험 결과

연구 질문

• RQ1 3D 분자 생성에서 정확한 유클리드 대칭성(E(n))에 대해 등변인 정규화 플로우 모델을 설계할 수 있는가?
• RQ2 E(n) 등변성은 비등변 기반 모델과 비교해 3D 분자 및 입자 시스템에서 가능도와 생성 품질을 어떻게 향상시키는가?
• RQ3 정확한 기하학적 불변성을 유지하면서도 가능도 기반 모델이 3D 분자 생성에서 최신 기준 성능를 달성할 수 있는가?
• RQ4 등변 GNN을 연속 정규화 플로우 프레임워크에 통합하여 역행 가능한 생성 모델링을 수행하는 것은 가능한가?

주요 결과

• E-NFs는 QM9 데이터셋에서 최신 기준 가능도 성능를 달성하며, 기존의 3D 분자 생성 방법을 능가한다.
• DW4 및 LJ13와 같은 벤치마크 3D 입자 시스템에서 E-NFs는 가능도 측면에서 뛰어난 성능를 보였다.
• E-NFs는 정확한 E(n) 등변성을 유지하는 최초의 가능도 기반 딥 생성 모델로, 3D 분자에 적용된다.
• E(n) 등변 GNN을 연속 정규화 플로우 프레임워크에 통합함으로써 고품질의 생성과 정확한 가능도 계산이 가능해졌다.