[论文解读] 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 Methods被引用 4

一句话总结

3D-MolT5 提出了一种基于 T5 的统一框架，通过使用 E3FP 指纹的新型 3D 分子分词方法，将 1D 分子序列（SELFIES）与 3D 分子结构整合。通过在原子级别对齐 1D 和 3D 分词，并在多个任务上联合预训练，该模型在分子性质预测、分子描述生成和文本到分子生成任务中均达到最先进性能，在 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 的任务至关重要。
• 通过在共享分词空间中统一 1D 和 3D 表示，克服现有 3D 结构增强语言模型中存在的模态错位和独立微调问题。
• 通过统一的预训练和指令微调框架，实现分子序列与 3D 结构信息的端到端学习。
• 通过整合细粒度的 3D 子结构表征，提升下游任务（如分子性质预测、分子描述生成和文本引导的分子生成）的性能。

提出的方法

• 提出一种基于 E3FP（扩展 3D 指纹）的 3D 分子分词方法，将 3D 分子结构转换为离散的 3D 分词，每个分词代表一个以特定原子为中心的子结构。
• 在原子级别对齐 1D SELFIES 分词与 3D 结构分词，通过求和其嵌入向量形成联合分子表征。
• 采用多任务预训练策略，包含五个目标：1D 去噪、1D+3D 联合去噪、3D 到 1D 翻译、3D 到文本翻译，以及文本到 1D 分子翻译。
• 采用基于 T5 的架构，联合编码 1D 序列、3D 分词和文本，支持指令微调以实现零样本和少样本泛化。
• 在多个下游数据集上应用指令微调，使模型适应多样化的分子生成与理解任务。
• 采用统一的分词词汇表，支持 1D 和 3D 分子表征，实现在单一模型架构内的无缝集成。

实验结果

研究问题

• RQ1 能否通过基于子结构的离散化表征，有效将 3D 分子结构分词并整合进预训练语言模型？
• RQ2 在 1D 和 3D 模态上联合预训练是否能提升依赖 3D 的分子任务（如性质预测和 3D 功能建模）的性能？
• RQ3 与仅使用 2D 或 1D 的模型相比，3D 结构信息在多大程度上提升了文本到分子生成的准确性？
• RQ4 在原子级别对齐 1D 和 3D 分词，如何影响模型根据文本描述生成化学上合理分子的能力？

主要发现

• 3D-MolT5 在 ChEBI-20 数据集上实现 48.7% 的精确匹配得分，显著优于所有基线模型，包括 GPT-4 和 BioT5。
• 在 ChEBI-20 上，模型获得 0.849 的 BLEU 得分和 0.574 的 text2mol 得分，表明其生成质量高且与参考分子高度对齐。
• 消融实验表明，若移除 3D 输入，HOMO-LUMO 能隙预测的 MAE 从 0.0791 上升至 0.0968，证实 3D 结构在 3D 依赖任务中的关键作用。
• 若移除 1D+3D 联合去噪与翻译预训练任务，性能下降，证明这些任务对 3D 敏感表征学习具有重要贡献。
• 3D-MolT5 在 ChEBI-20 数据集上实现 100% 的分子有效性率，表明生成分子具有高度化学正确性。
• 在所有指标上，3D-MolT5 均优于 GPT-4 和 BioT5 等大规模模型，凸显其高效性与有效性。