作者 | 程玉 审稿 | 李芬
今天给大家介绍的是纽约大学于2021年5月26日发表在Nature Communications的一篇论文，作者引入了一个遮掩图模型（masked graph model，MGM），它通过捕获未观察到的节点（原子）和边（键）上的条件分布来学习图上的分布，通过迭代遮盖和替换初始化图的不同部分来训练并从中采样，并利用GuacaMol 分布学习基准评估QM9和 ChEMBL数据集上的表现。该模型优于先前提出的基于图的方法，并且与基于 SMILES 的方法可比。最后，作者展示了MGM模型生成具有指定属性所需值的分子，同时保持与训练分布的物理化学相似性。
1
背景
设计具有所需性质的分子是药物发现和材料设计的重要组成部分，但由于可合成的分子化学空间巨大，因此设计分子仍具有很大的挑战性。而最近提出的分子生成模型，包括神经自回归模型、变分自编码器、对抗性自编码器、生成对抗网络和强化学习，这些方法背后的一个统一原则是它们对分子图的分布进行建模。一旦捕获了分布，就会相应地对新的分子图进行采样，模型基于此分布从而进行分子生成。
作者基于无需直接对联合分布建模就能够从中采样的思想，探索了一种新的概率图生成方法，提出了一个遮掩图模型MGM，这是一个基于图的生成模型，在给定图的其余部分的情况下学习被遮盖图组件的条件分布。MGM避免了先前提出的模型的许多缺点。例如避免考虑边缘化和变量的顺序。另外，作者在QM9 和 ChEMBL 数据集上进行实验，使用 GuacaMol分布学习基准进行评估，并且这些基准指标彼此相关。具有较高有效性、KL散度和Fréchet ChemNet距离分数的模型通常具有较低的新颖性。因此，基于不同度量之间的权衡评估模型可能比基于启发式评估模型提供更多信息。而作者发现通过改变生成超参数，MGM比以前最先进的基线模型更能有效地平衡这些指标。
2
模型框架
作者使用单个图神经网络来参数化由给定图产生的任何条件分布，并且假设遮掩部分η在给定剩余图的情况下彼此条件独立。为了训练模型，作者使用来自训练数据集D的全图，之后开始破坏一个全图G的完整性。受语言模型的影响，利用特殊符号MASK 随机替换了一些节点和边缘特征，从而可以获得条件分布。然后，给定图的其余部分，作者最大化遮掩部分 η 的对数概率。这一过程类似于带有mask的语言模型，即在给定不完整句子的情况下预测mask部分的单词。这一过程的损失函数如下图所示：
另外，在生成之前，需要对分子进行初始化，这里采用两种不同的初始化方法。第一种方式，称之为训练初始化，使用来自训练数据的随机图作为初始图。第二种方法，为边缘初始化，根据组件在训练集中类别分类结果来初始化每个图组件。例如，在图中，边类型为 r的概率等价于在训练集中边类型为 r的比例。
然后，利用Gibbs采样方法从学习到的条件分布中迭代更新图组件。在每个生成步骤中，随机均匀地采样遮掩部分η的片断α，并用MASK符号替换这些分量的值。作者通过将这些片段mask图传递给模型，根据预测分布对mask组件的新值进行采样，并将这些值放置在图中来计算条件分布整个步骤重复多次。
如果想要获得带有目标理化性质的分子，可以对原模型进行一定的修改，加入一个额外的条件输入。该输入y是分子目标属性的图级表示，y*是分子属性的真实值，具体的优化函数如下所示：
3
实验
（1）性质依赖性分析
由于作者注意到其中一些指标可能会衡量类似的属性，所以探究了Guacamol框架中的指标之间是否存在依赖性。通过计算QM9 数据集上所有指标之间的Spearman相关性，具体结果如下表3-1所示：
表3-1 MGM性质关联性分析
从上表，可以得到a）有效性、KL 散度和 Fréchet 距离分数彼此高度相关。b）这三个指标与新颖性得分呈负相关。c）独特性与任何其他指标没有很强的相关性。结果表明，我们可以查看指标的一个子集，即独特性、Fréchet 距离和新颖性，以衡量生成质量。
（2）和基线模型比较
作者将MGM的结果与基线模型进行比较。在QM9和ChEMBL数据集上的分布基准结果分别如表3-2和表3-3所示。在QM9上，MGM与现有的基于SMILES的方法表现相当。与CharacterVAE和GrammarVAE相比，MGM显示出更高的有效性和独特性，同时新颖性较低。与自回归LSTM和Transformer模型相比，MGM有效性、KL散度和Fréchet距离分数较低，但独特性略高，新颖性显著提高。
表3-2 QM9数据集模型比较结果
表3-3 ChEMBL数据集模型比较结果
（3）条件生成比较结果
作者以分子量(MolWt)和Wildman-Crippen分配系数(LogP)作为目标。MGM模型和基线模型的条件生成结果下表3-4所示。
表3-4 条件生成比较结果
MGM生成的分子的属性值接近所需属性的目标值。对于MolWt=120、MolWt=125、LogP=0.2和LogP=0.8，MGM生成的分子的平均目标属性比NAT GraphVAE生成的分子更接近目标值。
4
总结
在这项工作中，作者提出了一个遮掩图模型，通过对图组件的子集进行迭代采样，采样该模型中新的分子图。在未来，可以将原子间距离等附加信息纳入图表示中。也可以扩展到蛋白质生成上，以氨基酸为节点，以接触图为邻接矩阵。在这个框架中可以使用条件生成来重新设计蛋白质，以实现所需的功能。此外，该模型也适用于先导化合物优化。最后，由于该方法广泛适用于通用图结构，未来也可以将其应用在非分子数据集上。
参考资料
Mahmood, O., Mansimov, E., Bonneau, R. et al. Masked graph modeling for molecule generation. Nat Commun 12, 3156 (2021).
https://doi.org/10.1038/s41467-021-23415-2
代码：https://github.com/nyu-dl/dl4chem-mgm
