文献总结|通过随机化 SMILES 序列提升分子生成模型质量

Category 碎碎念

本文介绍于 2019 年发表在 Journal of Cheminformatics 上的一篇文章,文章原标题为 Randomized SMILES strings improve the quality of molecular generative models。

doi.org/10.1186/s13321-019-0393-0

引言

SMILES 是深度学习邻域最常用的化学分子描述方式,在众多研究中使用的 SMILES 都是规范形式(canonical format)。而规范形式的 SMILES 常常会使模型具有一定偏向,还常常导致模型生成大量极相似的分子。文章考查了不同形式的 SMILES 编码对于模型的影响,主要包括规范形式、无约束随机形式和有约束随机形式。

n

例如图片所示,对于阿司匹林分子具有三种不同的编码形式,规范形式(a)是 RDKit 工具中的默认形式。经过随机化后,起始原子发生了变化,无约束形式与有约束形式的区别在于,无约束形式先编码芳环,再编码支链,而有约束形式则先编码支链,最后编码芳环。

方法

预处理

对于 SMILES 的预处理在文中被称为 tokenization,其目的是使 SMILES 序列成为能够被神经网络识别的形式。其步骤包括在字列始末加上起始与终止符号、将两个字符的元素(如 Cl 和 Br)视作一个 token 等,最终将 SMILES 序列转化为 one-hot 形式。

模型结构

n

分子生成模型由 RNN 构成,经过处理后的 SMILES 序列将逐个 token 进入 \(m\) 维的 embedding 层,再进入 \(l\) 层大小为 \(w\) 的 LSTM 或 GRU 单元,在 RNN 层间也可以设置 dropout 层,最终得到 token 的概率分布 \(Y_{ij}\)

数据集

GDB-13 数据来自于官网,使用 RDKit 对 SMILES 规范化。

ChEMBL 25 数据来自于 ChEMBL 官网,使用 MolVS 0.1.1 数据库对其处理,主要的处理包括去重、去除含有重原子的分子等,最终数据集中的分子都编码为 one-hot 形式。

在使用随机化 SMILES 序列训练模型时,将数据集中的分子转化为随机化形式,也就是说,不论分子的编码形式是哪一种,数据集所使用的分子都是相同的。

训练

将训练集中的若干 SMILES 整合为批数据后用于训练 RNN,在训练过程中使用了 Teacher’s Forcing 方法。

未预训练的模型通常没有很好的预测能力,因此每个 RNN 单元产生的错误结果势必会导致模型训练时间大大增加。Teacher’s Forcing 方法就是为了解决这个问题,具体步骤是忽略每个 RNN 单元输出的预测结果,而是以正解的标签作为下一个 RNN 单元的输入,避免了一个 RNN 单元的错误预测影响整个网络的训练。

评价指标

对生成模型的评估

文章假设最好的模型对于验证集、训练集和生成分子集的 NLL 分布应当是均一且相等(uniform and equal)的,所以就用 JSD(Jensen–Shannon Divergence)衡量若干概率分布的差异:

$$JSD=H\left(\sum_{d\in D}\alpha_i\cdot d_i\right)-\sum_{d\in D}\alpha_iH(d_i)$$

其中 \(H(d)\) 是给定概率分布的信息熵,\(\alpha_d\) 为权值,满足 \(0<\alpha_d<1\)\(\sum\alpha_d=1\)。当 \(d_i=d_j\) 时,就有 \(JSD\rightarrow 0\) ,但此时还有可能 \(i\not =j\), 所以该指标只能保证概率分布是相等的而非均一的。

UC-JSD(Uniformity–Completeness JSD)解决了这个问题,考虑给定的三个 NLL 向量,有

$$NLLS=\{NLL_\mathrm{validation},NLL_\mathrm{training},NLL_\mathrm{sampled}\}$$

并指定 \(\alpha_i=1/3\)。当 \(UC_{JSD}\rightarrow0\) 时,生成的分子具有相同的 NLL 或三个概率分布均一。

对生成结果的评估

对于 GDB-13 数据集,文章对生成结果提出了三个要求:

  1. 均一(uniformity):对于各取样结果(生成分子)的概率相等
  2. 完备(completeness):能够从 GDB-13 数据集中取样出所有分子
  3. 封闭(closedness):只有 GDB-13 数据集中的分子被取样

针对这三个要求,文章定义了三种量化指标:

  1. \(\mathrm{uniformity}=ratio_\mathrm{unique}/\varphi(k)\)
  2. \(\mathrm{completeness}=ratio_\mathrm{unique}/\varphi(|in|)\)
  3. \(closedness=ratio_\mathrm{in}\)

\(\varphi(k)\) 表示在 GDB-13 的理想模型(每个分子概率完全均一)上生成 \(k\) 个分子中不同分子的比例。

此外,具有合法 SMILES 的分子就称为 valid,在 GDB-13 数据集中的分子就称为 in,在 GDB-13 中独一的分子就称为 unique。

$$UCC=\mathrm{completeness}\cdot\mathrm{uniformity}\cdot\mathrm{closedness}$$

作为对模型结果的整体打分。

工具

  • PyTorch 1.0.1
  • RDKit 2019_03_01
  • Spark 2.4.3:数据处理
  • MOSES 和 FCD:提供评价指标

结果

规范形式与随机化形式

n

在 GDB-13 中随机选择 1000000 个分子用于模型训练,结果表明随机化形式 SMILES 训练的模型生成了更多 GDB-13 数据集中的分子。从生成结果的其他指标上来看,与规范形式 SMILES 训练的模型相比,随机化形式模型产生了更好的结果。

n

超参数的优化

在 RNN 间的 dropout 比例、批处理数据的大小都属于模型的超参数,针对于这两个超参数,文章比较了改变这两个超参数对于 3 种模型 4 个指标的影响。结果表明,在增加 dropout 比例后,规范形式的模型表现出更好的效果,其完备性上升了,但封闭性下降了,也就是说模型以产生更多错误为代价,生成了更多 GDB-13 数据集中的分子。至于批大小,随着批处理数据的增加,所有模型的效果都改善了。

UC-JSD 评估

n

由于不同模型的输出空间不同,不同模型的 UC-JSD 不能相互比较,但 UC-JSD 可以用作为对于每个 epoch 训练结果的损失函数。UC-JSD-UCC 图像可以明确显示中模型在各个 epoch 中优化的过程,最后得到的最优模型应当是具有最低 UC-JSD 与 最高 UCC 的模型。

小分子集测试

n

文章中分别使用 1000 个分子与 10000 个分子训练模型。使用 1000 分子训练的规范形式模型仅生成了 14.5% 的 GDB-13 数据集分子,其合法性也仅有 0.504,可以认为模型没有学习到数据集的特征。当训练集增加到 10000 个分子,模型的效果也有所提高,但随机化形式的模型效果都要高于规范形式模型。

ChEMBL 数据测试

n

文章还使用 ChEMBL 的类药分子训练集测试了两种两种模型,最后绘制了两种模型对于 NLL 的核密度估计(Kernel Density Estimates)结果。结果表明,规范形式模型生成的规范形式 SMILES 能够很好地接近训练集和验证集,但是它不能很好地产生随机形式的 SMILES。与之相反,随机化形式模型不仅能产生随机化形式的 SMILES,产生规范形式 SMILES 的平均概率也很高。这个结果可以说明规范形式模型的输出空间是随机化形式模型的子集。

讨论

当以规范形式 SMILES 表示分子时,模型不仅要学习如何生成合法的 SMILES 序列,还需要学习如何以规范形式生成。相应地,如果使用随机化形式的 SMILES 序列,模型就可以不受限制地学习 SMILES 语法。

从另一方面理解,就是随机化 SMILES 序列具有数据增强(Data Augmentation)的功能,类似于对图片进行裁切、旋转以提升模型识别的准确度,随机化 SMILES 序列可以看作是从不同的角度「观察」一个分子,这不仅能提高模型应对噪声的能力,还有利于模型更好地学习到训练集化学空间的特征。

在文章中构建的模型很难生成较大的分子,可以将模型生成每个 token 的步骤看作一系列可选的行为(加原子、加键),这种模型的问题就在于随着分子增大,可选的行为激增,所以使用 SMILES 生成较大、较复杂的的分子仍有难度。

结论

文章基于 RNN 构建了分子生成模型,在不用修改仍何模型结构的情况下,使用随机化的 SMILES 序列编码分子可以显著提高模型质量。

在使用 ChEMBL 训练集的测试中,文章进一步确定了相比于常用的规范形式模型,随机化形式模型具有更大的输出空间,也是性能更优的原因。