Nature Machine Intelligence | “化学元素知识+功能提示”双驱动，探索分子预测新方法
论文题目：Knowledge graph-enhanced molecular contrastive learning with functional prompt
论文链接：https://doi.org/10.1038/s42256-023-00654-0
项目地址：GitHub - HICAI-ZJU/KANO: Code and data for the Nature Machine Intelligence paper "Knowledge graph-enhanced molecular contrastive learning with functional prompt".Code and data for the Nature Machine Intelligence paper "Knowledge graph-enhanced molecular contrastive learning with functional prompt". - GitHub - HICAI-ZJU/KANO: Code and data for the Nature Machine Intelligence paper "Knowledge graph-enhanced molecular contrastive learning with functional prompt".
https://github.com/HICAI-ZJU/KANO
发表期刊：Nature Machine Intelligence
引言
借助 AI的分子性质预测模型，可以评估候选药物的临床试验成功率和治疗潜力，极大地提升了药物研发的效率，能够显著降低因后期失败带来的经济损失。当前，大多数基于分子的自监督学习方法由数据驱动，通常仅考虑原子之间通过化学键建立的联系，未能充分挖掘分子图中原子间的潜在关联，从而使其结果缺乏可解释性。此外，预训练任务往往与下游任务有很大差异，这意味着直接将预训练得到的表示应用于下游任务，可能无法取得理想的性能表现。
因此，本文提供了一种有力的辅助药物发现的分子性质预测工具。通过融合化学元素知识图谱中的基础化学信息作为先验知识，激发模型在理解分子语义和拓扑结构方面的潜在能力。
方法
图1: KANO模型总览
图2: ElementKG概览
首先，本文基于元素周期表和维基百科，构建了一个“化学元素知识图谱”（ElementKG）。ElementKG 从化学元素的角度呈现了一个全面且标准化的视图，覆盖了元素的类层次结构、化学属性、元素间的关系、相关官能团及官能团与其组成元素间的连接。
在此基础上，本文提出了一种基于功能提示的知识图谱增强分子预训练方法。在预训练阶段，针对分子中存在的元素和官能团类型，从ElementKG中检索相应的实体和关系，创建一个增强的分子图。该分子图融合了基础领域知识，并捕获了拥有相同元素类型的原子间的基本关系，即使它们并未通过化学键直接相连。此方法用于构建对比学习的正样本对，并通过预训练获得分子的向量表示。该表示在保留拓扑结构的同时，整合了关键的化学语义。
在预训练后，为了弥补预训练对比任务和下游任务之间的差距，本文提出使用官能团知识作为提示，以激励预训练的图编码器，并从ElementKG的官能团知识中生成功能提示。首先，检测输入分子中的所有官能团，并检索其在ElementKG中的对应实体嵌入，构建可学习的中介变量，以捕捉每个官能团的重要性。接下来，将自注意力机制应用于中介变量和官能团实体的嵌入，以全面聚合其语义并获得功能提示。该方法有效提升了分子属性预测的准确性，同时确保了一定程度的可解释性。
实验
本文在14个分子性质预测的基准测试中对KANO进行了全面的评估，展现了其优越的性能。
图3: KANO在14个基准数据集上获得了优越的性能
为深入评价KANO在表示空间质量上的表现，研究主要聚焦于对齐性和均匀性两大纬度，并与现有最具代表性的有监督模型、基于预测的预训练模型及基于对比的预训练模型进行对比分析。实验结果揭示，KANO能够将拥有相同主要构造（scaffold）的分子映射到相近的表示空间，并确保所得表示向量在单位球面上均匀分布，进而保留更多的数据信息。
图4: KANO通过元素知识增强的预训练获得高质量的分子表示空间
进一步地，通过可视化官能团的注意力权重，本文探讨了KANO的预测可解释性。可解释性分析表明，KANO能够在分子中识别出具有特殊化学意义的子结构，并通过激活与下游任务密切相关的知识，来实现预训练目标与下游性质预测目标之间的对齐。
图5: KANO通过官能团提示指导的微调感知具有特殊化学意义的子结构
总结
本文提出了KANO，一种通过结合化学领域的专业知识来增进分子属性预测任务新策略。在知识图谱的指导下进行的预训练使KANO获得了高质量的分子表征空间，同时，官能团提示捕捉了与下游任务相关的、富有化学意义的子结构信息。当然，KANO也存在一些局限性。例如，ElementKG可能无法充分体现分子系统的复杂性，现有的官能团提示可能无法捕捉子结构之间的远程相互作用。面对这些挑战，未来的研究方向可能包括：首先，可以将ElementKG扩展到其他化学领域或与其他现有知识图谱融合，以提供更全面、更系统的理解分子的途径。其次，探讨KANO所学习的表征的可解释性和官能团提示捕捉的化学知识，为分子设计和优化提供深入的见解。最后，探讨将KANO与其他技术整合的方法，以增强其在小数据集上的性能，并进一步加速药物的发现过程。