概述
背景概述
本篇文章来自"CHEMICAL-REACTION-AWARE MOLECULE REPRESENTATION LEARNING",送审ICLR2022。
分子表征学习(Molecule representation learning,MRL)旨在将分子嵌入到实向量空间中。现有的基于SMILES(简化分子线性输入系统)或GNN的MRL方法要么以SMILES字符串作为输入,难以编码分子的结构信息,要么过度强调GNN对于结构感知的重要性,而忽略了图结构的泛化性和鲁棒性。
作者提出使用化学反应协助分子表征的学习,约束条件为保持分子在嵌入空间中的化学反应的等价性,可以理解为加入一个简单的先验条件:质量守恒定律,强制要求每个化学方程式的反应物嵌入和生成物嵌入的总和相等,该限制在保持嵌入空间的有序性和提高分子嵌入的泛化能力中被证明是有效的。此外,该模型可以使用任何GNN作为分子编码器,与GNN结构无关(论文提到的方法是可泛化的学习策略,适用于大多数机器学习模型)。
实验结果表明,这种方法在各种下游任务中都达到了最佳性能,超过了最佳baseline。
方法概述
在化学中,一般用IUPAC命名法,分子式,结构式,骨架式等形式来表示一个分子,然而这些形式最初是为人类而不是计算机设计的,为了便于机器理解分子,MRL被发展起来,MRL将分子映射到低维稠密的向量空间。分子的表征(也被称为embedding)可用于广泛的下游任务,比如:化学反应预测,分子性质预测,分子结构预测。
近年来,大部分方法以SMILES作为输入,语言模型擅长对序列建模,但却难以感知分子的结构信息,因为SMILES是分子结构的一维线性化表示。GNN在处理分子表示时局限于设计新颖的GNN结构(总是根据某个具体任务提出特定的GNN,比如针对多元信息采用异构图),忽略了MRL的主旨:可泛化可移植。出于以上考虑,激发了作者探索一个通用的学习策略。
在论文中,作者使用化学反应协助学习分子表征,提高泛化性。化学反应由化学反应方程式表示,反应物在左侧,生成物在右侧。作者的想法是保持分子在嵌入空间的等价性。作者将该模型称为MolR(chemical-reaction-aware molecule embeddings,化学反应感知分子嵌入)。
方法
结构分子编码器
分子图用 G = ( V , E ) G=(V,E) G=(V,E)表示,其中, V = { a 1 , . . . , a n } V=\left\{a_{1},...,a_{n}\right\} V={a1,...,an}是非氢原子集合, E = { b 1 , . . . , b m } E=\left\{b_{1},...,b_{m}\right\} E={b1,...,bm}是键的集合。在这项工作中,每个原子 a i a_{i} ai都有初始特征向量,该向量包括四个one-hot向量,每个one-hot向量表示一种原子属性,四个one-hot向量对应属性为:
- 原子类型;
- 电荷量;
- 原子是否是芳香环;
- 原子附着的氢原子数量;
每个one-hot向量都额外扩展一个条目用于保存属性中的未知类型,以在训练中处理未知值。
对于键,虽然有双键和单键的区别,但论文为了计算方便,故没有区分键的类型。
作者采用消息传递范式作为计算框架,即聚合原子的邻居和自身的表示迭代更新原子的表示。在消息传递的第 k k k层为: h i k = a g g r e g a t e ( { h j k − 1 } j ∈ N ( i ) ∪ { i } ) , k = 1 , . . . , K h_{i}^{k}=aggregate(\left\{h_{j}^{k-1}\right\}_{j\in N(i)\cup\left\{i\right\}}),k=1,...,K hik=aggregate({hjk−1}j∈N(i)∪{i}),k=1,...,K这里, h i k h_{i}^{k} hik是原子 a i a_{i} ai在第 k k k层的表征向量。
最后,readout函数用于聚合最后一个消息传递层输出的所有节点的表示,以获得整个分子的表示
h
G
h_{G}
hG:
h
G
=
r
e
a
d
o
u
t
(
{
h
i
K
}
a
i
∈
V
)
h_{G}=readout(\left\{h_{i}^{K}\right\}_{a_{i}\in V})
hG=readout({hiK}ai∈V)readout函数可以是简单的置换不变函数,例如求和,也可以是复杂的graph-level pooling算法。
- 上图a是GNN编码器处理脯氨酸(proline)分子的示意图(氢原子被省略)。
- 图b是丙酸与丙醇的费舍尔酯化反应示意图,以及由模型预测得到的反应模板,反应中心用橙色表示,距离反应中心1跳或2跳的原子用浅橙色表示。
- 图c是小批量化学反应的对比损失(contrastive loss), d i j d_{ij} dij表示embedding h R i h_{R_{i}} hRi和 h P j h_{P_{j}} hPj的欧氏距离。
化学反应等价性策略
作者希望在分子嵌入空间中保持反应物和生成物的等价性,即: ∑ r ∈ R h r = ∑ p ∈ P h p \sum_{r\in R}h_{r}=\sum_{p\in P}h_{p} r∈R∑hr=p∈P∑hp在一个化学反应中(在一个化学方程式中), R R R是反应物的集合, P P P是生成物的集合。这个限制可以有效提高分子嵌入的质量。
作者在论文中提出了命题1:
设
M
M
M是分子的集合,
R
R
R和
P
P
P分别是反应物集合和生成物集合,如果
R
→
P
⇔
∑
r
∈
R
h
r
=
∑
p
∈
P
h
p
R\rightarrow P\Leftrightarrow \sum_{r\in R}h_{r}=\sum_{p\in P}h_{p}
R→P⇔∑r∈Rhr=∑p∈Php对于所有化学反应都成立,那么
→
\rightarrow
→是
2
M
2^{M}
2M上满足下面3个属性的等价关系:
- 自反性:对于所有 A ∈ 2 M , A → A A\in 2^{M},A\rightarrow A A∈2M,A→A;
- 对称性:对于所有 A , B ∈ 2 M , A → B ⇔ B → A A,B\in 2^{M},A\rightarrow B\Leftrightarrow B\rightarrow A A,B∈2M,A→B⇔B→A;
- 传递性:对于所有 A , B , C ∈ 2 M A,B,C\in 2^{M} A,B,C∈2M,如果 A → B A\rightarrow B A→B且 B → C B\rightarrow C B→C,则 A → C A\rightarrow C A→C;
以上三个属性正好符合化学反应质量守恒的规则,化学反应等价性对分子的嵌入表达施加了更加强大的约束,分子嵌入的可行解将更加稳健,整个嵌入空间将更加有序。
接着,作者针对广义的反应中心(广义的反应中心被定义为将反应物转换为生成物所需的一组最小图集合)提出以下命题2:
设
R
→
P
R\rightarrow P
R→P是一个化学反应,
R
R
R是该反应的反应物集合,
P
P
P是该反应的生成物集合,
C
C
C是该化学反应的反应中心,对于反应物中的任何一个原子
a
a
a,经过
K
K
K次消息传递后表达为
h
a
K
h_{a}^{K}
haK,当且仅当原子
a
a
a和反应中心
C
C
C之间的距离小于
K
K
K跳时,该原子的化学反应剩余项被定义为:
f
(
h
a
K
)
=
∑
r
∈
R
h
r
−
∑
p
∈
P
h
p
f(h_{a}^{K})=\sum_{r\in R}h_{r}-\sum_{p\in P}h_{p}
f(haK)=r∈R∑hr−p∈P∑hp命题2表明,反应物嵌入和生成物嵌入之间的剩余仅取决于距离反应中心小于
K
K
K跳的原子。比如上图b,剩余物完全取决于反应中心(橙色)以及距离反应中心1跳或者2跳的原子(浅橙色)。这意味着方程:
R
1
−
C
H
2
C
O
O
H
+
R
2
−
C
H
2
C
H
2
O
H
→
R
1
−
C
H
2
C
O
O
C
H
2
C
H
2
−
R
2
+
H
2
O
R_{1}-CH_{2}COOH+R_{2}-CH_{2}CH_{2}OH\rightarrow R_{1}-CH_{2}COOCH_{2}CH_{2}-R_{2}+H_{2}O
R1−CH2COOH+R2−CH2CH2OH→R1−CH2COOCH2CH2−R2+H2O适合用于任何官能团
R
1
R_{1}
R1和
R
2
R_{2}
R2,上面的方程被称为反应模板,它概括了同一类别内的一组化学反应,反应模板可以适应训练集内没有出现过但符合已知反应类型的化学反应。
训练方法
论文中使用小批量对比学习框架进行优化。对于小批量数据 B = { R 1 → P 1 , R 2 → P 2 , . . . , R n → P n , } B=\left\{R_{1}\rightarrow P_{1},R_{2}\rightarrow P_{2},...,R_{n}\rightarrow P_{n},\right\} B={R1→P1,R2→P2,...,Rn→Pn,},首先使用编码器处理批量数据 B B B中所有的反应物和生成物分子,得到embedding,将匹配的"反应物-生成物" ( R i , P i ) (R_{i},P_{i}) (Ri,Pi)标记为正对,不匹配的 ( R i , P j ) , i ≠ j (R_{i},P_{j}),i\neq j (Ri,Pj),i=j标记为负对,目标是最小化正对反应物和生成物之间的embedding之差,最大化负对反应物和生成物之间的embedding之差。为了避免优化目标被大量负对主导,使用边距作为正则项,优化目标 m i n L B minL_{B} minLB如下: L B = 1 ∣ B ∣ ∑ i ∣ ∣ ∑ r ∈ R i h r − ∑ p ∈ P i h p ∣ ∣ + 1 ∣ B ∣ ( ∣ B ∣ − 1 ) ∑ i ≠ j m a x ( γ − ∣ ∣ ∑ r ∈ R i h r − ∑ p ∈ P j h p ∣ ∣ , 0 ) L_{B}=\frac{1}{|B|}\sum_{i}||\sum_{r\in R_{i}}h_{r}-\sum_{p\in P_{i}}h_{p}||+\frac{1}{|B|(|B|-1)}\sum_{i\neq j}max(\gamma-||\sum_{r\in R_{i}}h_{r}-\sum_{p\in P_{j}}h_{p}||,0) LB=∣B∣1i∑∣∣r∈Ri∑hr−p∈Pi∑hp∣∣+∣B∣(∣B∣−1)1i=j∑max(γ−∣∣r∈Ri∑hr−p∈Pj∑hp∣∣,0)其中, γ \gamma γ是边距的超参数,我们可以使用随机梯度下降SGD最小化上述目标训练模型。
实验结果
化学反应预测
论文使用了USPTO-479k数据集,每个反应实例包含5种反应物和1种生成物。化学反应预测的任务可以表述为一个排序问题,在推理时,给定化学反应的反应物集合 R R R,将测试集中所有生成物视为候选者 C C C,根据反应物embedding h R h_{R} hR和候选生成物embedding h C h_{C} hC之间的L2距离对所有生成物排序,选出置信度最高的生成物。
化学反应的预测结果比较见下表,MolR策略优于其他方法
分子性质预测
作者在5种数据集中测试MolR,每个数据集包含数千个SMILES分子。预测的AUC结果如下表所示,MolR在4个数据集中表现最好,作者将MolR在分子性质预测方面的优异性能归因于,MolR在USPTO-479k上预训练,因此具备对于命题2对反应中心的感知。注意,反应中心通常由化学活性官能团组成,这些官能团对确定分子性质至关重要。
图编辑距离预测
图编辑距离(GED)是两个图之间相似性的度量,定义为将一个图转换为另一个图的最小编辑操作数。这项任务根据两个分子图的嵌入来预测它们之间的GED,目的是显示学习到的分子嵌入是否能够保持分子之间的结构相似性。结果如下表:
embedding可视化
为了直观展示分子embedding空间,论文使用预训练的MolR-GCN输出BBBP数据集中的分子embedding,然后使用t-SNE可视化这些向量,如下图。
在图a中,分子根据渗透性被着色,可以发现两个非渗透性的分子集群,这代表MolR可以捕获感兴趣的分子性质。
在图b中,分子根据GED着色,首先从BBBP中随机选择分子(No.1196)遍历BBBP中的分子,与No.1196号分子计算GED,从结果看出,在嵌入空间中,与1196号分子(橙色表示)结构相似的分子同样在距离上很相近,而与1196号分子(用红色表示)结构不同的分子也与它距离很远。结果表明,MolR可以很好地捕捉分子间的结构相似性。
在图c中,分子根据其大小(非氢原子数量)着色,可见嵌入空间被完美分割为小分子区域(上半部分)和大分子区域(下半部分)。换句话说,二维嵌入空间的纵轴表征了分子的大小。
在图d中,可以发现,水平轴实际上与最小环(环中不包含其他环)数量相关。无环分子(蓝色)仅仅在左侧的集群里,单环分子(黄色)仅仅在左侧与中间的集群里,双环分子(橙色)基本上位于中间的集群中,而右边的集群主要由2个以上环的分子组成(红色)。
解释具体的化学反应类型
论文以醇氧化和醛氧化为例,说明了MolR编码化学反应的过程,结果如下图:
醇氧化和醛氧化的反应模板分别为:
R
−
C
H
2
O
H
+
O
2
→
R
−
C
H
O
+
H
2
O
R-CH_{2}OH+O_{2}\rightarrow R-CHO+H_{2}O
R−CH2OH+O2→R−CHO+H2O
R
−
C
H
O
+
O
2
→
R
−
C
O
O
H
R-CHO+O_{2}\rightarrow R-COOH
R−CHO+O2→R−COOH作者使用预训练的MolR-GCN模型输出乙醇(
C
H
3
C
H
2
O
H
CH_{3}CH_{2}OH
CH3CH2OH),1-辛醇(
C
H
3
(
C
H
2
)
7
O
H
CH_{3}(CH_{2})_{7}OH
CH3(CH2)7OH),乙二醇(
(
C
H
2
O
H
)
2
(CH_{2}OH)_{2}
(CH2OH)2)以及对应的醛和羧酸的嵌入物,然后使用主成分分析PCA对其进行可视化。
从上图可以看出(对应橙色与红色的箭头): h C H 3 C H O − h C H 3 C H 2 O H ≈ h C H 3 ( C H 2 ) 6 C H O − h C H 3 ( C H 2 ) 7 O H h_{CH_{3}CHO}-h_{CH_{3}CH_{2}OH}\approx h_{CH_{3}(CH_{2})_{6}CHO}-h_{CH_{3}(CH_{2})_{7}OH} hCH3CHO−hCH3CH2OH≈hCH3(CH2)6CHO−hCH3(CH2)7OH h C H 3 C O O H − h C H 3 C H O ≈ h C H 3 ( C H 2 ) 6 C O O H − h C H 3 ( C H 2 ) 6 C H O h_{CH_{3}COOH}-h_{CH_{3}CHO}\approx h_{CH_{3}(CH_{2})_{6}COOH}-h_{CH_{3}(CH_{2})_{6}CHO} hCH3COOH−hCH3CHO≈hCH3(CH2)6COOH−hCH3(CH2)6CHO注意到,蓝色箭头的长度大约是相应红色或橙色箭头的两倍,这是由于 ( C H 2 O H ) 2 / ( C H 2 C H O ) 2 (CH_{2}OH)_{2}/(CH_{2}CHO)_{2} (CH2OH)2/(CH2CHO)2有两个羟基/醛被氧化。
总结
在这项工作中,作者使用消息传递模型作为分子编码器,并使用化学反应来协助学习分子表征,强制让反应物的嵌入总和等于生成物的嵌入总和。该模型能够学习反应模板,这是提高模型泛化能力的关键。该模型能够胜任广泛的下游任务,可视化的结果表明,学习到的嵌入是有组织的以及带有反应感知的。