Nat. Comm. | 语言模型可以学习复杂的分子分布

编译 | 吴庆文 审稿 | 苏延森

本文介绍由加拿大多伦多大学的Daniel Flam-Shepherd和Alán Aspuru-Guzik共同通讯发表在Nature Communications的研究成果：作者研究了语言模型学习复杂的分子分布的能力。通过编译更大、更复杂的分子分布，作者引入几个挑战性的分子生成任务评估语言模型的学习能力。结果表明，语言模型具有强大的生成能力，能够学习复杂的分子分布。语言模型可以准确生成：ZINC15数据集中惩罚 LogP得分最高分子的分布、PubChem数据集中多模态分子及最大分子的分布。

1

简介

针对化学空间的有效探索对于药物和材料发现有着重要作用，然而目前只探索了可合成化学空间的一小部分，开发新工具至关重要。最近，深度生成模型已成为有效探索化学空间最有前景的工具之一。这些模型在化学空间的相关子集上进行训练，可以生成与训练数据相似的新分子。其中语言模型循环神经网络 (RNN)以SMILES（Simplified Molecular Input Line Entry Specification）字符串的方式生成分子；其它一些模型或以图的方式生成分子（图生成模型），或将分子生成为3D空间中的点云。深度生成模型学习训练集的分布和生成有效的相似分子的能力对下游应用非常重要。

最初由于SMILES字符串表示的脆弱性，导致RNN模型经常生成无效分子。随后，研究人员使用鲁棒的分子字符串表示SELFIES（SELF-referencIng Embedded Strings），或者改进训练方法，RNN模型也能够持续生成高比例的有效分子。目前尚未有研究将语言模型用于生成更大更复杂的分子，或者从大规模化学空间进行生成。为了测试语言模型的生成能力，作者通过构建比标准的分子数据集更复杂的训练集来制定一系列挑战性的生成任务。作者在所有任务上训练语言模型，并和图生成模型进行比较。结果表明，语言模型的生成能力更强，比图生成模型能更好地学习复杂的分子分布。

2

结果

作者定义了三个生成任务：（1）具有高惩罚LogP值的分子分布（图 1a，d），（2）分子的多模态分布（图 1b，e），和（3）PubChem中最大的分子（图 1c，f）。

图1 生成建模任务

对比标准数据集Zinc和Moses，表1统计了所有数据集中的原子数和环数。所有任务都涉及具有更多子结构的更大分子，并且每个分子包含更大范围的原子数和环数。

表1原子数和环数的统计信息

对于每项任务，作者通过绘制训练分子属性的分布以及语言模型和图生成模型学习的分布来评估生成模型的性能。作者对训练集和生成样本的分布使用直方图表示，并通过调整其带宽参数来拟合高斯核密度估计（KDE）。所有模型生成1万个分子，计算它们的属性并生成所有绘图和指标。为了公平比较，在删除重复分子后，作者使用来自所有模型的相同数量的生成分子。

为了定量评估模型学习训练集分布的能力，作者计算生成分子的属性值和训练分子的属性值之间的Wasserstein距离。作者还计算了不同训练分子样本之间的Wasserstein距离，以确定最佳基准（TRAIN）。作者使用了药物相似性的定量估计 (QED)、合成可及性评分 (SA)、辛醇-水分配系数 (Log P)、精确分子量(MW)、Bertz复杂度 (BCT)、天然产物相似度（NP）等分子属性。此外，作者还使用有效性、唯一性、新颖性等标准指标来评估模型的生成能力。

作者使用带有长期短期记忆的RNN模型，基于SMILES表示 (SM-RNN) 或 SELFIES表示 (SF-RNN)进行分子生成实验。作者训练了两个图生成模型：连接树变分自动编码器 (JTVAE)和约束图变分自动编码器 (CGVAE)进行对比。

惩罚 LogP 任务

该任务目标是学习具有高惩罚LogP分数的分子的分布。作者在ZINC数据库中筛选惩罚LogP值超过 4.0 的分子构建训练数据集。

结果如图2所示。图2a中，SF-RNN产生的分布与训练分布更接近。CGVAE 和 JTVAE生成的大量分子的分数远低于训练集中最低分数。在图2b-d 中，JTVAE和CGVAE能够生成比训练数据更大的 SA分数的分子，所有模型都能学习到LogP属性的主要分布，但RNN生成更接近的分布，QED属性同样如此。

图2 惩罚LogP任务结果I

表 2展示了Wasserstein距离指标结果，两个RNN生成的Wasserstein距离低于CGVAE和JTVAE，其中SM-RNN最接近最佳基准TRAIN。

表2 Wasserstein距离指标的结果

作者进一步研究了训练数据的最高惩罚LogP区域，即值超过6.0的训练分布的尾部。图3e中，两个RNN都能学习到训练数据的这一区域，而图生成模型几乎完全忽略，只学习到更接近主模式的分子。此外，训练数据中惩罚LogP得分最高的分子通常包含非常长的碳链和更少的环（图 3b，d），RNN能够学习到这一点。这在模型产生的样本中非常明显，图 3f 显示的样本中，RNN主要产生具有长碳链的分子，而CGVAE和JTVAE产生的分子则具有许多环，语言模型学习的分布接近于图3a-d直方图中的训练分布。总体而言，语言模型比图模型可以更好地学习具有高惩罚LogP的分子分布。

图3 惩罚LogP任务结果II

多分布任务

作者通过组合以下子集创建了一个数据集：(1) 分子量 (MW) ≤ 185的GDB13分子，(2) 185 ≤ MW ≤ 425的ZINC分子，(3) 哈佛清洁能源项目( CEP）分子，460 ≤ MW ≤ 600，和（4）POLYMERS分子，MW > 600。

多分布任务结果如图4所示，RNN模型能很好地捕获数据分布，并学习到训练分布中的每种模式（图4a）。JTVAE未能学习到GDB13的分布，对ZINC和CEP的学习也很差。同样，CGVAE学习到了GDB13的分布，但低估了ZINC的分布，未能学习到CEP的分布。图4e中表明RNN模型能更紧密地学习训练分布，但CGVAE和JTVAE几乎没有区分主要模式。除CGVAE之外的所有模型都捕获了QED、SA和Bertz复杂度的训练分布（图4b-d）。

图4 多分布任务结果

大规模任务

该任务测试深度生成模型学习大分子的能力。作者从PubChem中筛选含有100多个重原子的分子，约有30万个，分子量从1250到5000，大多数分子量在1250-2000之间（图1c）。

该任务中，CGVAE和 JTVAE都未能完成训练并且完全无法学习训练数据。对于非常长的分子字符串表示，SELFIES字符表示拥有额外的优势，SF-RNN可以更紧密地匹配数据分布（图5a）。使用SMILES语法生成有效的大分子更加困难，因为要为这些分子生成更多字符，并且模型出错并产生无效字符串的可能性更高。图5b显示训练分子具有较低和较高LogP值的模式，RNN模型能够学习训练数据分布的双模态性质。

图5 大规模任务结果I

此外，作者针对训练数据包含各种不同的分子和子结构进行实验。图6a中，RNN模型充分学习到训练分子中出现的子结构的分布。即使训练分子变得越来越大并且出现的次数越来越少，两个RNN模型仍然能够生成这些分子（图5a，当分子量 >3000 时）。数据集中还包含许多肽和环肽，作者对RNN模型生成的样本进行可视化分析，以评估它们是否能够保留主链结构和天然氨基酸。图6c展示了SM-RNN和SF-RNN生成的两个肽段示例。作者对RNN学习生物分子结构的能力进行了研究，图6b表明两个RNN模型都能学习必需氨基酸的分布。RNN模型也有可能用于设计环肽，作者展示了由RNN生成的分子，该分子与粘菌素和万古霉素具有最大的Tanimoto相似性（图6d）。

图6 大规模任务结果II

作者还评估了所有模型生成分子的标准指标，（1）有效性：有效分子数与生成分子数之比，（2）唯一性：独特分子（非重复）与有效分子数之比，（3）新颖性：不在训练数据中的独特分子与独特分子总数的比率。每个任务生成1万个分子，结果如表3所示。JTVAE 和CGVAE具有更好的指标，具有非常高的有效性、唯一性和新颖性（均接近1），SM-RNN和 SF-RNN表现较差，但SF-RNN比较接近图生成模型。

表3 标准指标对比结果

3

讨论

在这项工作中，为了测试化学语言模型的能力，作者引入了三个复杂的分子生成任务，使用语言模型和图生成模型从具有挑战性的数据集中生成分子。结果表明，语言模型是非常强大、灵活的模型，可以学习各种不同的复杂分布，而图生成模型在很多方面表现较弱。

语言模型SM-RNN和SF-RNN在所有任务中都表现良好，优于基线方法。实验结果表明SF-RNN在每项任务中都有更好的标准指标，但SM-RNN有更好的Wasserstein距离指标。此外，SF-RNN比 SM-RNN具有更好的新颖性，这可能由于SELFIES语法导致语言模型只需记忆更少的训练数据。这也有助于解释为什么SF-RNN比SM-RNN具有更好的标准指标但更差的Wasserstein指标。此外，数据增强和随机SMILES可用于提高SM-RNN的新颖性得分。

图生成模型JTVAE和CGVAE不如语言模型灵活。对于惩罚LogP任务，得分为2的分子与得分为4的分子之间的差异很小。有时改变单个碳或其他原子会导致分数大幅下降，这可能解释了为什么CGVAE非常不适于该任务。对于多分布任务，JTVAE和 CGVAE表现较差的原因是：JTVAE必须学习广泛的树类型，其中许多没有类似环的较大子结构（GDB13分子），而另一些则完全是环（CEP和 POLYMERS）；CGVAE必须学习大量不同的生成轨迹，这很困难，尤其是因为它在学习过程中只使用一个样本轨迹。出于同样的原因，这些模型无法训练 PubChem中大的分子。虽然语言模型可以灵活地生成更大的分子，但图生成模型更易于解释。

未来方向：对分子SMILES和SELFIES表示在深度生成模型中的使用进行更全面的评估；语言模型的改进，因为这类模型无法解释其他重要信息，如分子几何；探索语言模型在学习越来越大的化学空间方面的能力。

参考资料

Flam-Shepherd, D., Zhu, K. & Aspuru-Guzik, A. Language models can learn complex molecular distributions. Nat Commun 13, 3293 (2022). https://doi.org/10.1038/s41467-022-30839-x

数据

https://github.com/danielflamshep/genmoltasks

代码

JTVAE

https://github.com/wengong-jin/icml18-jtnn

CGVAE:

https://github.com/microsoft/constrained-graph-variational-autoencoder

RNN

https://github.com/molecularsets/moses