[논문 리뷰] 3D-MolT5: Leveraging Discrete Structural Information for Molecule-Text Modeling

Qizhi Pei, Rui Yan외 3명|arXiv (Cornell University)|2024. 06. 09.

Synthetic Organic Chemistry Methods인용 수 4

한 줄 요약

3D-MolT5는 E3FP 지문을 사용하는 새로운 3D 분자의 토큰화 방법을 통해 1D 분자 서열(SELFIES)과 3D 분자 구조를 통합하는 통합 T5 기반 프레임워크를 제안한다. 1D와 3D 토큰을 원자 수준에서 정렬하고 다중 작업에서 동시 사전 훈련을 수행함으로써 분자 성질 예측, 분자 캡션 생성, 텍스트 기반 분자 생성에서 최신 기술 수준(SOTA) 성능을 달성하며, ChEBI-20에서 정확한 매칭 점수 48.7%를 기록한다.

ABSTRACT

The integration of molecular and natural language representations has emerged as a focal point in molecular science, with recent advancements in Language Models (LMs) demonstrating significant potential for comprehensive modeling of both domains. However, existing approaches face notable limitations, particularly in their neglect of three-dimensional (3D) information, which is crucial for understanding molecular structures and functions. While some efforts have been made to incorporate 3D molecular information into LMs using external structure encoding modules, significant difficulties remain, such as insufficient interaction across modalities in pre-training and challenges in modality alignment. To address the limitations, we propose extbf{3D-MolT5}, a unified framework designed to model molecule in both sequence and 3D structure spaces. The key innovation of our approach lies in mapping fine-grained 3D substructure representations into a specialized 3D token vocabulary. This methodology facilitates the seamless integration of sequence and structure representations in a tokenized format, enabling 3D-MolT5 to encode molecular sequences, molecular structures, and text sequences within a unified architecture. Leveraging this tokenized input strategy, we build a foundation model that unifies the sequence and structure data formats. We then conduct joint pre-training with multi-task objectives to enhance the model's comprehension of these diverse modalities within a shared representation space. Thus, our approach significantly improves cross-modal interaction and alignment, addressing key challenges in previous work. Further instruction tuning demonstrated that our 3D-MolT5 has strong generalization ability and surpasses existing methods with superior performance in multiple downstream tasks. Our code is available at https://github.com/QizhiPei/3D-MolT5.

연구 동기 및 목표

• 기존 언어 모델이 입체화학과 3D 의존 작업에 핵심적인 3D 분자 구조를 효과적으로 모델링하지 못하는 데에 대응하기 위해.
• 기존 3D 구조 보강 언어 모델에서 발생하는 모odal 비일치 및 별도의 튜닝 문제를 해결하기 위해, 공통 토큰 공간 내에서 1D와 3D 표현을 통합함으로써.
• 통합 사전 훈련 및 지시 훈련 프레임워크를 통해 분자 서열과 3D 구조 정보의 엔드 투 엔드 학습을 가능하게 하기 위해.
• 세밀한 3D 부분구조 표현을 통합함으로써 분자 성질 예측, 분자 캡션 생성, 텍스트 유도 분자 생성과 같은 후행 작업에서 성능 향상을 이루기 위해.

제안 방법

• 특정 원자를 중심으로 한 부분구조를 나타내는 3D 토큰으로 3D 분자 구조를 이산화하는 데에 E3FP(확장된 3D 지문) 기반의 3D 분자 토큰화 방법을 도입한다.
• 1D SELFIES 토큰과 3D 구조 토큰을 원자 수준에서 정렬하고, 그들의 임베딩을 합쳐 공동 분자 표현을 형성한다.
• 다섯 가지 목표를 포함하는 다중 작업 사전 훈련 전략을 적용한다: 1D 노이즈 제거, 1D+3D 공동 노이즈 제거, 3D에서 1D로의 번역, 3D에서 텍스트로의 번역, 텍스트에서 1D 분자로의 번역.
• T5 기반 아키텍처를 활용하여 1D 서열, 3D 토큰, 텍스트를 동시에 인코딩함으로써 지시 훈련을 통해 제로샷 및 희소한 few-shot 일반화를 가능하게 한다.
• 다양한 분자 생성 및 이해 작업에 대응하기 위해 후행 데이터셋 여러 개에 지시 훈련을 적용한다.
• 1D와 3D 분자 표현을 모두 지원하는 통합 토큰 어휘를 활용하여 단일 모델 아키텍처 내에서 원활한 통합을 가능하게 한다.

실험 결과

연구 질문

• RQ1 E3FP 기반의 이산적이고 부분구조 기반의 표현 방식을 사용하여 3D 분자 구조를 효과적으로 토큰화하고 사전 훈련된 언어 모델에 통합할 수 있는가?
• RQ2 1D와 3D 모odal을 동시에 사전 훈련하면, 성질 예측 및 3D 기능 모델링과 같은 3D 의존 작업에서 성능 향상이 이루어지는가?
• RQ3 2D 또는 1D 전용 모델에 비해 3D 구조 정보는 텍스트 기반 분자 생성 정확도에 얼마나 기여하는가?
• RQ4 1D와 3D 토큰을 원자 수준에서 정렬함으로써, 텍스트 기반 설명에서 화학적으로 타당한 분자를 생성하는 데에 모델의 능력은 어떻게 향상되는가?

주요 결과

• 3D-MolT5는 텍스트 기반 분자 생성 작업에서 ChEBI-20 데이터셋에서 48.7%의 정확한 매칭 점수를 기록하여, GPT-4와 BioT5를 포함한 모든 기준 모델을 크게 앞서며 성능을 냈다.
• ChEBI-20에서 BLEU 점수 0.849와 text2mol 점수 0.574를 기록하여 높은 생성 품질과 기준 분자와의 높은 일치도를 입증했다.
• 제거 실험 결과, 3D 입력을 제거할 경우 HOMO-LUMO 갭 예측의 MAE가 0.0791에서 0.0968로 증가하여 3D 구조 정보가 3D 의존 작업에서 핵심적인 역할을 한다는 것을 확인했다.
• 1D+3D 공동 노이즈 제거 및 번역 사전 훈련 작업을 제거할 경우 성능 저하가 발생하여, 이 작업들이 3D 인식 표현 학습에 기여한다는 점을 입증했다.
• ChEBI-20 데이터셋에서 3D-MolT5는 100%의 유효성 비율을 기록하여 생성된 분자의 화학적 정확도가 매우 높다는 것을 보여주었다.
• 텍스트 기반 분자 생성 분야에서 GPT-4나 BioT5와 같은 대규모 모델들조차도 모든 지표에서 3D-MolT5가 뛰어난 성능을 보이며, 효율성과 효과성의 우수성을 입증했다.