VitaMol：离散空间中的轻量级分子模态扩展，用于多任务迁移学习

《Pattern Recognition》：VitaMol: Lightweight Molecular Modality Expansion in Discrete Space for Multi-Task Transfer Learning

针对大分子模型在下游任务中泛化能力不足的问题，本文提出VitaMol框架，通过级联向量量化码本注入化学功能模态，统一结构化学嵌入空间，显著提升迁移学习效果。

　　

刘云帆|高志峰|吴立荣|周静波|徐同月|谭成|黄宇飞|池长熙|余昌|Stan Z. Li

中国浙江省杭州市浙江大学

摘要

最近在3D分子构象的自监督预训练方面取得的进展已经产生了具有数十亿参数的模型，在属性预测基准测试中创下了新的最佳记录。然而，这些以结构为中心的模型在转移到特定任务时表现可能不如简单的可训练GNN，因为它们忽略了化学功能信息。为了解决这一差距，我们提出了VitaMol，这是一个基于向量量化（VQ）的轻量级多模态增强框架。VitaMol使用级联VQ码本来（1）保留基础模型中获得的丰富结构表示，（2）以最小的计算开销注入新的化学模态，以及（3）在统一的离散嵌入空间中对齐两种模态。我们使用增加了药效团特征的1300万个分子对VitaMol进行微调，教会模型同时关注功能化学线索和3D几何结构。在迁移学习设置下，我们用两个不同的骨干网络在三个多样化任务上评估了VitaMol：酶底物预测、单细胞药物扰动和分子光谱模拟，显示出比各种预训练方法更一致的改进。我们的结果证实，VitaMol能够以轻量级的方式有效结合结构和化学模态，从而提高了实际分子应用的泛化能力。所有代码可以在https://anonymous.4open.science/r/VitaMol-B9E4/readme.md中找到。

引言

近年来，利用3D分子构象的自监督预训练模型取得了显著进展[1]、[2]、[3]、[4]。通过在包含数亿个构象的越来越大的数据集上进行预训练，并将模型规模扩展到数十亿个参数，这些方法大大提高了它们的表示能力。其中最大的模型，如Uni-Mol2[5]，在广泛的下游分子属性预测任务中展示了最佳性能[6]、[7]。与基于图[8]、[9]或序列[10]的早期预训练方法相比，性能提升显著，预测准确性大大提高。这些方法和所使用的数据的惊人规模有效地重新定义了化学信息学中预测建模的标准，开启了一个高精度分子属性预测的新时代。

然而，尽管这些基于结构的预训练模型规模庞大且准确性强，但它们往往无法有效地转移到专门的下游任务中。在某些情况下，它们的表现甚至不如简单的端到端可训练的GNN。如图1所示，尽管在许多数据集上进行了预训练并且参数数量众多，但在单细胞药物扰动任务上的测试中，所有预训练模型的表现都比只有0.6M参数的简单GNN还要差。这种不足是因为它们的预训练目标几乎完全集中在捕捉3D结构模式上，而忽视了潜在的化学功能模态。然而，在许多实际应用中[11]、[12]，例如单细胞药物扰动[13]、酶底物预测[14]，对化学语义的深入理解至关重要。如果不将这些功能线索纳入预训练，模型在预测分子在复杂化学环境中的行为时就会显得不足。

将额外模态纳入现有模型的常见策略包括在新任务上对预训练模型的所有参数进行微调[15]，或者简单地连接（或以其他方式融合）从每个模态中提取的特征。然而，对预训练模型的所有参数进行微调在计算和内存使用方面都代价高昂。此外，下游任务通常只提供有限数量的分子——通常从几百到几千不等——这不足以有效地微调大型模型。另一方面，简单的特征融合方法在平衡和协调来自不同来源的信息时常常面临挑战。这种融合过程可能导致一种模态主导其他模态，或者相反，一种模态中更微妙但重要的线索可能被另一种模态的更强信号淹没。这些限制表明需要一种更高效、更轻量级的模态扩展机制：一种能够在不需要完全重新训练或依赖简单且容易失衡的融合方案的情况下，用新模态的额外信号增强基础模型。这样的机制将为将多模态信息整合到现有模型中提供更可扩展和实用的解决方案。

在这项工作中，我们直接针对三个基本挑战：（1）如何在不降级的情况下保留从原始模态中学到的丰富表示，（2）如何以轻量级、计算高效的方式引入新模态，以及（3）如何对齐两种模态，使它们的互补信号相互增强而不是冲突。过去，模态扩展技术（如特征融合和语义空间合并）大多是连续的。然而，这种连续方法存在潜在问题，包括过拟合、缺乏可解释性以及处理复杂数据结构的难度。尽管存在这些挑战，离散解决方案在很大程度上尚未得到探索。向量量化（VQ）提供了一种强大的机制，可以仅保留连续表示中最显著的特征，同时丢弃虚假噪声。通过将高维嵌入映射到从共享码本中抽取的有限离散代码集，VQ强制进行特征选择：只有跨分子重复出现的典型模式被保留，而罕见或噪声较大的变化则被有效过滤掉。此外，由于每个输入仅由最接近的代码向量表示，整个模态可以以最小的开销存储和访问——码本条目捕获了一个压缩但信息丰富的特征词汇表。

为了实现这些目标，我们引入了VitaMol¹，这是一个基于VQ的轻量级模态增强框架，旨在明确保留现有模态，以最小的开销注入新模态，并在统一的离散表示空间中对齐两种模态。如图2所示，为了弥补仅依赖结构去噪预训练任务导致的化学语义缺失，我们引入了一个级联码本，以实现结构模态的保留和化学模态在离散空间中的注入。具体来说，级联码本旨在保留训练有素的结构模态，并为拟合由药效团特征表示的化学功能提供额外的模型自由度。关键的是，模态扩展是在离散域内实现的，而不是在连续空间中，利用了码本条目的结构化粒度。接下来，我们在1300万个分子上训练VitaMol，使模型能够关注与药效团相关的特征。这些设计有效地结合了大型模型的特征，并从结构和化学模态两方面增强了全面理解，从而能够直接适应下游迁移学习场景。为了验证VitaMol，我们选择了两个预训练的骨干网络以及三种迁移学习场景：酶底物预测、单细胞药物扰动和分子光谱模拟。通过多种指标的广泛实验，证实了我们设计的有效性，显示出其在多样化下游任务中的巨大潜力。我们的主要贡献可以总结如下：

•

困境 我们发现当前的分子大型模型在下游迁移学习场景中表现较弱，因为缺乏对化学功能模态的认识，阻碍了多样化学场景的增强。

•

提取与对齐 我们提出了VitaMol框架，能够在保留训练有素的结构模态的同时，在离散表示空间中进行轻量级的化学模态增强和跨模态对齐。

•

迁移学习 我们在迁移学习设置下，使用两个骨干网络和三个基于分子的下游任务评估了VitaMol的泛化能力。广泛的结果展示了VitaMol的巨大潜力。

相关工作

相关工作

向量量化 向量量化（VQ）最初被提出作为一种压缩高维数据的方法[16]。近年来，VQ已被应用于各种领域，以压缩嵌入并获得离散表示[17]、[18]、[19]。残差VQ[20]提出分层量化数据，递归地保留更多详细信息。多头VQ[21]量化独立头以捕获数据的不同方面。随机投影VQ[22]使用随机

整体框架

计算机视觉[31]、[32]和自然语言处理[33]、[34]中的大规模模型展示了强大的涌现能力，在广泛的任务中展示了深度理解和生成能力。相比之下，分子领域尚未看到类似的突破。例如，迄今为止最大的3D构象预训练模型Uni-Mol2[5]在单细胞药物扰动[13]等专门任务上的表现仍不如简单的可训练GIN（一种GNN）[9]。这

训练后

我们从ZINC15数据库[35]下载了13,153,938个分子，通过中性电荷和具有3D构象的在库分子进行筛选。我们应用RDKit中的“ChemicalFeatures”模块和“BaseFeatures”模块从分子中提取7维化学特征。选择了两个代表性的结构去噪基础模型Uni-Mol2[5]和3D-Denoising[2]来验证我们的模态扩展范式是否可以移植到不同的骨干网络中。我们将这两个变体表示为

结论

在本文中，我们指出了当前大型分子表示模型面临的一个关键挑战：它们在特定迁移学习场景中难以有效泛化。这一限制主要源于缺乏化学功能模态，这阻碍了它们捕捉和利用重要化学信息的能力。为了解决这个问题，我们提出了一种新方法，通过添加新的化学模态来扩展模型的能力，同时仔细保留

未引用图表

图7、表5

CRediT作者贡献声明

刘云帆：撰写——原始草稿、可视化、验证、方法论。高志峰：监督。吴立荣：方法论。周静波：方法论、调查。徐同月：资源。谭成：撰写——审阅与编辑、方法论。黄宇飞：撰写——审阅与编辑、方法论。池长熙：可视化。余昌：撰写——审阅与编辑、监督。Stan Z. Li：监督。

利益冲突声明

作者声明他们没有已知的竞争财务利益或个人关系可能影响本文报告的工作。

致谢

本工作得到了国家重点研发计划（项目编号2022ZD0115101）、国家自然科学基金项目（项目编号U21A20427）和多模态人工智能系统国家重点实验室开放研究基金（项目编号MAIS2025064）的支持。