[논문 리뷰] Graph Diffusion Transformers for Multi-Conditional Molecular Generation

Gang Liu, Jiaxin Xu외 2명|arXiv (Cornell University)|2024. 01. 24.

Machine Learning in Materials Science인용 수 5

한 줄 요약

다중 조건 확산(MCD)을 제안하여 다중 속성 제약으로 그래프 확산 모델을 안내하고, 고분자와 소분자 모두에서 수치적 및 범주형 분자 속성을 동시에 제어할 수 있도록 한다.

ABSTRACT

Inverse molecular design with diffusion models holds great potential for advancements in material and drug discovery. Despite success in unconditional molecular generation, integrating multiple properties such as synthetic score and gas permeability as condition constraints into diffusion models remains unexplored. We present the Graph Diffusion Transformer (Graph DiT) for multi-conditional molecular generation. Graph DiT integrates an encoder to learn numerical and categorical property representations with the Transformer-based denoiser. Unlike previous graph diffusion models that add noise separately on the atoms and bonds in the forward diffusion process, Graph DiT is trained with a novel graph-dependent noise model for accurate estimation of graph-related noise in molecules. We extensively validate Graph DiT for multi-conditional polymer and small molecule generation. Results demonstrate the superiority of Graph DiT across nine metrics from distribution learning to condition control for molecular properties. A polymer inverse design task for gas separation with feedback from domain experts further demonstrates its practical utility.

연구 동기 및 목표

• 여러 속성이 충족되어야 할 때 확산 모델을 활용한 역분자 설계의 동기를 부여한다.
• 다양한 수치적 및 범주형 제약을 표현하고 통합하기 위한 다중 조건 가이던스 메커니즘을 개발한다.
• 생성을 가이드하기 위해 조건 표현을 사용하는 그래프 의존 확산 과정과 Transformer 기반 디노이징 모델을 설계한다.
• 고분자 및 소분자 데이터세트에서 다중 조건 생성의 시연을 보이고, 가스 분리용 역고분자 설계를 통한 실용적 활용도를 평가한다.

제안 방법

• 다중 수치적 및 범주적 조건에 대한 표현을 학습하는 조건 인코더를 포함한 다중 조건 확산(MCD)을 도입한다.
• Transformer 기반 구조 인코더에서 분자 통계치를 조건 통계치로 대체하기 위해 적응형 레이어 정규화를 사용한다.
• 노드-에지 확산 행렬 Q_G를 포함하는 그래프 의존 노이즈 모델을 제안하여 노이즈를 분자 그래프 구조에 더 잘 맞춘다.
• 예측기 없는 가이던스 전략을 사용하여 무조건적 및 조건부 디노이징 확률을 가변 스케일 매개변수 s와 함께 결합하는 것을 적용한다.
• 누락된 조건을 처리하기 위해 드롭핑 임베딩을 사용하고, 하나의 그래프 트랜스포머를 학습시켜 무조건적 및 조건부 디노이징을 모두 수행한다.
• 조건 인코더, 구조 인코더, 구조 디코더의 세 부분으로 구성된 아키텍처와 그래프로부터 분자를 생성하기 위한 실용적 변환 단계를 제시한다.

(a) Existing work’s limitation: A median rank of 30 showed that on fewer than half test polymers, the sets of generated graphs from different single conditions intersected, indicating a failure to generate polymers meeting multiple properties.

실험 결과

연구 질문

• RQ1 다중 수치적 및 범주형 속성을 확산 기반 분자 생성에 어떻게 통합하되 스케일이나 유형이 서로 뒤섞이지 않도록 할 수 있는가?
• RQ2 그래프 인지적 노이즈 모델이 다중 조건 가이던스 하에서 생성된 분자 그래프의 현실성 및 타당성을 향상시킬 수 있는가?
• RQ3 다중 조건 확산 가이던스가 단일 조건 기준선보다 여러 속성 제약을 더 잘 만족하는 생성을 가능하게 하는가?
• RQ4 확산 기반 분자 설계에서 다중 속성 결과를 제어하는 데 예측기 없는 가이던스가 효과적인가?
• RQ5 이 접근법이 고분자와 소분자 모두에 확장되어 역설계 작업에서 실용적 활용도를 제공할 수 있는가?

주요 결과

• MCD는 단일 조건 기준선보다 다중 속성 제약에 더 잘 맞춘 고분자를 생성했으며, 테스트된 시나리오에서 중간 순위가 30을 크게 상회했다.
• 모델은 다수의 수치 속성에 대해 고분자 데이터셋에서 평균 MAE를 더 낮게 달성하여 최상의 기준선에 비해 오차를 17.86% 감소시켰다.
• 소분자에 대해서는 MCD가 작업 관련 범주 조건에서 정확도 0.9 이상을 달성하여 기준선 정확도(0.6 이하)를 상회했다.
• O2/N2 가스 분리용 역고분자 설계는 도메인 전문가의 피드백이 다중 조건 설계 이점을 뒷받침하는 실용적 가치를 보여준다.
• 본 접근법은 고분자 및 약물 설계 데이터셋에서 분포 학습 및 조건 제어 지표에서 강력한 성능을 입증한다.

(b) Proposed work: Our idea, multi-conditional guidance for diffusion models, successfully generated polymers that satisfied multi-property constraints. It achieved a higher rank than 30 in any set of the single-conditional generated graphs.