[논문 리뷰] SemlaFlow -- Efficient 3D Molecular Generation with Latent Attention and Equivariant Flow Matching

Ross W. Irwin, Alessandro Tibo외 2명|arXiv (Cornell University)|2024. 06. 11.

Innovative Microfluidic and Catalytic Techniques Innovation인용 수 4

한 줄 요약

SemlaFlow는 확장 가능한 E(3)-동변성 메시지 전파 아키텍처인 Semla와 스케일 최적 운반 및 플로우 매칭을 결합하여 효율적인 3D 분자 생성을 가능하게 한다. 단 20개의 샘플링 스텝만으로도 기존 방법 대비 100배 빠른 속도를 달성하면서도 화학적으로 타당하고 낮은 비틀림 에너지를 가진 분자를 생성한다.

ABSTRACT

Methods for jointly generating molecular graphs along with their 3D conformations have gained prominence recently due to their potential impact on structure-based drug design. Current approaches, however, often suffer from very slow sampling times or generate molecules with poor chemical validity. Addressing these limitations, we propose Semla, a scalable E(3)-equivariant message passing architecture. We further introduce an unconditional 3D molecular generation model, SemlaFlow, which is trained using equivariant flow matching to generate a joint distribution over atom types, coordinates, bond types and formal charges. Our model produces state-of-the-art results on benchmark datasets with as few as 20 sampling steps, corresponding to a two order-of-magnitude speedup compared to state-of-the-art. Furthermore, we highlight limitations of current evaluation methods for 3D generation and propose new benchmark metrics for unconditional molecular generators. Finally, using these new metrics, we compare our model's ability to generate high quality samples against current approaches and further demonstrate SemlaFlow's strong performance.

연구 동기 및 목표

• 3D 분자 생성 모델에서 느린 샘플링이 강화 학습 및 단백질 포켓 유도 설계에 한계를 초래하는 핵심적 문제를 해결하기 위해.
• 기존 메시지 전파 모델보다 뛰어난 성능을 보이는 새로운 E(3)-동변성 아키텍처인 Semla를 도입하여 3D 분자 생성기의 확장성과 효율성을 향상시키기 위해.
• 분자의 크기를 고려해 동변성 최적 운반을 확장하여 분자 크기 기반 스케일 최적 운반을 도입함으로써 최적 운반 기반 훈련의 운반 비용을 감소시키기 위해.
• 기존의 유효성 및 유사도 지표를 넘어서 생성된 분자 구조의 품질을 더 잘 평가하기 위해 에너지 및 비틀림 에너지와 같은 새로운 평가 지표를 제안하기 위해.
• 높은 품질의 화학적으로 타당한 3D 분자 샘플을 단 20단계로 생성할 수 있음을 입증하여 약물 발굴 파이프라인에서의 실용적 구현 가능성을 보여주기 위해.

제안 방법

• SemlaFlow는 노이즈 분포를 분자 좌표의 데이터 분포로 운반하는 데 플로우 매칭을 사용하여 생성 모델을 훈련한다.
• 모델은 기존의 전치합 모델 대비 파rameter 효율성과 추론 속도를 향상시키는 새로운 E(3)-동변성 메시지 전파 아키텍처인 Semla를 활용한다.
• 분자의 크기를 고려한 새로운 OT 공식인 스케일 최적 운반을 도입하여 운반 비용을 감소시키고 노이즈 및 데이터 분포 간의 정렬을 향상시킨다.
• 신경망이 ODE 기반 생성 과정에서 노이즈에서 데이터로 가는 벡터장을 추정하도록 조건부 플로우 매칭을 사용하여 훈련한다.
• 시간에 따라 변화하는 조건부 플로우를 사용하여 기대 벡터장 오차를 최소화함으로써 시뮬레이션 없이 엔드 투 엔드 훈련이 가능하다.
• 추론 과정에서는 학습된 벡터장을 ODE 솔버를 사용하여 통합하며, 샘플링 스텝 수를 조절하여 속도와 품질의 균형을 맞출 수 있다.

Figure 1 : An overview of our approach. Firstly, prior coordinate noise is sampled according to the respective scale of the molecules (atom and bond type noise is not shown for brevity). Our trained model, SemlaFlow , transports this noise to the data distribution to generate a molecular sample.

실험 결과

연구 질문

• RQ1 확장 가능한 E(3)-동변성 메시지 전파 아키텍처가 기존의 전치합 기반 모델 대비 3D 분자 생성 모델의 효율성과 확장성에 크게 기여할 수 있는가?
• RQ2 스케일 최적 운반을 통해 동변성 최적 운반에 분자의 크기를 포함시킬 경우 운반 비용이 감소하고 생성 품질이 향상되는가?
• RQ3 스케일 최적 운반을 활용한 플로우 매칭이 단 20개의 샘플링 스텝으로도 SOTA 성능을 달성할 수 있는가?
• RQ4 생성된 분자의 에너지 및 비틀림 에너지는 실제 데이터 및 SOTA 모델과 비교해 어떻게 다른가? 이러한 지표들이 구조적 품질 평가에 더 물리적으로 의미 있는가?
• RQ5 현재 평가 벤치마크가 생성된 분자 구조의 물리적 타당성을 충분히 반영하지 못하는 정도는 어느 정도이며, 이를 어떻게 개선할 수 있는가?

주요 결과

• SemlaFlow는 벤치마크 데이터셋에서 단 100개의 샘플링 스텝만으로도 유효성, 안정성 및 구조적 품질 측면에서 기존 모델을 능가하는 SOTA 성능을 달성한다.
• 단 20개의 샘플링 스텝으로도 SemlaFlow는 SOTA 모델의 성능을 재현하면서도 샘플링 시간에서 100배 빠른 속도를 기록하여 추론 효율성에서 2개의 지수 단위 향상이 이루어졌다.
• SemlaFlow가 생성한 분자는 비틀림 에너지가 127.5 ± 2.1로 EQGAT-diff(150.2 ± 1.8)보다 낮아 더 뛰어난 구조적 안정성을 나타낸다.
• 스케일 최적 운반의 사용은 핵심적이다. 추론 분석 결과, 운반 비용을 크게 감소시키고 에너지 및 비틀림 에너지 지표를 향상시킨다.
• 100단계에서 SemlaFlow는 최소화된 에너지가 88.9 ± 1.3로 EQGAT-diff(95.2 ± 0.22)보다 낮아 더 나은 구조 최적화를 보였다.
• 제안된 평가 지표인 에너지 및 비틀림 에너지는 현재의 평가 기준이 구조적 품질을 과소 평가하고 있음을 드러내며, SemlaFlow는 이러한 더 물리적으로 의미 있는 지표에서 SOTA를 초월한다.