【深度学习】RetroKNN：kNN增强的化学逆合成预测模型

——背景——

化学逆合成分析是药物发现过程中的关键环节，而深度学习方法的应用极大提高了逆合成预测的准确性。不过，深度学习方法会受到数据本身的限制。对基于反应模板（template-based）的深度学习逆合成预测来说，其预测任务是从待分解产物分子到特定反应模板的分类任务，并通过监督学习进行模型训练。然而，和其他分类任务类似，化学反应数据中不同的反应模板所对应的反应数量不同，同样存在数据不平衡问题。如何实现few-shot甚至zero-shot的反应模板准确预测，是基于反应模板的逆合成预测需要解决的关键问题之一。

最近，中国人民大学高瓴人工智能学院与微软亚洲研究院合作，提出了通过kNN方法增强的逆合成预测模型RetroKNN，在few-shot和zero-shot情景下改善深度学习模型的预测能力。相关工作目前发表在arXiv中¹。

——方法和结果——

LocalTemplate与LocalRetro模型架构

作者在本文中使用kNN方法对Chen等人提出的LocalRetro架构²进行增强。LocalRetro是目前在逆合成预测领域常用数据集USPTO-50K上表现最佳的模型，其主要特点在于通过局域反应模板LocalTemplate捕捉化学反应的局域性特征。大多数化学反应总是发生在分子中小部分原子和化学键上，也就是反应中心（reaction center）上。具体来说，反应中心仅涉及原子的LocalTemplate被称为atom template，如果涉及化学键则称为bond template。如图1所示，LocalRetro首先通过一个MPNN架构得到输入分子图各节点的隐藏表示，然后用单层全连接层从一对相连节点的隐藏表示获得连接它们的边的隐藏表示，全部节点和边的隐藏表示经过一个全局注意力层（Global Reactivity Attention, GRA）后再通过全连接层，得到各节点和边应用各LocalTemplate的概率。

图1 LocalRetro模型架构

通过kNN增强LocalRetro预测能力

图2 RetroKNN架构

当所需预测的LocalTemplate在训练数据中出现频率较低时，预测的准确性将会下降。为解决这一问题，作者提出了RetroKNN架构（图2）。作者用训练好的LocalRetro模型的MPNN部分（图1红色框内GRA前的部分）作为Feature Extraction GNN计算分子图节点和边隐藏表示h_v和h_e，并与分子对应的atom/bond template中的反应中心进行比对，将属于反应中心的隐藏表示归类为对应的LocalTemplate，不属于反应中心的隐藏表示标记为0，即不对应任何LocalTemplate。这样就构建了用kNN对节点和边所属LocalTemplate进行分类的数据集atom/bond store。kNN可以通过输入的新分子图上节点和边的隐藏表示h_a和h_b预测其所属的LocalTemplate，预测的概率与h_a和h_b到数据集中其他隐藏表示h_i的距离d有关，如公式(1)(2)所示。公式中N_a和N_b分别表示atom store和bond store数据集；t_a和t_b分别表示atom template和bond template； 和 仅当t_a=t_i和t_b=t_i时为1，否则为0。

最终，预测的概率由LocalRetro与kNN的预测概率经过公式(3)(4)加权计算得到。公式(1)(2)中的两个温度参数T_A、T_B以及公式(3)(4)中LocalRetro预测概率的权重λ_A、λ_B通过一个神经网络（图2中Adapter）由输入分子图节点和边的隐藏表示h_v、h_e和h_v、h_e到kNN数据集中数据点的距离计算得到。

实验结果

作者首先在逆合成分析中常用的数据集USPTO-50K和USPTO-MIT上对比了RetroKNN与其他逆合成预测模型的表现。可以看出在top-1直至top-50准确率上，RetroKNN均达到了最高值，并在top-1准确率上相对于原始的LocalRetro模型提升明显。这充分证明kNN可以有效增强LocalRetro模型的逆合成预测准确性。

图3 不同逆合成预测模型在USPTO-50K数据集上未知反应类别时的top-k准确率。TPL.一列，实心圆表示基于模板的方法（template-based），半实心圆表示部分基于模板的方法（semi-template），空心圆表示不基于模板的方法（template-free）。

图4 不同逆合成预测模型在USPTO-50K数据集上已知反应类别时的top-k准确率。TPL.列图例同图3。

图5 不同逆合成预测模型在USPTO-MIT数据集上的top-k准确率。TPL.列图例同图3。

接下来，作者检验了kNN方法对于分子图中节点和边所对应LocalTemplate进行分类的合理性。如图所示，作者在两个测试分子中分别选择了一根化学键和一个原子，找出在kNN数据集中与它们的隐藏表示最接近的化学键或原子，观察化学键或原子在分子中所处的化学环境。可以发现，在图中第一行，所有查找到的化学键都位于三氟乙酰基和芳环上的氨基之间，与测试分子中一致；在图中第二行，所有原子均为酚羟基，与测试分子中一致。这说明kNN方法能通过LocalRetro给出的隐藏表示有效识别化学环境接近的原子和化学键。

图5 kNN通过类似的隐藏表示预测化学环境相近的原子和化学键

作者最后在few-shot和zero-shot情景下测试了RetroKNN的预测能力。作者从USPTO-50K构建的训练/验证集中删除全部反应分类在6-10中的反应，构建了zero-shot训练/验证集；对于few-shot下的训练/验证集，作者仅保留反应分类在6-10中反应的10%。作者测试了RetroKNN和LocalRetro对反应分类在6-10中的产物分子的预测准确率。作者发现zero-shot对于基于模板的预测方法非常有挑战性，且RetroKNN的预测能力显著优于作为baseline的LocalRetro，但在few-shot情景下，分类任务会容易很多，不过kNN方法也能略微提高LocalRetro的预测准确性。这表明RetroKNN可以解决基于模板的逆合成预测中few-shot和zero-shot反应模板的预测问题，具有一定实际意义。

图6 zero-shot和few-shot下RetroKNN与LocalRetro的预测效果比较。(a)(b)zero-shot & few-shot top-5准确率；(c)(d) zero-shot & few-shot top-10准确率

——小结——

本文介绍了一种通过kNN方法增强的GNN逆合成预测模型RetroKNN，其通过kNN方法分析LocalRetro计算得到的分子图中原子和化学键隐藏表示的相似性，可以在zero-shot和few-shot反应模板预测情境下给出更准确的结果。这种思路可以推广到其他可应用GNN解决的问题中以缓解数据不平衡问题，如正向反应预测、反应选择性预测等等，同时对基于GNN逆合成预测模型中隐藏表示的含义提供了一种直观的解释。

参考文献：

1. Shufang Xieet al. “Retrosynthesis Prediction with Local Template Retrieval.” arXiv (2023) DOI: 10.48550/arXiv.2306.04123

2. Shuan Chen and Yousung Jung. “Deep Retrosynthetic Reaction Prediction using Local Reactivity and Global Attention.”JACS Au 1.10 (2021): 1612-1620. DOI: 10.1021/jacsau.1c00246

---

作者：郭家盛

审稿：黄志贤

编辑：黄志贤

GoDesign

ID：Molecular_Design_Lab