[논문 리뷰] Score-Based Generative Models for Molecule Generation

Dwaraknath Gnaneshwar, Bharath Ramsundar외 3명|arXiv (Cornell University)|2022. 03. 07.

Machine Learning in Materials Science인용 수 7

한 줄 요약

이 논문은 150만 개의 ZINC 분자를 대상으로 한 SELFIES 표현을 기반으로 한 Transformer 기반 점수 함수를 훈련시켜 분자 생성을 위한 점수 기반 생성 모델을 제안한다. 모델은 다양하고 유효하며 새로운 화합물을 생성하지만, 훈련 분포에서 크게 벗어나 있어 복잡한 화학적 공간을 모델링하기 위한 아키텍처 개선의 잠재적 여지를 보여준다.

ABSTRACT

Recent advances in generative models have made exploring design spaces easier for de novo molecule generation. However, popular generative models like GANs and normalizing flows face challenges such as training instabilities due to adversarial training and architectural constraints, respectively. Score-based generative models sidestep these challenges by modelling the gradient of the log probability density using a score function approximation, as opposed to modelling the density function directly, and sampling from it using annealed Langevin Dynamics. We believe that score-based generative models could open up new opportunities in molecule generation due to their architectural flexibility, such as replacing the score function with an SE(3) equivariant model. In this work, we lay the foundations by testing the efficacy of score-based models for molecule generation. We train a Transformer-based score function on Self-Referencing Embedded Strings (SELFIES) representations of 1.5 million samples from the ZINC dataset and use the Moses benchmarking framework to evaluate the generated samples on a suite of metrics.

연구 동기 및 목표

• 점수 기반 생성 모델이 신규 분자 생성에 대해 실현 가능성과 효과성을 탐구하기 위해.
• GAN 및 정규화 흐름의 단점(예: 훈련 불안정성, 아키텍처 제약)을 극복하기 위해 점수 기반 모델의 아키텍처 유연성을 활용하기 위해.
• SELFIES 표현을 사용하여 분자 데이터에 대한 점수 기반 모델 훈련을 위한 기본 최적 실천 방안을 수립하기 위해.
• MOSES 벤치마크 세트를 사용하여 유효성, 유일성, 다양성, 분포 충실도 지표를 통해 모델 성능을 평가하기 위해.
• 향후 3차원 분자 생성 향상을 위해 등변성 아키텍처(예: SE(3)-등변 모델)를 점수 함수에 통합할 잠재력을 탐색하기 위해.

제안 방법

• 자기 잠재 공간 내에서 분자 데이터의 로그 밀도의 기울기(즉, 점수)를 추정하기 위해 Transformer 기반 점수 함수를 훈련한다.
• 다중 노이즈 수준(σ₁ = 10, σₗ = 0.01, 350개 수준)을 사용한 노이즈 제거 점수 매칭을 통해 훈련 안정성 향상 및 점수 추정 개선을 도모한다.
• 냉각된 랭글레인 역동법을 사용한 샘플링: 초기 노이즈를 무작위로 설정하고 점차 감소하는 노이즈 수준을 거쳐 점수 함수를 반복적으로 적용하여 노이즈 제거를 수행한다.
• 분자를 일괄 인코딩된 SELFIES 문자열로 표현하고 이를 벡터로 평탄화하여 점수 네트워크의 입력으로 사용한다.
• 다음 목적 함수를 사용하여 점수 함수를 최적화한다: ∑ᵢ λ(i)𝔼[‖sθ(x,i) − ∇ₓlog pσᵢ(x)‖²], 여기서 λ(i)는 다양한 노이즈 수준에 대한 가중치이다.
• 단일 V100 GPU에서 약 12시간 동안 훈련하고, T=10개의 시간 단계로 약 4시간 내에 3만 개의 분자를 샘플링한다.

실험 결과

연구 질문

• RQ1 SELFIES로 인코딩된 분자 데이터에 대해 훈련된 점수 기반 생성 모델이 화학적으로 유효하고 다양한 분자를 효과적으로 생성할 수 있는가?
• RQ2 기존의 모델들(예: VAE, GAN)과 비교해 볼 때, Transformer 기반 점수 함수는 표준 분자 생성 벤치마크에서 어떤 성능을 보이는가?
• RQ3 생성된 분자들이 훈련 데이터(160만 개의 ZINC)의 통계 분포와 얼마나 유사한가? 분포 이탈을 유도할 수 있는 요인들은 무엇인가?
• RQ4 점수 기반 모델의 아키텍처의 민첩성은 향후 등변성 레이어(예: SE(3)-등변 GNN)를 통합하여 3차원 분자 생성을 향상시키는 데 어떤 잠재력을 지닐 수 있는가?
• RQ5 점수 기반 분자 생성에 있어 가장 효과적인 하이퍼파rameter와 훈련 전략(예: 노이즈 스케줄, 모델 깊이)은 무엇인가?

주요 결과

• 모델은 생성된 3만 개의 분자 중 100% 유효성과 100% 신규성을 달성하여 구조적으로 유효하고 완전히 새로운 화합물을 생성함을 확인하였다.
• 1,000개 기준으로 88%, 10,000개 기준으로 82%의 유일성 점수는 최첨단 VAE보다 낮지만 여전히 경쟁 가능하며, 하이퍼파ram터 튜닝을 통해 향상 가능할 잠재성을 시사한다.
• IntDiv₁ 및 IntDiv₂ 점수 각각 0.90과 0.88은 생성된 집합에서 강력한 화학적 다양성과 모드 붕괴의 부재를 나타낸다.
• FCD/Test 및 FCD/TestSF 점수 각각 39.84와 40.92는 훈련 데이터로부터의 분포적 이탈이 심각하다는 것을 보여주며, 생성된 분자가 실제 분자 분포와 유사하지 않음을 시사한다.
• 모델의 성능은 높은 다양성과 신규성의 대가로 분포 충실도가 떨어지는 상충관계를 드러내며, 점수 함수의 능력 향상 또는 아키텍처의 인덕티브 바이어스 개선이 필요함을 시사한다.
• 결과는 점수 기반 모델을 사용한 분자 생성의 실현 가능성을 확인하였으며, 향후 등변성 아키텍처의 통합을 통해 분자 대칭성을 더 잘 포착하고 분포 모델링을 향상시킬 수 있는 길을 열어 놓았다.