[논문 리뷰] MUDiff: Unified Diffusion for Complete Molecule Generation

Chenqing Hua, Sitao Luan외 5명|arXiv (Cornell University)|2023. 04. 28.

Machine Learning in Materials Science인용 수 15

한 줄 요약

MUDiff는 2D 그래프(엣지)와 3D 좌표(원자) 모두에서 확산을 통해 완전한 분자 표현을 생성하고, 회전-병진 대칭 보정을 수행하는 새로운 MUformer를 사용해 노이즈 제거를 수행합니다.

ABSTRACT

Molecule generation is a very important practical problem, with uses in drug discovery and material design, and AI methods promise to provide useful solutions. However, existing methods for molecule generation focus either on 2D graph structure or on 3D geometric structure, which is not sufficient to represent a complete molecule as 2D graph captures mainly topology while 3D geometry captures mainly spatial atom arrangements. Combining these representations is essential to better represent a molecule. In this paper, we present a new model for generating a comprehensive representation of molecules, including atom features, 2D discrete molecule structures, and 3D continuous molecule coordinates, by combining discrete and continuous diffusion processes. The use of diffusion processes allows for capturing the probabilistic nature of molecular processes and exploring the effect of different factors on molecular structures. Additionally, we propose a novel graph transformer architecture to denoise the diffusion process. The transformer adheres to 3D roto-translation equivariance constraints, allowing it to learn invariant atom and edge representations while preserving the equivariance of atom coordinates. This transformer can be used to learn molecular representations robust to geometric transformations. We evaluate the performance of our model through experiments and comparisons with existing methods, showing its ability to generate more stable and valid molecules. Our model is a promising approach for designing stable and diverse molecules and can be applied to a wide range of tasks in molecular modeling.

연구 동기 및 목표

• 완전한 분자 표현을 위한 2D(그래프 연결성)와 3D(기하학적) 분자 데이터를 공동 생성하려는 동기 부여.
• 연속(원자 특징, 좌표)과 이산(엣지 타입) 구성 요소를 모두 디노이즈하는 확산 기반 프레임워크를 개발.
• MUformer를 도입하여 3D 회전-병진 대칭을 보존하면서 2D와 3D 정보를 통합하는 등변 그래프 트랜스포머를 제시.
• 생성 및 학습 중 2D 또는 3D 데이터 누락에 대한 강건성을 확보.
• 기존 방법에 비해 생성된 분자의 안정성 및 다양성 향상을 입증합니다.

제안 방법

• MUDiff 제안: 원자 특징과 좌표에 연속 노이즈를, 엣지 타입에 이산 노이즈를 추가하여 모든 구성 요소를 공동으로 디노이즈하는 확산 모델.
• MUformer 도입: 2D 및 3D 분자 데이터를 회전-병진 제약하에서 처리하기 위한 불변 및 등변 채널을 갖춘 통합 트랜스포머.
• 원자 특징과 좌표에 대한 노이즈를 예측하고 엣지 타입을 분류하는 학습 목표를 정의.
• 2D 이웃 정보, 3D 이웃 정보, 전역 그래프 특징을 융합하는 인코딩 스킴을 도입하여 강건한 표현을 형성.
• 공간 정보를 포착하고 등변성을 보장하기 위해 3D 전용 방사 기저 및 코사인 컷오프 함수를 적용.
• 완전한 분자에서 점진적으로 노이즈를 제거하는 샘플링 절차를 제공.

Figure 1: The figure showcases our MUformer for processing 2D and 3D molecular data. Within the Transformer backbone, two channels exist: purple for 2D data and brown for 3D data. The blue part encodes 2D molecular structures, while the green part handles atom-level information and the red part proc

실험 결과

연구 질문

• RQ1 2D 그래프 구조와 3D 좌표에 대한 공동 확산이 단일 표현 모델보다 더 안정적이고 타당한 분자를 생성할 수 있는가?
• RQ2 2D와 3D 분자 정보를 효과적으로 통합하면서 3D 회전-병진에 대해 등변성을 가지도록 트랜스포머를 설계할 수 있는가?
• RQ3 생성 또는 학습 중 2D 또는 3D 데이터가 누락된 경우 모델은 얼마나 강건한가?
• RQ4 2D 및 3D 구조를 공동 모델링함으로써 얻는 성능 향상(안정성, 고유성)은 무엇인가?

주요 결과

• MUDiff는 기존 방법보다 분자가 7.9% 더 안정적임(Sec 6.2).
• MUDiff는 기존 방법과 비교하여 분자의 고유성을 2% 증가시킴(Sec 6.2).
• MUformer는 회전-병진 등변성을 갖춘 원자 특징, 좌표, 엣지 타입의 동시 예측을 가능하게 함.
• 모델은 3D 구조가 제한적으로 학습되어도 효과를 유지하며, 완전한 3D 데이터로 학습된 방법과 비교해 경쟁력 있는 성능을 보임(Sec 6.1).
• 2D 또는 3D 정보가 누락되더라도 독립적으로 작동하여 강건한 완전 표현을 가능하게 함.