近年来，分子表示学习(Molecular Representation Learning)获得了广泛关注，目前已有方法已在分子特性预测和靶点识别等任务中均有出色表现。然而，现有方法的模型设计或实验评估过程都是基于训练和测试数据是独立同分布的这样的假设。然而在实际应用中，这样的假设往往并不成立，因为测试分子极有可能来自模型训练阶段未见过的数据分布，这将导致严重的性能下降。
在本工作中，受“不同环境下(例如不同分子骨架、不同分子尺寸等)的分子们的生物化学性质通常与某些分子子结构相关”这样一个现象的启发，我们提出了一个名为 MoleOOD 的新分子表示学习框架，以增强分子表示学习模型对这种分布变化的鲁棒性。具体来说，我们引入了一个环境推理模型，以完全数据驱动的方式识别影响数据生成过程的潜在因素，即环境变量。我们还提出了一个新的学习目标来指导分子编码器识别、编码这些“跨环境性质稳定”的子结构。我们在十个真实数据集上进行实验，实验结果结果表明即使缺少事先人为标注好的环境标签，在各种分布外场景下，我们的模型也具有比现有方法更强的泛化能力。
论文题目：Learning Substructure Invariance for Out-of-Distribution Molecular Representations
论文链接：
https://openreview.net/pdf?id=2nWUNTnFijm
代码链接：
https://github.com/yangnianzu0515/MoleOOD
预测分子特性对于药物发现和材料设计等许多相关应用相当重要，这类任务通常需要来自化学和药理学等领域的专家付出巨大的努力。近年来，人们提出一系列基于机器学习的分子表示学习方法来加快这些任务的进程，并且在基于机器学习的分子表示学习领域取得了坚实的进展。一般来说，基于机器学习的分子表示学习任务旨在将分子嵌入到隐空间中的向量中，以此表征为基础再用于各种下游任务，例如靶点识别(target identification)、逆合成分析(retrosynthetic analysis)、虚拟筛选(virtual screening)等等。
然而，现有的分子表示学习方法主要基于一个假设：训练和测试的分子数据服从独立同分布(i.i.d.)。但现实世界中的分子数据的分布通常是不确定的，这就需要现有的分子表示学习能够有效地应对分布变化。
本文我们将介绍一项被 Thirty-sixth Conference on Neural Information Processing Systems (NeurIPS’22) 会议接收的新工作，我们于该工作中提出了一种名为 MoleOOD 的模型以解决上述问题，该工作有幸被选为 Spotlight presentation。
该工作的亮点总结如下：
- 我们首次探索了不变性原理与分子领域知识结合的可能性，并提出了一种能够有效应对分布变化的通用分子表示学习框架叫做 MoleOOD。
- 该框架理论上可以使用任意现有的分子表示学习模型作为骨干模型以提升他们的泛化能力。此外，MoleOOD 能够对训练分子数据自行进行环境划分，并不依赖人为标注的环境标签，所以 MoleOOD 也极具实用性。
- 通过在十个公开数据集上的实验，我们提出的 MoleOOD 的有效性也得到了充分验证。
一、背景
分布外泛化 (Out-of-Distribution Generalization)：深度神经网络当面临分布变化(distribution shifts)的时候性能很容易就会显著下降，现很多工作都在致力于探索分布外泛化(Out-of-Distribution Generalization)问题。近期的很多相关研究都假设有一个潜在的环境变量e导致了训练数据和测试数据间的分布变化 [1,2,3,4]。一般情况下，任务目标都是给定一个输入样本x我们需要去预测其对应的类别标签y 。然后，这样的OoD问题可以被正式地表述为：
不变学习 (Invariant Learning)： 这是一个试图解决OoD泛化问题的一个新兴研究方向。这类方法提出去学习一个能够揭示输入样本和对应标签之间相对环境不变的一种关联的不变预测器。所以我们的模型就需要去学习得到满足不变性原理(Invariance Principle)的不变表征(invariant representation)。不变性原理的具体内容如下：
- 预测充分性(Sufficiency)：表征包含的信息足够准确预测出类别标签。
- 环境不变性(Invariance)：对于任意环境 ，预测器给出的预测分布应该保持不变。
分子表示学习 (Molecular Representation Learning)：现有的分子表示学习方法大致可以被归为两类。第一类是基于分子SMILES表达式的，这类方法用自然语言处理领域的一些模型来处理SMILES文本表示，例如Transformer、BERT。该类方法的一个局限性在于分子的结构信息很难通过SMILES提取出来。为此，第二类旨在于学习分子的结构信息。由于我们可以将一个分子化合物视为一张图G=(V,E) ，图中V是图上的节点集合，对应着构成分子的原子，而E则表示图中的边的集合，即对应着分子中的化学键，所以这类方法的主流是利用图神经网络来学习分子表征用于下游任务。
二、动机
我们的建模动机来源于一个已经被来自生物信息学、药理学、数据挖掘等领域的相关工作提出的观察发现 [5,6,7,8]：分子的生物化学特性通常与一些特殊的分子子结构相关。我们正是基于此先验知识来设计我们的模型。下面我们又给出了两组具体的例子：
先看左边这一组分子，可以看到他们的分子骨架(scaffold)不同，一个是由3个碳原子构成的环，而另一个是由6个碳原子构成的环。而分子骨架可以被视作一种环境划分的依据，所以这两个分子来自于不同的环境，即来自不同的分布，但他们都含有一个共同的子结构羟基(-OH)，所以他们都易溶于水。通过这个例子，我们可以知道羟基和水溶性之间存在一种相对于环境稳定不变的关联性。再看右边的一组分子，可以看到他们的分子尺寸(size)差异明显，分子尺寸可以理解为所含原子数目，而分子尺寸也可以作为环境划分的依据，所以这两个分子也是来自于不同的环境(分布)。
类似地，我们可以发现这两个分子共同含有的子结构羧基(-COOH)也和水溶性之间存在一种相对于环境稳定不变的关联性。我们把分子中这类相对于环境能稳定不变地决定分子性质的子结构叫做不变(invariant)子结构，而分子中其余和分子性质在跨环境时并不不变地稳定相关的则称作虚假(spurious)子结构。因此，一个可行的方法就是从跨环境的和分子性质具有不变关联性的子结构中学习关于某种性质的不变性，从而达到 OoD 泛化的目的。
三、方法
对于分子表示学习问题，我们提出了一种OoD泛化框架。下面我们先给出分子表示学习上的OoD泛化问题的具体问题定义。
问题定义
问题的难点在于训练数据只涵盖了 ε 中的有限的环境类别，然后我们希望的是模型在所有的环境上都表现好。
现有的分子表示学习方法不区分前文中定义过的 invariant 和 spurious 子结构。所以，和性质不相关的子结构和分子性质之间的虚假的关联将会被编码到学到的分子表征中。当在来自未见过的环境的数据上进行测试的时候，下游的分类器将很容易被这些虚假的关联误导。基于分子的生物化学特性通常与一些特殊子结构相关联的先验知识，我们希望我们的模型在训练的过程中能抑制虚假的相关性，并利用相对于环境不变的子结构，即这些子结构与跨环境的标签更稳定地相关，以学习不变的分子表征。
模型设计
在实际中，由于给分子标注环境标签非常麻烦，因此在许多情况下人工标注的环境标签并不能获取。当环境标签缺失时，我们可以根据分子骨架(scaffold)直接对分子进行环境划分。但这其实是不合理的，因为最终的总环境数目会很大。以 Open Graph Benchmark 发布的用于分子特性预测任务的数据集HIV为例，如果我们对于 OGB 使用 scaffold 将分子划分到不同的环境中的策略，假设我们直接将每种分子骨架看成一个环境，那么HIV中的 41127 个分子被划分到 19076 个环境中。它的环境数目要远比别的领域的 OoD 数据集的环境数目要大得多。例如 Camelyon17 是一个用于肿瘤检测的 OoD 数据集，它就只将 455954 个组织切片图像划分到了 5 个环境上。尽管有些数据集会提供事先标注好的环境标签，但他们的环境数目还是很大，这对现有的一些 OoD 方法并不友好。所以，我们提出设计一个环境推理(environment-inference)模型将训练集中的分子划分到一些总数相对而言小得多的环境上。我们将划分得到的环境数目设定为一个超参数k.
以下是我们提出的 MoleOOD 模型流程图：
模型训练阶段主要分成两步：
四、实验结果
我们在Open Graph Benchmark (OGB) [9] 和 DrugOOD Benchmark [10] 提供的共 10 个数据集上都进行了实验来验证我们提出的 MoleOOD 框架的有效性。
数据集：我们选择了 OGB 中的四个数据集 BACE，BBBP，SIDER 和 HIV，这四个数据集 OGB 是按照分子骨架进行划分的。剩余 6 个数据集来自 DrugOOD，即 IC50/EC50-assay/scaffold/size，DrugOOD 相比于 OGB 提供了更多的划分数据集的方式，即不仅仅可以按照分子骨架(scaffold)定义分布，也可以按照分子测定(assay)或者分子尺寸(size)。所以 DrugOOD 可以提供更全面的 OoD 泛化性能评估。
评价指标：与之前的分子表示学习相关工作一致，我们采用 ROC-AUC 为评价指标。
Baseline：理论上，任何现有的分子表示学习模型都可以嵌入到我们的框架中作为我们的骨干模型(backbone)来提升他们原本的泛化能力。我们选取了三种模型：GCN，GIN 和 GraphSAGE 作为骨干模型来验证我们的方法是否会给他们的泛化能力带来提升，并且我们还和他们的增加版模型进行了比较，即加上虚拟节点，也就是下面实验结果表格中的"+ virtual node"。此外，我们还和现有的六种 OoD 泛化方法进行了比较：ERM，IRM, DeepCoral, DANN, MixUp 和 GroupDro。
实验结果：下面两张表格分别展示了我们提出的 MoleOOD 模型以及 baselines 在 OGB 和 DrugOOD 上的表现。
可以看到我们的模型在 OGB 数据集上均给 baselines 带来显著提升，而在 DrugOOD 的六个数据集上，除了 IC50-size 上，MoleOOD 一致地比 baselines 的表现要好。此外，我们还做了一些消融实验来分析模型各个组件对最终表现的影响，比如探究提出的环境推理模型的有效性，这一部分的详细内容可以在论文正文部分看到。
五、总结
在该文章中，我们提出了一个通用框架，可以将任何现有的分子表示学习方法作为骨干模型，以提高它们对分布变化的泛化能力。具体来说，我们设计了一种新的学习目标及其等效的具体实现。我们还设计了一个环境推理模型能够自行对训练数据的环境重新划分，使得我们的方法不依赖于人为事先标注好的环境标签，这一点也使得我们的模型实用性极高。在十个公开数据集的实验结果表明，我们的模型对作为主干模型的各种现有分子表示学习方法带来了一致且显著的泛化能力的提升。此外，与现有的针对 OoD 泛化问题而设计的需要人为标注的环境标签作为额外输入的最先进模型相比，我们的模型表现相当甚至能表现得更加优越。
写在最后
We’re always open for possible collaborations and feel free to contact us. 欢迎一起交流与进步！
参考文献
[1] Rojas-Carulla, Mateo, et al. “Invariant models for causal transfer learning.” The Journal of Machine Learning Research 19.1 (2018): 1309-1342.
[2] Arjovsky, Martin, et al. “Invariant risk minimization.” arXiv preprint arXiv:1907.02893 (2019).
[3] Bühlmann, Peter. “Invariance, causality and robustness.” Statistical Science 35.3 (2020): 404-426.
[4 ]Wu, Qitian, et al. “Handling Distribution Shifts on Graphs: An Invariance Perspective.” International Conference on Learning Representations. 2022.
[5] Klekota, Justin, and Frederick P. Roth. “Chemical substructures that enrich for biological activity.” Bioinformatics 24.21 (2008): 2518-2525.
[6] Phanus-Umporn, Chuleeporn, et al. “Privileged substructures for anti-sickling activity via cheminformatic analysis.” RSC advances 8.11 (2018): 5920-5935.
[7] Zhu, Jiajing, et al. “DGDFS: Dependence guided discriminative feature selection for predicting adverse drug-drug interaction.” IEEE Transactions on Knowledge and Data Engineering 34.1 (2020): 271-285.
[8] Kao, Yu-Ting, et al. “A substructure-based screening approach to uncover N-nitrosamines in drug substances.” Journal of Food & Drug Analysis 30.1 (2022).
[9] Hu, Weihua, et al. “Open graph benchmark: Datasets for machine learning on graphs.” Advances in neural information processing systems 33 (2020): 22118-22133.
[10] Ji, Yuanfeng, et al. “DrugOOD: Out-of-Distribution (OOD) Dataset Curator and Benchmark for AI-aided Drug Discovery–A Focus on Affinity Prediction Problems with Noise Annotations.” arXiv preprint arXiv:2201.09637 (2022).
作者：杨念祖
文章来源：公众号【 sjtuThinklab】
快抢沙发！
