排版 | X
本文介绍西南交通大学杨燕/江永全团队发表于《Applied Intelligence》的研究成果,第一作家是硕士生胡昊哲。
论文链接:https://doi.org/10.1007/s10489-023-05048-8
相关代码:https://github.com/AILBC/BiG2S
作家以目前无模板逆分解领域兴起的图至序列(Graph-to-Sequence)模型框架为基础,进一步在同参数量规模下尝试建立一类在单个模型中同时解决逆分解猜测与正向反响猜测工作的模型 BiG2S(Bidirectional Graph-to-Sequence)。
同时,作家初步分析了模型在支流逆分解数据集 USPTO-50k 上训练时分歧 SMILES 片段的猜测难度差异以及模型在训练期间对验证集数据 Top-k 匹配率的波动,并对此针对性的引入了不平衡损失函数以及对模型集成与束搜寻(beam search)策略的改进。
在三个支流反响猜测数据集上对逆分解与正向反响猜测工作的测试以及对上述模块全面的消融尝试证明了 BiG2S 能够在合适的参数规模下以单一模型处理逆分解与正向反响猜测工作,且完全猜测能力比肩已有的鉴于预训练和数据增强的无模板方法。
研究背景
逆分解(Retrosynthesis)与正向分解(Forward Synthesis)是目前有机化学、计算机辅助分解规划(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 在三个分歧规模的支流数据集上均表现出了较好的双工作猜测能力,而进一步的消融尝试也证明了额定引入的训练与推理策略的有效性。