[논문 리뷰] 3D-Transformer: Molecular Representation with Transformer in 3D Space

Fang Wu, Qiang Zhang외 6명|arXiv (Cornell University)|2021. 10. 04.

Machine Learning in Materials Science참고 문헌 79인용 수 10

한 줄 요약

3D-Transformer는 다중 척도 자기주의 메커니즘과 적응형 위치 인코딩을 통해 공간 기하학을 통합함으로써 분자의 표현 학습을 위한 3D 인식 Transformer 아키텍처를 도입한다. 주어진 분자의 크기 변화에 걸쳐 원자 간 거리의 맥락 인식 모델링과 주목사의 다운샘플링을 활용하여 결정성질 예측 및 단백질-리간드 결합 친화도 과제에서 최신 기술 수준(SOTA) 성능을 달성한다.

ABSTRACT

Spatial structures in the 3D space are important to determine molecular properties. Recent papers use geometric deep learning to represent molecules and predict properties. These papers, however, are computationally expensive in capturing long-range dependencies of input atoms; and have not considered the non-uniformity of interatomic distances, thus failing to learn context-dependent representations at different scales. To deal with such issues, we introduce 3D-Transformer, a variant of the Transformer for molecular representations that incorporates 3D spatial information. 3D-Transformer operates on a fully-connected graph with direct connections between atoms. To cope with the non-uniformity of interatomic distances, we develop a multi-scale self-attention module that exploits local fine-grained patterns with increasing contextual scales. As molecules of different sizes rely on different kinds of spatial features, we design an adaptive position encoding module that adopts different position encoding methods for small and large molecules. Finally, to attain the molecular representation from atom embeddings, we propose an attentive farthest point sampling algorithm that selects a portion of atoms with the assistance of attention scores, overcoming handicaps of the virtual node and previous distance-dominant downsampling methods. We validate 3D-Transformer across three important scientific domains: quantum chemistry, material science, and proteomics. Our experiments show significant improvements over state-of-the-art models on the crystal property prediction task and the protein-ligand binding affinity prediction task, and show better or competitive performance in quantum chemistry molecular datasets. This work provides clear evidence that biochemical tasks can gain consistent benefits from 3D molecular representations and different tasks require different position encoding methods.

연구 동기 및 목표

• 기존 기하학적 딥 러닝 모델이 장거리 상관관계를 포착하는 데서 비효율하고 분자 내 비균일한 원자 간 거리 분포를 처리하는 데에 한계를 보이는 문제를 해결하기 위해.
• 다양한 공간 척도와 분자 크기에 적응 가능한 3D 인식 주목사 메커니즘을 개발하기 위해.
• 분자의 표현을 집계하는 데 있어 고정 또는 거리 중심의 다운샘플링 방법의 단점을 극복하기 위해.
• 분자 크기에 따라 동적으로 적응하는 위치 인코딩 전략을 설계하여 표현 품질을 향상시키기 위해.
• 3D-Transformer의 효과성을 양자 화학, 재료 과학 및 단백질체학을 포함한 다양한 과학 분야에서 검증하기 위해.

제안 방법

• 모든 원자 상호작용을 모델링하기 위해 완전히 연결된 그래프를 사용하여 3D 공간 내 직접적인 원자 간 연결을 유지한다.
• 수렴 영역이 점차 증가하는 방식으로 주목사에 의해 작동하는 다중 척도 자기주의 모듈을 도입하여 세밀한 국소 패턴과 광범위한 맥락적 의존성을 모두 포착한다.
• 분자 크기에 따라 다른 인코딩 전략을 선택하는 적응형 위치 인코딩 모듈을 활용하여 소분자에는 상대적 위치 인코딩을, 대분자에는 학습된 절대 인코딩을 사용한다.
• 주목사 점수를 사용하여 대표성을 갖춘 원자를 선택하는 주목사 기반 가장 먼 점 샘플링 방법을 제안하여 가상의 노드나 거리 기반 풀링 전략을 능가한다.
• 3D 좌표를 Transformer 인코더의 입력 특징으로 활용하여 주목사 계산에 공간 유도 편향을 제공한다.
• 표준 손실 함수를 사용하여 분자 성질 예측 과제에서 모델을 종단 간(end-to-end)으로 훈련시킨다.

실험 결과

연구 질문

• RQ1 3D 인식 Transformer 아키텍처가 공간 기하학과 장거리 상관관계를 명시적으로 모델링함으로써 분자의 표현 학습을 향상시킬 수 있는가?
• RQ2 다중 척도 주목사는 다양한 원자 간 거리 분포에서 맥락 의존 패턴을 어떻게 향상시키는가?
• RQ3 적응형 위치 인코딩 전략은 고정 인코딩 방식에 비해 크기가 다른 분자에서 더 높은 성능을 내는가?
• RQ4 주목사 기반 다운샘플링은 전통적인 거리 기반 또는 가상 노드 풀링 전략보다 분자의 표현 학습에서 더 나은 성능을 내는가?
• RQ5 3D 공간 정보를 통합할 경우, 양자 화학, 재료 과학 및 단백질체학과 같은 다양한 과학 분야에서 성능 향상이 얼마나 이루어지는가?

주요 결과

• 3D-Transformer는 결정성질 예측 벤치마크에서 이전의 최신 기술 수준(SOTA) 모델을 능가하는 최신 기술 수준 성능을 달성한다.
• 단백질-리간드 결합 친화도 예측 과제에서 뚜렷한 향상이 나타나 생물학적 및 화학적 분야에서의 강력한 일반화 능력을 입증한다.
• 양자 화학 데이터셋에서 3D-Transformer는 기존 모델과 비교해 경쟁력 있거나 더 낫거나 성능을 발휘하여 전자 구조 예측의 효과성을 입증한다.
• 적응형 위치 인코딩 모듈은 분자 크기에 따라 인코딩 전략을 맞춤화함으로써 소분자와 대분자 양쪽에서 성능 향상을 이룬다.
• 다중 척도 자기주의 메커니즘은 비균일한 원자 간 거리 분포를 가진 분자에서도 장거리 상관관계와 세밀한 공간 패턴을 효과적으로 포착한다.
• 주목사 기반 가장 먼 점 샘플링 방법은 가상의 노드나 거리 기반 풀링 방법보다 더 정보가 풍부하고 대표적인 분자 임베딩을 생성한다.