[论文解读] A Large Encoder-Decoder Family of Foundation Models For Chemical Language

Eduardo Soares, Victor Shirasuna+4 more|arXiv (Cornell University)|Jul 24, 2024

Advanced Computational Techniques and Applications被引用 4

一句话总结

该论文介绍了 SMI-TED289M，一个在 PubChem 中 9100 万条精心筛选的 SMILES（对应 40 亿个分子标记）上预训练的大规模编码器-解码器基础模型，在分子性质预测和分子重构任务中达到最先进性能。该模型在潜在空间中展现出强大的组合结构，支持基于分子片段之间线性关系的少样本推理，并在生成任务中表现出较高的 Tanimoto 相似度。

ABSTRACT

Large-scale pre-training methodologies for chemical language models represent a breakthrough in cheminformatics. These methods excel in tasks such as property prediction and molecule generation by learning contextualized representations of input tokens through self-supervised learning on large unlabeled corpora. Typically, this involves pre-training on unlabeled data followed by fine-tuning on specific tasks, reducing dependence on annotated datasets and broadening chemical language representation understanding. This paper introduces a large encoder-decoder chemical foundation models pre-trained on a curated dataset of 91 million SMILES samples sourced from PubChem, which is equivalent to 4 billion of molecular tokens. The proposed foundation model supports different complex tasks, including quantum property prediction, and offer flexibility with two main variants (289M and $8 imes289M$). Our experiments across multiple benchmark datasets validate the capacity of the proposed model in providing state-of-the-art results for different tasks. We also provide a preliminary assessment of the compositionality of the embedding space as a prerequisite for the reasoning tasks. We demonstrate that the produced latent space is separable compared to the state-of-the-art with few-shot learning capabilities.

研究动机与目标

• 开发一种可扩展、开源的基础模型，用于化学语言表征，通过在大规模未标注分子数据上进行自监督预训练。
• 通过编码器-解码器 Transformer 架构学习上下文感知、可组合的表征，以改进分子性质预测与生成。
• 探究所学分子嵌入的组合结构，作为化学推理与少样本泛化的前提条件。
• 为药物发现和材料科学等多样化科研应用，提供一个灵活且高性能的模型系列（289M 和 8×289M 参数变体）。
• 发布一个包含 9100 万条高质量 SMILES 的精心筛选数据集，以供未来化学基础模型的预训练使用。

提出的方法

• 在 PubChem 中的 9100 万条精心筛选 SMILES 上预训练基于 Transformer 的编码器-解码器模型，相当于 40 亿个分子标记，采用自监督掩码语言建模方法。
• 采用双编码器机制，共享词汇表并使用分层分词策略，通过 SMILES 字符串建模分子结构与连接性。
• 在较大变体（SMI-TED8x289M）中引入专家混合（MoE）架构，实现稀疏激活下的参数高效扩展。
• 使用 t-SNE 可视化和嵌入投影上的线性回归，量化潜在空间中的组合结构，其中三元组 SMILES 表示片段加和关系。
• 通过从少量示例三元组生成嵌入并利用 Tanimoto 相似度重建目标 SMILES，评估少样本推理能力。
• 应用标准基准测试（QM9、MOSES）以及回归/分类任务，验证模型在量子、物理、生物物理和生理性质预测方面的性能。

实验结果

研究问题

• RQ1 大规模编码器-解码器基础模型能否通过在未标注 SMILES 上进行自监督预训练，学习到组合型分子表征？
• RQ2 SMI-TED289M 所学得的潜在空间在少样本推理与分子重构方面表现如何？
• RQ3 该模型的嵌入空间在多大程度上反映了分子片段之间的分层与线性关系？
• RQ4 在精心筛选的大规模 SMILES 语料上进行预训练，是否能相比先前方法显著提升分子性质预测性能？
• RQ5 该模型能否在极少微调的情况下泛化至多样化化学任务，展现出强大的少样本能力？

主要发现

• SMI-TED289M 在 QM9 基准测试中 12 项任务中的 11 项达到最先进性能，用于预测量子力学性质。
• 在片段加和任务的线性回归中，R² 达到 0.99，MSE 为 0.002，表明潜在空间具有强大的组合结构。
• 在少样本生成任务中，模型的平均 Tanimoto 相似度得分为 0.52，最佳结果达到 0.92，对应正确 SMILES 的重构。
• SMI-TED289M 嵌入的 t-SNE 可视化显示，按分子家族和层级距离清晰分离，优于 MoLFormer（R²=0.55，MSE=0.237）。
• 模型展现出稳健的少样本泛化能力，即使仅使用少量三元组作为输入生成嵌入，也能获得较高的相似度得分。
• 从 PubChem 筛选的 9100 万条 SMILES 数据集被证明对预训练有效，显著提升了模型性能与组合推理能力。