近期刚刚在NeurlPS-21上发表了一篇文章,用于分子性质预测的图预训练策略。
本文亮点在于①提出一套分子破裂的方法;②提出构建官能团树的方法;②同时在节点级和官能团级进行预训练,自动分配不同的权重。
NeurIPS(简称 NIPS),全称神经信息处理系统大会(Conference and Workshop on Neural Information Processing Systems),是一个关于机器学习和计算神经科学的国际会议。该会议固定在每年的12月举行,由NeurIPS基金会主办。
1.Introduction 1.1 背景Motif:是一段典型的序列或者一个结构。一般来说,我们称为基序。一般情況下是指构成任何一种特征序列的基本结构。通俗来讲,即是有特征的短序列,一般认为它是拥有生物学功能的保守序列,可能包含特异性的结合位点,或者是涉及某一个特定生物学过程的有共性的序列区段。比如蛋白质的序列特异性结合位点,如核酸酶和转录因子,在这里指的是官能团。
GNN在各种分子生成和预测任务中取得了显著的成功。在标记数据稀少的情况下,GNN可以在未标记的分子数据上先进行预训练,以便在针对特定任务进行微调之前先学习一般的语义和结构信息。然而,大多数现有的GNN自监督预训练框架只关注节点级或图级任务,无法捕获子图或官能团中的丰富信息
1.2 本文工作本文通过为GNN引入一种 自监督官能团生成框架,提出了基于官能团的图自监督学习,可以有效学习分子中的官能团信息。
2. 模型介绍官能团(分子图中经常出现的子图)通常携带有关分子性质的指示信息
Overview
模型框架包含三个部分:
① Chemistry-inspired Molecule Fragmentation
② Motif Generation
③ Multi-level Self-supervised Pre-training
2.1 Chemistry-inspired Molecule Fragmentation①首先,为了从分子图中提取motif,本文设计了一种分子碎片化方法,该方法利用基于(逆合成的算法) BRICS 和 一些附加规则 提取官能团。
②设计了一个通用的基于官能团生成的预训练框架,让 GNN 进行拓扑和标签预测,这个生成框架可以通过两种不同的方式实现,BF和DFS。
③最后,为了考虑分子图中的多尺度信息,本文引入了多级自监督预训练,在节点级和官能团级进行预训练。
分子裂解方法要达到的目标:
(1) The union of all motifs should equals graph;
(2) In a motif tree, the motifs should have no intersections;
(3) The induced motifs should capture semantic meanings, e.g., similar to meaningful functional groups in the chemistry domain.
(4) The occurrence of motifs in the dataset should be frequent enough so that the pre-trained GNNs can learn semantic information of motifs that can be generalized to downstream tasks.
①官能团的集合组成图;(每个官能团都是图G的子图)
②在官能团树中,分解出来的官能团应该互不相交;
③产生的官能团能够捕捉语意,就像化学领域有意义的功能群一样;
④数据集中官能团出现频率应足够高,这样PreGNN就能学习功能团的语义信息,以便推广到下游任务。
方法如下:
(1)为了划分分子官能团,首先使用BRICS方法;
(2)为了去除冗余官能团,进行进一步的分解,
①化学键两边的原子一个在环上,另外一个不在环上,破裂化学键
②一个原子不在环上,且邻居数≥3,破裂化学键
(3)构建官能团树
最后:处理完整个数据集后,得到预训练数据集的官能团词典
基于化学反应模板 BRICS:定义了16条规则,破裂了 与一系列 化学反应 相匹配的 分子中的键。虚拟原子连接到切割位点的每一端,标志着两个片段可以连接在一起的位置。 BRICS 裂解规则旨在保留 具有 有价值的 结构和功能内容的 分子组分,例如芳香环。
2.2 Motif Generation为什么要进一步分解?
①BRICS 仅根据一组有限的化学反应破坏键,并且往往会产生几个大碎片;
②BRICS 产生了许多具有相同底层结构的官能团,如果不进一步分解,就无法达到“官能团不相交“这个要求。(结合不同氧原子的呋喃环)
Goal:Let GNNs learn the data distribution of graph motifs so that the pre-trained GNNs can easily generalize to downstream tasks with several finetuning steps on graphs from similar domains.
让GNN学习图官能团的数据分布,以便预训练的GNN 可以通过 对来自相似域的图 进行几步微调,然后泛化到下游任务.
本文旨在通过 最大化 官能团树的 可能性 来预训练GNN模型
给定分子图G = (V, E) 和GNN模型fθ,首先将分子图转换为官能团树 T(G) = (V,ε,χ)。然后通过 GNN 模型对这个motif tree 的似然性建模为 p(T(G);θ) ,表示这些motifs 是如何被标记和连接的。
本文方法旨在通过最大化motif tree的可能性来预训练 GNN 模型 fθ,即 θ∗= argmaxθp(T(G);θ)。为了对motif tree的可能性进行建模,本文还设计了拓扑预测和motif标签预测 的 特殊预测头 并与 fθ 一起优化
预训练后,只有 GNN 模型 fθ 被转移到下游任务。
大多数现有的图生成工作都用自回归方式来分解概率目标,也就是本文的p(T(G); θ)。对于每个分子图,将其分解为一系列生成步骤,每次迭代添加一个官能团。用π来表示官能团的添加顺序,那么p(T(G); θ)等于所有可能的排序的期望:
自回归模型(Autoregressive model),是统计上一种处理 时间序列 的方法,用同一变量 例如x 的之前各期,亦即x1至xt-1来预测本期xt的表现,并假设它们为 线性关系。为这是从 回归分析的线性分析而来,只是不用x预测y,而是用x预测 x(自己),所以叫做自回归。
公式: 在参数θ下,所有可能的排列π的期望,也就是各种官能团树的可能性 的 期望,Vπ代表官能团的label,επ代表所有官能团之间的边。
为了简单,假设所有的排列顺序都存在而且是等概率的,上标π暂时省略,对于一个排列顺序,生成motif treeT(G)的步骤为:在每一个时间步i,利用在第i步之前生成的 所有官能团 以及他们的边 来生成新的官能团 还有他们与 现有官能团 之间的边。官能团树的自回归生成过程为下面公式。此时问题转化成了如何选择有效的生成顺序以及如何对条件概率进行建模。
在motif tree中,节点代表官能团,边代表官能团之间的关系。从头开始生成 motif tree的步骤:
①选择根(canonical order规范顺序 中的第一个)
②用BFS或者DFS生成:
BFS:MGSSL 逐层生成motif节点。对于第k层的motif节点,MGSSL进行拓扑预测和标签预测。如果第k层的motifs的所有孩子节点都生成了,MGSSL将移动到下一层。
DFS:对于每个访问的motif节点,MGSSL首先进行拓扑预测,即该节点是否有要生成的孩子节点。如果生成了一个新的孩子节点,就对它进行标签预测,然后递归这个过程。
③由于兄弟节点的顺序是不明确的,所以BFS和DFS中的motif节点顺序不是唯一的。在实验中,我们以一个顺序进行预训练,并将这个问题和其他潜在的生成顺序留给未来探索。
2.3 Multi-level Self-supervised Pre-trainingTopology Prediction :在DFS和BFS时用来预测是否生成孩子,通过一个隐层+sigmoid来预测。(根据message和motif embedding)
Label Prediction : 使用softmax来预测。
多层次自监督预训练包含两层:分别是官能团层和分子层
在原子层,本文mask节点/边属性,并让GNN基于相邻结构预测这些属性。
在官能团层,构建了官能团树,进行官能团生成预训练。在每一步,根据现有的官能团和连接方法,迭代地进行拓扑官能团预测。(在2.2中)
对于原子级预训练,我们利用attribute masking让GNN首先学习节点属性和边属性的规则:
①随机采样的节点和键属性(例如原子数、键类型),把它们替换为特殊的掩蔽指示符;
②用GNN获得相应的节点或边嵌入(边嵌入可以通过结合 边两端的 节点的嵌入 获得);
③将embedding输入到全连接层中预测节点或边类型。
节点属性与边属性的损失用交叉熵损失计算,分别表示为Latom与Lbond。目标是让预训练损失最小:
3. 实验 3.1 数据集λ代表各个损失的权重。本文采用多任务学习中的MGDA-UB算法(多重梯度下降)来有效地解决优化问题。由于MGDA-UB通过Frank Wolfe算法在每个训练步骤计算权重λi,因此我们不必明确给出权重。
【(MGDA-UB的idea其实很简单:通过链式法则的推导, 证明了, 在梯度为full rank的情况下, 我们不需要对每个任务的向后传播都算到底(所有层), 只需要算到共享模型的最后一层, 用这个去解出各个任务的权重, 既能加快速度, 又能得到一个比较好的解)】
self-supervised pre-training tasks: 25w molecules sampled from the ZINC15
downstream finetune tasks: 8 binary classification benchmark datasets contained in MoleculeNet
3.2 BeselinesZINC 数据库:是目前最大的有机小分子化合物库之一,不少类药分子的前期虚拟筛选都是基于这个数据库的(取出的25W个分子都是未标注的)
MoleculeNet:适合分子机器学习的大型基准 MoleculeNet,MoleculeNet 提供多个公共数据集、建立了评估度量,并提供多个分子特征化(molecular featurization)和学习算法的高质量开源实现
对比五种先进的自监督预训练方法:
①Deep Graph Infomax (2019):以无监督的方式去学习节点的嵌入向量。
②Attribute masking (from PreGNN, 2020):屏蔽节点/边的属性,并让GNN预测这些属性。
③GCC (2020):一种自监督GNN预训练框架,可以捕获跨多个网络的通用拓扑结构特性。
④Grover (2020):基于节点嵌入预测上下文属性,将上下文信息编码到节点嵌入中。
⑤GPT-GNN (2020):利用图生成作为自监督任务来预训练 GNN,可预测屏蔽边和节点属性。
①只要用了PreTrain,效果就比直接fine-tune效果要好;
②在这8个任务中,本文方法在7个中表现最好,证明了基于官能团的自监督预训练的有效性;
③BFS比DFS有点优势,故接下来的实验采用BFS(可能是因为MGSSL需要按BFS顺序分层生成基序,这有助于GNN了解更多基序的结构信息)
MGSSL比基线方法在训练集和测试集上收敛更快。
值是(1)中8个数据集上的平均值得到的。
不论使用哪种GNN结构,效果都有提升,提升最高的是GIN。
不同层次预训练的消融研究
①只在原子级别预训练
②先原子级别+后官能团级别
③两个级别同时训练,在2.3部分的损失函数地方,给与不同的权重,性能最好
分子切割方法的影响
①仅 BRICS 倾向于产生含有大量原子的图。由于组合爆炸,其生成的motif词汇超过10w,而90%以上的motif频率小于5;
②JT-VAE将分子分割成环和键,其官能团词汇表数量小于800。本文方法产生了大约12k个不同的图案。通过结合不同的裂解策略,我们能够裂解具有中等粒度的分子
本文模型效果最好,原因如下:
①当motif分割太粗糙且基序词汇量太大时,生成的motif tree具有较少的节点,GNN很难捕捉到motif的结构信息。此外,生成的motif出现频率较低,这使得GNN无法学习模体的一般语义信息,而这些信息可以推广到下游任务。
② 当motif分割太精细时,许多生成的基序是单原子或键,这阻碍了GNN通过motif生成任务学习更高层次的语义信息。
本文为提出了基于官能团的图自监督学习MGSSL:首先使用BRICS和两个附加规则分割分子图,并推导出有意义的官能团词典;其次,设计了一个官能团生成预训练框架,并考虑了BFS和DFS两种顺序,在每一步中,需要预先训练的GNN进行拓扑和模体标签预测;此外,同时在节点级和官能团级进行预训练。



