[Paper Review] 3D-MolT5: Leveraging Discrete Structural Information for Molecule-Text Modeling

Qizhi Pei, Rui Yan+3 more|arXiv (Cornell University)|Jun 9, 2024

Synthetic Organic Chemistry Methods4 citations

TL;DR

3D-MolT5 proposes a unified T5-based framework that integrates 1D molecular sequences (SELFIES) and 3D molecular structures via a novel 3D molecular tokenization method using E3FP fingerprints. By aligning 1D and 3D tokens at the atomic level and jointly pre-training on multiple tasks, it achieves state-of-the-art performance in molecular property prediction, molecule captioning, and text-to-molecule generation, with a 48.7% exact match score on ChEBI-20.

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.

Motivation & Objective

• To address the limitation of existing language models in modeling 3D molecular structures, which are critical for stereochemistry and 3D-dependent tasks.
• To overcome modality misalignment and separate tuning issues in existing 3D-structure-augmented language models by unifying 1D and 3D representations in a shared token space.
• To enable end-to-end learning of molecular sequence and 3D structural information through a unified pre-training and instruction tuning framework.
• To improve performance in downstream tasks such as molecular property prediction, molecule captioning, and text-guided molecule generation by integrating fine-grained 3D substructure representations.

Proposed method

• Introduces a 3D molecular tokenization method based on E3FP (Extended 3D Fingerprint) to convert 3D molecular structures into discrete 3D tokens, each representing a substructure centered on a specific atom.
• Aligns 1D SELFIES tokens and 3D structure tokens at the atomic level, summing their embeddings to form a joint molecular representation.
• Employs a multi-task pre-training strategy with five objectives: 1D denoising, 1D+3D joint denoising, 3D-to-1D translation, 3D-to-text translation, and text-to-1D molecule translation.
• Utilizes a T5-based architecture to jointly encode 1D sequences, 3D tokens, and text, enabling instruction tuning for zero-shot and few-shot generalization.
• Applies instruction tuning on multiple downstream datasets to adapt the model for diverse molecular generation and understanding tasks.
• Employs a unified token vocabulary that supports both 1D and 3D molecular representations, enabling seamless integration within a single model architecture.

Experimental results

Research questions

• RQ1 Can 3D molecular structures be effectively tokenized and integrated into a pre-trained language model using a discrete, substructure-based representation?
• RQ2 Does joint pre-training on 1D and 3D modalities improve performance on 3D-dependent molecular tasks such as property prediction and 3D function modeling?
• RQ3 To what extent does 3D structure information enhance text-to-molecule generation accuracy compared to 2D or 1D-only models?
• RQ4 How does the alignment of 1D and 3D tokens at the atomic level affect the model’s ability to generate chemically valid molecules from textual descriptions?

Key findings

• 3D-MolT5 achieves a 48.7% exact match score on the ChEBI-20 dataset for text-guided molecule generation, significantly outperforming all baselines, including GPT-4 and BioT5.
• The model achieves a 0.849 BLEU score and 0.574 text2mol score on ChEBI-20, demonstrating strong generation quality and alignment with reference molecules.
• Ablation studies show that removing 3D input increases the MAE for HOMO-LUMO gap prediction from 0.0791 to 0.0968, confirming the critical role of 3D structure in 3D-dependent tasks.
• Removing the 1D+3D joint denoising and translation pre-training tasks leads to performance degradation, proving their contribution to 3D-aware representation learning.
• 3D-MolT5 achieves a 100% validity rate on the ChEBI-20 dataset, indicating high chemical correctness of generated molecules.
• The model outperforms even large-scale models like GPT-4 and BioT5 across all metrics on text-based molecule generation, highlighting its efficiency and effectiveness.