基于BERT的化学空间映射

背景介绍与方法概述

化学反应分类可以帮助我们探索更复杂的化学概念，但化学反应的分类任务是困难的。在专家的角度看，化学反应分类需要通过标注反应中的分子数量，反应中心以及反应物和试剂之间的区别来识别相应的反应类别模板。IBM和伯尔尼大学的研究人员提出基于Transformer的模型，可以直接从化学反应的文本表示中推断反应类别。他们的工作表明，学习化学反应分类的表达可以捕获分子在化学反应空间下的细粒度差异，这有助于我们深入了解化学反应空间。

Related Work

在现有的方法里，使用SMILES结合NLP或者其他文本表示分子和化学反应的方法在化学领域十分有效。近期大量的实验证明，神经网络可以在化学反应中捕捉从前体（precursors）到产物（products）的原子重排（atom rearrangements），并且无需监督。

下图a展示了化学反应的例子，对应的文本表示如图b所示。图c表示：反应类别可以根据反应中的变化结构特征进行命名，比如"从硝基到氨基（nitro to amino）"，这也被作为反应的模板或者中心（Reaction template or center）。或者，可以根据发现该反应的人名来命名，比如图c中的第二个例子，这是一个 chloro Suzuki coupling reaction，我们以 Akari Suzuki 命名，Akari Suzuki 因为在钯催化交叉偶联反应方面的工作获得2010年诺贝尔奖。

最近，Schneider等人开发了一种基于传统反应指纹（reaction fingerprints）的反应分类器。分子和反应指纹（Molecular and reaction fingerprints）是离散分子结构和化学反应的固定长度的向量编码。Schneider等人开发的指纹图谱都需要确保反应物与试剂（reactants-reagents）两者分离，这限制了其适用性。

传统上，反应指纹是使用反应中心或反应物的结合，试剂和产物指纹的组合手工制作的。在Duvenaud等人提出的可微分子指纹图谱的基础上，Wei等人提出了第一个学习的反应指纹例子，并用于预测化学反应。不幸的是，他们的指纹只限于一个固定的反应方案：两个反应物和一个试剂，因此只适用于符合该方案的反应。同样的，Sandfort提出的多重指纹特征是通过连接固定数量的分子的多重指纹而得到的，也只能限于固定的反应方案。

方法概述

在方法中的第一部分，基于BERT预测化学反应类别，该深度学习模型不依赖于定制好的专家规则（这些规则要求每个反应被适当的原子映射）。相反，模型学习原子的动机（在没有反应物-试剂作用注释的情况下，从原始反应中区分不同类别的反应），如上图d。

实验中，基于Transformer的Seq2Seq模型匹配Groud Truth的精度为95.2%，基于BERT的分类器的分类精度达到98.2%。进一步，作者分析了Seq2Seq模型的注意力，以及基于BERT的分类器的注意力，观察到，参与反应中心的原子，以及特定于反应类别的试剂，具有更大的注意力权重。

在方法的第二部分，作者证明了BERT学习到的表示，无监督和有监督，都可以作为反应指纹。经过化学反应数据训练的BERT模型可以将任何化学反应的smile文本表示转换为embedding编码，而不需要进行传统方法比如原子映射或反应物-试剂分离操作。该方法的反应指纹因此普遍适用于任何化学反应数据库。

结果与讨论

化学反应分类

论文使用一组具有标签的化学反应数据集训练两个模型，一个是基于Transformer的Seq2Seq模型，另一个是基于BERT编码的模型。对于第二种模型，在输出端有一个分类器，通常在NLP中用于单个句子的分类任务，这种基于BERT的模型如下图所示。

论文提出一个基于BERT的模型，利用分类器将模型输出的上下文表达用于化学反应分类任务。分词[CLS]对应的编码可以用作反应指纹的表达（RXNFP）。

Ground Truth数据由以文本格式（SMILES）表示的化学转换（chemical transformations）组成。这些标签来自一个很不平衡的数据集：Pistachio dataset，该数据集本来是使用NameRxn进行化学反应分类的。在另一个实验中，作者使用一个叫做 USPTO 1k TPL 的开源数据集导出反应标签。论文分析了模型在测试集上的分类性能，其中包含了来自Pistachio的792个不同类别的13.2万个反应，以及来自USPTO 1k TPL的1000个类别的4.5万个反应。下表中提供了对结果的总结。

在上表中，有两种模型的效果很好，一种是基于BERT的classifier直接对上下文表达（Contextual representations）进行分类，另一种是从BERT提取反应指纹RXNFP输入K近邻（令K=5）进行分类。

---

RXNFP(pretrained)代表传统手工定义的反应指纹，RXNFP代表使用BERT得到的反应指纹。

---

表中结果表明，基于BERT，我们可以得到具备明显细粒度差异的反应指纹表达，从而有利于在化学反应的相关模式识别任务中带来效果提升。

可视化注意力分布

下图展示了 layer-wise 的注意力分布。横轴代表化学反应的SMILES标记。颜色越深，意味着当前处理层下，SMILES序列中对应的原子就被关注越多。

可以发现，具有较大注意力的原子是反应中心的一部分，或与特定于反应类别的前体相关。就像人们期望在分类的基础上看到一组特定的原子一样。在Seq2Seq模型中，解码器学会了关注重排所涉及的原子来对反应进行分类。BERT分类器中，初始层对反应中的所有原子关注较弱。中间层倾向于产物。最后一层重点关注了对分类很重要的与反应中心相关的前体。

化学空间映射（Mapping the chemical reaction space）

模型可以得到数据驱动的反应指纹。分子指纹被广泛应用于筛选具有相似性质的分子或绘制化学空间。在BERT的训练中，对反应SMILES中的个体标记进行掩盖，然后由模型进行预测。由于预先的[CLS]令牌（token）不会被掩盖，模型总是能够加入该令牌，以恢复那些被屏蔽的令牌。另外来看，模型使用[CLS]令牌可以嵌入对反应的全局描述。在微调之前，[CLS]令牌嵌入完全是通过自我监督学习的。对于有监督的微调，[CLS]的embedding被作为分类器的输入，并进一步细化。我们将经过Pistachio训练集微调的指纹称为RXNFP。在论文中，[CLS]令牌embedding是一个大小为256的向量，对应于BERT模型的hidden大小。

反应指纹的另一个主要用例是在数据库中搜索类似的反应。一个独立于原子映射的反应指纹是极其强大的，因为它提高了不需要知道反应中心的反应检索的可能性。例如，当像正向反应预测模型或逆向合成模型这样的黑箱模型预测一个反应时，可以从这些模型的训练集中检索出最相似的反应。这种相似反应的检索不仅可以增加深度学习模型的可解释性。它还允许化学家访问最接近的反应的元数据(包括产率和反应条件)。