研究背景
- 分子表示学习的重要性:在药物发现、材料科学和化学工程等领域,分子性质预测是一个关键任务,而准确的分子表示对于预测分子性质和设计新分子至关重要。现有方法的局限性:传统的分子表示方法主要基于分子指纹和描述符等特征驱动方式,虽能快速获取分子的物理、化学或拓扑特征,但难以捕捉分子结构中的复杂模式和关系。近年来,基于图神经网络(GNNs)和图变换器(GTs)的数据驱动方法在学习分子结构表示方面取得了进展,但大多忽略了分子结构与电子信息之间的关系以及分子内部的语义推理,导致分子表示不完整,缺乏对结构和电子数据的整合.
研究方法
- MOL-Mamba框架的提出:为解决现有方法的局限性,作者提出了MOL-Mamba框架,通过结合结构和电子见解来增强分子表示学习。该框架包括用于层次结构推理的原子和片段Mamba图(MG)以及用于整合分子结构和电子相关学习的Mamba变换器(MT)融合器。
图1展示了MOL-Mamba框架的预训练结构,包含三个核心模块:片段级图GNNF、原子级结构推理的Mamba-Graph(MG)以及分子结构与电子关联学习的Mamba-Transformer(MT)融合器。在预训练过程中,先通过结构分布协作训练提升分子结构学习效果,再利用电子语义融合训练整合分子结构与电子数据,以增强分子表示学习能力。
- 双图协作学习:
- 片段图构建:使用主子图挖掘(PSM)算法将分子结构分解为多个片段子图,构建片段级图GF,以捕捉分子内大组分之间的高阶连通性。原子级Mamba图构建:为增强模型对分子图中结构上下文的归纳关联能力,结合GNN和Mamba提出Mamba图(MG)。使用GNN层初始化原子图节点嵌入,引入新的图节点排序策略以适应Mamba的上下文感知推理,并添加位置编码以增强原子位置感知能力。此外,设计新的图状态空间模型(GSSM)机制,使模型能够从图上下文中自适应地选择相关信息。
- 分子电子视图表达:选择具有预测相关性的分子描述符来表征分子的整体性质,形成分子的电子视图表达DE。统一的Mamba变换器融合器:采用Mamba变换器(MT)作为骨干网络,将分子的结构和电子表示进行融合。通过引入自监督掩蔽预测任务,增强结构和电子特征之间的相互作用,并去除下游任务预测中冗余的电子描述符的影响。
图2展示了Mamba-Graph(MG)模块的结构和工作流程,包括用于初始化节点嵌入的GNN层、处理序列数据的Mamba块,以及结合图结构信息和Mamba模型上下文推理能力的新图状态空间模型(GSSM)机制,通过图节点排序策略和位置编码增强分子结构信息的提取和处理。
实验
- 数据集:使用GEOM数据集进行分子预训练,该数据集包含50k个合格分子。在下游任务中,对来自MoleculeNet的11个基准数据集进行实验,涵盖物理化学、生物物理、生理学和量子力学等领域,包括7个分类基准和4个回归基准。基线方法:与多个监督和自监督/预训练基线方法进行比较,包括SchNet、GIN、AttentiveFP、DMPNN等监督GNN方法,以及PretrainGNN、GROVER、GEM、GraphMVP、MolCLR、Uni-Mol和MOLEBLEND等预训练方法。性能评估:
- 整体比较:MOL-Mamba框架在11个分类和回归基准中的8个上表现优于其他方法。例如,在BBBP数据集上达到了75.0%的ROC-AUC,在ESOL数据集上达到了最低的平均MAE 0.63。与监督方法相比,MOL-Mamba在多个数据集上展现出更优越的性能。模型复杂度比较:MOL-Mamba在参数量和运行时间上具有显著的低成本优势,与GNNs(如GIN和SchNet)相比运行时间相近,但比GTs(如Uni-Mol和MOLEBLEND)更高效,平衡了复杂度和性能。主模块分析:通过消融实验验证了MOL-Mamba的各个模块的重要性。去除片段图GNNF模块、Mamba图(MG)模块或MT融合器模块都会导致性能下降,说明每个模块对MOL-Mamba的鲁棒性性能都有独特贡献。自监督损失分析:不同的预训练损失设置对MOL-Mamba的性能有显著影响。实验表明,引入结构分布协作训练损失(Ld)、片段相关结构损失(Ls和Lf)以及掩蔽E语义融合训练损失(Lmask)都能提升模型性能,完整模型整合所有损失组件时表现最佳。Mamba图模块分析:Mamba图(MG)模块的不同设置对性能有重要影响。去除图节点排序策略(SORT)、位置编码(PE)或图状态空间模型(GSSM)都会导致性能下降,凸显了这些组件对MG模块准确预测分子性质的关键作用。特征分离分析:通过t-SNE可视化展示了不同预训练方法在BBBP和ESOL基准上的特征分离效果。MOL-Mamba的“结构PT”和“E语义PT”预训练方法能显著降低错误率(ER),提高正负样本之间的分离度,表明模型能更好地捕捉复杂语义细节。特征可视化分析:可视化特征权重显示了模型决策过程中各输入特征的相对重要性。MG和MT模块对特征的转换表明模型能够关注关键的分子区域,如反应中心和立体化学,从而提高预测的准确性和可靠性。
总结
本文提出了MOL-Mamba框架,旨在通过结合分子结构和电子特性来增强分子表示学习。框架包含用于层次结构推理的原子和片段Mamba图(MG)以及用于整合分子结构和电子相关学习的Mamba变换器(MT)融合器。通过双图协作学习和电子语义驱动的融合策略,MOL-Mamba在多个分子性质预测基准数据集上取得了优于现有方法的性能,为药物发现和材料科学等领域的分子性质预测提供了有力工具。