西南交大杨燕/江永全团队：基于双任务的端到端图至序列无模板反应预测模型

排版 | X

本文介绍西南交通大学杨燕/江永全团队发表于《Applied Intelligence》的研究成果，第一作者是硕士生胡昊哲。

相关代码：https://github.com/AILBC/BiG2S

作者以目前无模板逆合成领域兴起的图至序列(Graph-to-Sequence)模型框架为基础，进一步在同参数量规模下尝试构建一类在单个模型中同时解决逆合成预测与正向反应预测任务的模型 BiG2S(Bidirectional Graph-to-Sequence)。

同时，作者初步分析了模型在主流逆合成数据集 USPTO-50k 上训练时不同 SMILES 片段的预测难度差异以及模型在训练期间对验证集数据 Top-k 匹配率的波动，并对此针对性的引入了不平衡损失函数以及对模型集成与束搜索(beam search)策略的改进。

在三个主流反应预测数据集上对逆合成与正向反应预测任务的测试以及对上述模块全面的消融实验证明了 BiG2S 能够在合适的参数规模下以单一模型处理逆合成与正向反应预测任务，且整体预测能力比肩已有的基于预训练和数据增强的无模板方法。

研究背景

逆合成(Retrosynthesis)与正向合成(ForwardSynthesis)是目前有机化学、计算机辅助合成规划(CASP)以及计算机辅助药物设计(CADD)领域的基础性挑战。

其中，前者旨在为目标产物搜索可用于合成该产物的反应及对应的反应物，后者则需要为给定的反应物集合预测其发生反应后的主要产物。

早期的逆合成规划系统直接依赖于领域专家预先编码的反应规则，或者是基于物理化学的计算，而随着深度学习的快速发展。目前领域内的主流方法则是构建一个任务特异的神经网络框架以从数据驱动的角度完成反应预测任务。其中，不依赖于特定先验化学知识的无模板法通过其类似于端到端机器翻译的简洁思路以及灵活性逐渐成为了领域内的主流发展方向之一。

目前，大多数无模板逆合成模型的输入与输出均为分子的 SMILES 字符串，即采用了序列至序列(Seq2Seq)的流程。这种方法能够很好的利用在 NLP 领域内已有的模型框架，以及针对于 SMILES 表示方法的成熟的数据处理流程。

然而，SMILES 作为一维的字符串序列无法很好的表征与利用分子图所包含的二维/三维结构信息。因此，领域内逐渐出现了采用分子图代替 SMILES 作为模型输入的图至序列(Graph2Seq)方法，亦或是将分子图的额外结构信息嵌入 SMILES 的序列至序列方法；这两类方法均能很好的受益于来自分子图的丰富结构特征。

基于此，本文以新兴的图至序列方法为基础，在原基于 SMILES 的模型对逆合成与正向反应预测任务同时训练的相关探索的基准上，进一步全面的探究对此类双任务模型的构建与实验，同时也初步的探索与分析了模型在训练过程中所展现的难度不平衡以及 Top-k 匹配率波动的问题；在此基础上构建的 BiG2S 模型能够较好的处理主流数据集中的逆合成与正向反应预测任务，并在不使用数据增强的情况下取得与其他无模板逆合成模型一致的反应预测能力。

总体框架

如图 1 所示，BiG2S 整体是一个端到端的编码器-解码器结构，其中编码器端通过局部的定向消息传递图网络以及融入图结构偏置信息的全局图 Transformer 生成最终的分子图节点表征；解码器则通过标准的 Transformer 解码器以自回归式的生成目标分子的 SMILES 序列。

值得注意的是，为了同时学习逆合成与正向反应预测，解码器端的输入额外包含了不添加位置信息的双任务标签，同时解码器端的归一化层以及最终的线性层均包含有两套参数，用于分别学习逆合成任务与正向反应预测任务。

图 1：BiG2S 整体框架图

双任务训练框架

逆合成与正向反应预测作为目标相对的两个任务，它们之间存在有非常紧密的联系；因为将以产物作为输入，反应物作为目标输出的逆合成任务中的输入与目标输出互换即可转换至正向反应预测任务。

因此，部分基于 SMILES 的无模板模型已经尝试通过将逆合成与正向反应预测同时作为训练目标来提升模型对化学反应的理解，并取得了一定的成效。基于此，作者进一步尝试在图至序列的模型中融入双任务训练。

具体来说，作者基于之前已在其它方法上使用过的参数共享策略，仅在解码器的归一化层与最终的线性层内构建了任务特异的两套参数，而在其它模块中对两类任务共享一套参数，同时额外在输入的分子图节点以及解码器的初始输入序列中额外加入了双任务标签，以此在控制整体模型规模的情况下使模型能够区分两类任务并分别学习两类任务的不同数据分布。

训练与推理优化

在训练过程中，作者进一步记录并分析了模型在训练过时所反映出了两类问题。

首先，作者记录了不同 SMILES 字符在 USPTO-50k 中的出现频次以及其在训练时对应的预测准确率，如图 2 所示。在训练过程中，对于在训练集中占比分别为 0.4% 和 0.3% 的 S 与 Br，它们之间整体预测准确率的绝对差异达到了 8%。这初步表明了不同的分子结构/片段间预测的难度存在明显的差异，由此，作者通过引入不平衡损失函数(如 Focal Loss)来缓解此类问题，从而使模型能够更加关注训练时准确率更低的分子片段。

图 2：USPTO-50k 训练集中不同SMILES字符的出现频次以及其在训练时的整体预测准确率

此外，作者进一步记录了模型在训练期间于验证集中的预测结果质量变化，如图 3 所示。作者发现在 USPTO-50k 的中后期训练期间，模型整体在验证集上的 Top-1 准确率仍然呈现一定的上升趋势，但在 Top-3,5,10 的预测质量上存在有明显的下滑。

为了在提升模型 Top-1 预测质量的同时保持模型前十位反应物生成结果的整体质量，作者额外构建了一类基于自定义评价指标的模型集成策略。具体来说，作者构建了一类存储模型的队列，同时依据预定义的评价指标(如 Top-1 准确率，加权的 Top-k 准确率等)对存入的模型进行排序；由此在整个训练过程中动态的存入待选模型并自动生成基于队列中前 3-5 位的集成模型，从而保留 Top-k 预测质量最高的模型。在推理阶段，作者也基于新的框架重新构建了更加注重于搜索广度的束搜索策略以提升模型 Top-k 生成结果的整体质量。

图3：在 USPTO-50k 数据集中训练期间模型对验证集的 Top-k 匹配率以及 Top-k 无效生成分子比例的变化曲线

基准数据集双任务实验

作者在分别包含 5 万、50 万以及 100 万条化学反应的数据集 USPTO-50k、USPTO-MIT、USPTO-full 中分别测试了双任务模型以及单任务模型在逆合成任务与正向反应预测任务中的表现，测试结果如图4所示。

可以发现，在小规模数据集中，基于双任务训练的 BiG2S 在逆合成任务中取得了无模板逆合成模型中领先的预测精度，同时也保持了较高的正向反应预测精度；而在偏向于正向反应预测的 USPTO-MIT 数据集以及大规模数据集 USPTO-full 中。由于模型整体参数量的限制，导致引入双任务训练后的模型在更大规模数据集中的表现出现了降低。然而，从双任务模型以几乎一致的参数量与小幅度的反应预测能力降低( Top-k 准确率的绝对差值位于 0.5% 左右)获得了同时处理逆合成任务与正向反应预测任务的能力这一角度来看，BiG2S 模型已经达到了预期目标。

图 4：BiG2S 的双任务模型与单任务模型在三个基准数据集上的实验结果，其中上标 b 表示采用单任务模型分别完成两类任务

消融实验分析

作者同时进一步通过消融实验验证了新的束搜索算法以及采用不平衡损失后 BiG2S 在不同数据集中进行预测的最适温度超参数。这里的温度超参数指 Softmax 中用于控制输出概率分布的温度参数 T，具体实验结果分别如图 5 与图 6 所示。

在针对束搜索算法的实验中，可以发现 OpenNMT 在搜索宽度扩大至 3 倍的同时搜索耗时仅扩大至 1.74 倍，而新束搜索算法在 Top-1 精度与 OpenNMT一致的情况下整体的搜索耗时扩大了 1-2 倍；但在 Top-10 预测结果的质量上，新的束搜索算法与 OpenNMT 相比具有至少 3% 的绝对精度优势以及 2% 的有效分子比例优势，可以说新的束搜索算法以搜索耗时为带来明显提升了模型整体Top-k搜索结果的质量。

而在针对温度超参数的实验中，作者发现在小规模数据集上使用较大的温度参数可以明显提升整体的 Top-k 预测精度，而在更大规模的数据集中，由于 BiG2S 的模型规模无法完全的拟合所有的反应数据，此时选用更小的温度参数往往有利于模型的搜索。

图 5：BiG2S 采用的束搜索算法与之前无模板逆合成模型常用的基于 OpenNMT 的束搜索算法在搜索结果质量与搜索耗时上的比较

图 6：使用不平衡损失后的 BiG2S 在不同数据集上采用不同温度超参数(T)进行搜索的 Top-k 预测准确率

结论

本文中，作者提出了一类同时处理逆合成任务与正向反应预测任务的无模板反应预测模型 BiG2S。基于合适的参数共享策略与额外的双任务标签，BiG2S 能够以较小的参数量在不同规模的数据集上以单一模型完成逆合成任务与反应预测任务，且整体预测能力与主流模型达到了同一水平。

而针对模型训练时所反映出的不同 SMILES 字符预测难度不均衡以及 Top-k 预测精度波动的问题，作者额外引入了不平衡损失，基于自定义评价指标的模型自动集成策略，以及基于新框架的束搜索算法以缓解这两类问题。

最终，BiG2S 在三个不同规模的主流数据集上均表现出了较好的双任务预测能力，而进一步的消融实验也证明了额外引入的训练与推理策略的有效性。

人工智能×[ 生物 神经科学 数学 物理 化学 材料 ]