[논문 리뷰] One Transformer Can Understand Both 2D & 3D Molecular Data

Shengjie Luo, Tianlang Chen외 5명|arXiv (Cornell University)|2022. 10. 04.

Machine Learning in Materials Science인용 수 4

한 줄 요약

이 논문은 2D 분자 그래프와 3D 기하학적 구조를 별도이지만 공유되는 어텐션 채널을 사용해 통합적으로 인코딩하는 Transformer-M이라는 통합형 트랜스포머 기반 모델을 제안한다. 자기 어텐션 기반 구조에서 2D 또는 3D 구조 인코딩을 동적으로 활성화하는 바이어스 항목으로 활용함으로써, 모델은 2D 및 3D 분자 성질 예측 작업 모두에서 뛰어난 성능을 달성하며, 아키텍처 재학습 없이도 다양한 모odal 간 일반화 성능을 보여준다.

ABSTRACT

Unlike vision and language data which usually has a unique format, molecules can naturally be characterized using different chemical formulations. One can view a molecule as a 2D graph or define it as a collection of atoms located in a 3D space. For molecular representation learning, most previous works designed neural networks only for a particular data format, making the learned models likely to fail for other data formats. We believe a general-purpose neural network model for chemistry should be able to handle molecular tasks across data modalities. To achieve this goal, in this work, we develop a novel Transformer-based Molecular model called Transformer-M, which can take molecular data of 2D or 3D formats as input and generate meaningful semantic representations. Using the standard Transformer as the backbone architecture, Transformer-M develops two separated channels to encode 2D and 3D structural information and incorporate them with the atom features in the network modules. When the input data is in a particular format, the corresponding channel will be activated, and the other will be disabled. By training on 2D and 3D molecular data with properly designed supervised signals, Transformer-M automatically learns to leverage knowledge from different data modalities and correctly capture the representations. We conducted extensive experiments for Transformer-M. All empirical results show that Transformer-M can simultaneously achieve strong performance on 2D and 3D tasks, suggesting its broad applicability. The code and models will be made publicly available at https://github.com/lsj2408/Transformer-M.

연구 동기 및 목표

• 특정 아키텍처를 전문화하지 않고도 2D 분자 그래프와 3D 기하학적 구조를 모두 처리할 수 있는 일반 목적의 딥러닝 모델을 개발하는 것.
• 기존 모델들이 하나의 분자 데이터 포맷 전용으로만 훈련되어 있어 모달 간 제로샷 전이 성능이 열악한 문제를 해결하는 것.
• 단일 모델을 설계하여 2D 및 3D 데이터를 동시에 학습할 수 있도록 모달 적응형 어텐션 메커니즘을 도입하는 것.
• 공유 파rameter와 모달 특화 구조 인코딩을 사용함으로써 단일 신경망이 분자 데이터 모달 간 효과적으로 학습하고 일반화할 수 있음을 입증하는 것.

제안 방법

• 표준 트랜스포머 백본을 사용하며, 2D 및 3D 구조 인코딩을 위한 별도의 모달 특화 채널을 두 개 병렬로 운영한다.
• 2D 채널은 차수 인코딩, 최단 경로 거리 인코딩, 엣지 인코딩을 입력 바이어스로 사용하여 위상 및 결합 특징을 포착한다.
• 3D 채널은 원자 간 3D 유클리드 거리에 기반한 가우시안 커널 기반 거리 인코딩을 어텐션 바이어스 항목으로 사용한다.
• 추론 시 입력 모달에 따라 활성화된 채널(2D 또는 3D)만 사용하며, 공유 어텐션 및 피드포워드 레이어는 양쪽 모두에 동일하게 적용된다.
• 일괄 처리에서 2D, 3D 또는 양방향 데이터 유형을 동시에 학습할 수 있도록 공동 훈련 전략을 도입하며, 비활성화된 채널은 활성화되지 않도록 한다.
• 3D 위치 노이즈 제거 목표를 사용해 PCQM4Mv2 데이터셋에서 사전 훈련하고, 최종적으로 2D 및 3D 데이터셋에서 미세 조정한다.

실험 결과

연구 질문

• RQ1 단일 신경망 아키텍처가 2D 분자 그래프와 3D 기하학적 구조 양쪽에서 의미 있는 표현을 효과적으로 학습할 수 있는가?
• RQ2 어텐션 기반 구조에서 모달 특화 구조 인코딩이 분자 데이터 포맷 간 모델 일반화에 어떤 영향을 미치는가?
• RQ3 2D 및 3D 데이터를 공동으로 사전 훈련하면 다양한 분자 성질 예측 작업의 최종 성능 향상 정도는 어느 정도인가?
• RQ4 모달 간 공유 파rameter 사용이 모달 특화 모델 대비 더 나은 일반화 성능을 이끌어내는가?

주요 결과

• PDBBind 코어 세트에서 Transformer-M은 피어슨 상관계수(R) 0.830을 기록하여 최고의 베이스라인(SIGN) 및 사전 훈련된 Graphormer 모델을 능가한다.
• QM9 데이터셋에서 2D-3D 공동 사전 훈련 전략은 초기 학습 대비 평균 제곱오차(MAE)를 최대 30% 감소시키며, 모든 6개의 분자 타겟에서 성능 향상을 보였다.
• DFT로 생성한 3D 커파이머를 사용할 경우 성능 향상이 뚜렷했으며(예: R 0.804에서 0.830으로 상승), RDKit로 생성한 커파이머의 경우는 미미한 향상만 보였다.
• 3D 위치 노이즈 제거를 통해 사전 훈련한 경우, 일관된 최종 성능 향상이 관찰되어, 3D 분자 표현 학습에 효과적인 사전 훈련 목표임을 입증했다.
• 3D 데이터로 훈련한 경우에도 2D 작업에서 뛰어난 성능 유지를 보이며, 강력한 교차 모달 전이 능력을 입증했다.
• 제거 실험 결과, 모달 특화 구조 인코딩이 필수적임을 확인하였으며, 어느 한 채널도 비활성화할 경우 해당 모달에서 성능 저하가 발생했다.