NeurIPS22 | 分子图分布外泛化与因果不变性

点击蓝字

关注我们

AI TIME欢迎每一位AI爱好者的加入！

分布外泛化（Out-of-Distribution Generalization，简称OOD泛化）旨在解决训练数据与测试数据分布不一致的问题，即模型需要基于有限观测数据学习出背后稳健的规律与关联性，从而泛化到未知数据分布。OOD泛化是目前机器学习领域的研究热点，也是各种实际场景（如推荐系统、自动驾驶、药物设计等领域）亟待解决的关键问题。

在上一篇文章《漫谈图上的分布外泛化：不变性视角下的求解》中我们围绕一篇ICLR22的工作EERM[1]介绍了图上分布外泛化问题和基于因果不变性的求解。本文将进一步介绍另一篇最近刚发表于NeurIPS22的工作，考虑分子图性质预测场景下的分布外泛化问题，这项工作也入选了NeurIPS Spotlight Presentation（前5%）。

论文题目：Learning Substructure Invariance for Out-of-Distribution Molecular Representations

论文链接：

https://openreview.net/pdf?id=2nWUNTnFijm

代码链接：

https://github.com/yangnianzu0515/MoleOOD

前言

预测分子特性对于药物发现和材料设计等许多相关应用相当重要，也是近期图机器学习在AI4Science领域的热门应用之一。分子性质预测（molecular property prediction）本质是一个图分类/回归问题，即输入样本为一个图，模型需要输出图的标签。与常规的图表示学习类似，分子表示学习旨在将分子图嵌入到隐空间中的向量中，以此表征为基础再用于各种下游预测任务。然而，现有的分子表示学习方法主要基于一个假设：训练和测试的分子数据来自同样的分布。但现实世界的环境是充满变化的，这就需要分子表示学习能够有效应对数据分布的变化。

这项工作的贡献总结如下：

• 首次将分子表示学习的分布外泛化问题做了形式化定义，并探索了因果不变性原理与分子图结合的可能性

• 提出了一种能够有效应对分布变化的通用分子表示学习框架叫做 MoleOOD，该框架理论上可以使用任意现有的分子表示学习模型（如各类图神经网络）作为骨干模型以提升他们的泛化能力

• MoleOOD在多个公开数据集取得了SOTA效果。

分子图预测任务的分布外泛化

分子图性质预测问题本质上可以看作一个图分类/回归问题，即每个分子 Gi 是一个图样本（其中每个原子是一个节点，原子之间的化学键组成了连边），待预测的分子性质就是样本的标签 yi 。使用 G 和 y 表示图样本和标签对应的随机变量。

这里的 f 表示分子图分类器， l 表示损失函数，E 表示所有环境的集合。上式的直观理解就是我们希望学习到的分类器能够最小化最坏情况下的分类误差，即不管测试数据来自什么分布分类器都能给出满意的结果。然而这一问题的难点在于，由于训练数据的有限性，往往只能涵盖 E 中有限的环境（即训练数据的分布无法实现对全局的覆盖），而我们希望的是模型在所有的环境下都表现好，即具备良好的分布外泛化能力。

分子子结构的不变性

要实现有效的分布外泛化，一种可行的思路就是利用不变性原理 (Invariance Principle)，从有限观测数据中学习到数据背后关于环境的不变性，从而泛化到新的数据分布上。在《漫谈图上的分布外泛化：不变性视角下的求解》中我们对这一原理做了详实的说明，这里再简单提一下其核心思想。基于（因果）不变性学习的方法目标是学习一个对分布变化鲁棒的分类器，即一个相对环境变化不发生改变的从样本到标签的映射关系。具体的，希望分类器能够学习到输入样本的不变表征(invariant representation)，它满足以下两个性质：

• 预测充分性(Sufficiency)：表征包含的信息足够准确预测出类别标签。

• 环境不变性(Invariance)：对于任意环境e，基于表征给出的预测分布应该保持不变。

在分子图性质预测这一具体场景下，要如何定义或者建模这种不变性呢？我们的建模动机来源于一个已经被来自生物信息学、药理学等领域的相关工作提出的观察发现：分子的生物化学特性通常与一些特殊的分子子结构相关。下面我们又给出了两组具体的例子：

• 左图例子：可以看到他们的分子骨架(scaffold)不同，一个是由3个碳原子构成的环，而另一个是由6个碳原子构成的环。而分子骨架可以被视作一种环境划分的依据，所以这两个分子来自于不同的环境，即来自不同的分布，但他们都含有一个共同的子结构羟基(-OH)，所以他们都易溶于水。通过这个例子，我们可以知道羟基和水溶性之间存在一种相对于环境稳定不变的关联性。
  

• 右图例子：右图中两个分子的尺寸(size)存在明显差异，分子尺寸可以理解为所含原子数目，而分子尺寸也可以作为环境划分的依据，所以这两个分子也可以看作来自于不同的环境(分布)。类似地，我们可以发现这两个分子共同含有的子结构羧基(-COOH)也和水溶性之间存在一种相对于环境稳定不变的关联性。

我们把分子中这类相对于环境能稳定不变地决定分子性质的子结构叫做不变(invariant)子结构，而分子中其余的子结构（即与分子性质的关联在跨环境时会发生变化）则称作虚假关联(spurious correlation)子结构。

基于上述观察，一个可行的方法就是利用这种不变性，引导分类器学习到分子图中的不变子结构对应的表征，从而实现分布外泛化的目的。

学习子图结构不变性的分布外泛化求解

这项工作的模型包含了一个分子编码器 Φ 用来学习分子的不变表征以及一个预测器 ω 用以做出最后的预测。回想我们之前所说的，我们希望给定一个分子G, 分子编码器能够利用其中相对于环境不变的子结构学习得到一个分子表征 Φ(G)。我们的目标1是最大化 Φ(G) 对于 y 的预测能力，这可以通过最大化 Φ(G) 和 y 之间的互信息来实现。同时，目标2是希望分子表征尽可能提取分子中的不变子结构信息，即在给定 Φ(G) 之后的 y 和 e 条件独立，这可以通过最小化他们之间的互信息实现。为了展示的方便，我们用 z 表示 Φ(G)。基于上述讨论，我们的学习目标可以写为

如果我们现在再把ω 和 Φ 的输出分别看成分布 qθ(z|G) 和 qθ(y|z)，公式(2)又可以进一步写成：

上式的目标需要用到明确的环境信息，即我们需要知道每个分子对应的环境。在实际中，由于给分子标注环境标签非常麻烦，因此在许多情况下人工标注的环境标签并不能获取。尽管有些数据集会提供事先标注好的环境标签（例如基于分子骨架信息scaffold的环境划分），但对应的环境数目过于庞大造成每种环境对应的观测数据过于稀疏，这对训练一个相对环境鲁棒的分类器并不友好。所以，这项工作引入了一个环境推理(environment-inference)模型，用于从数据中学习归纳出环境信息，将训练集中的分子划分到k个不同的环境上（k可以人为调整）。这种做法和ICLR22工作[1]中利用k个图结构生成器来生成不同环境的数据是殊途同归的，即都是额外引入了一个模型“模拟”产生环境，以数据驱动的方式自适应的提供不同环境下的“监督”信息。

环境推理模型为数据集中的每个分子都估计了对应的环境标签，我们将其带入(3)式所定义的针对分子编码器和预测器的目标函数，可以得到如下基于不变性的优化目标：

我们在论文中给出了两个定理表明最小化公式(5)里的第一项和第二项分别对应着使得学到的分子表征 z 满足不变性原理的环境不变性(Invariance)和预测充分性(Sufficiency)。以下是提出的 MoleOOD 数据流程图：

模型训练阶段采用类似EM算法的迭代更新：

实验结果

这项工作在Open Graph Benchmark (OGB) 和 DrugOOD Benchmark提供的共 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 的表现要好。此外，我们还做了一些消融实验来分析模型各个组件对最终表现的影响，比如探究提出的环境推理模型的有效性，具体可以阅读论文。

总结与讨论

最后做一个简单的总结和延伸性的讨论。ICLR22的工作[1]和本文的工作都是关注在图结构数据的OOD问题，前者处理了节点级预测任务，后者是图级别预测任务。在方法层面，他们都利用了不变学习的思想，引导模型学习相对环境不变的输入-输出映射关系，前者引入了k个数据生成器来模拟产生不同环境下的数据，后者引入了一个环境推理模型来将数据样本划分为k个环境。

读完这两个工作的介绍，大家可能也会发现他们的一些共性，以及分析和求解分布外泛化问题的一般性思路。笔者认为分布外泛化问题的关键在于：1）发现和定义数据生成过程中的不变性因素；2）基于不变性设计有效的算法引导模型实现分布外泛化；3）设计/采用合理的实验评估方法验证模型在OOD场景下的有效性。这类问题处理起来具有较高的灵活度，不管在问题定义、解决方案、实验评估上可以开脑洞的空间都很大，目前极具探索空间。感兴趣的读者可以进一步阅读另一篇NeurIPS22工作[2]的详细解读：

https://zhuanlan.zhihu.com/p/603012617

这几个工作都是在一般方法层面给出了形式化的建模和求解，落到具体的模型设计层面，仍然有很大的创新空间，例如采用更高级的编码器和环境推理的模型，进一步优化设计目标函数，或是考虑特定的分布偏移等等。

[1] Qitian Wu, et al., Handling Distribution Shifts on Graphs: An Invariance Perspective, ICLR22.

[2] Chenxiao Yang, et al., Towards Out-of-Distribution Sequential Event Prediction: A Causal Treatment, NeurIPS22.

往期精彩文章推荐

记得关注我们呀！每天都有新知识！

 关于AI TIME 

AI TIME源起于2019年，旨在发扬科学思辨精神，邀请各界人士对人工智能理论、算法和场景应用的本质问题进行探索，加强思想碰撞，链接全球AI学者、行业专家和爱好者，希望以辩论的形式，探讨人工智能和人类未来之间的矛盾，探索人工智能领域的未来。

迄今为止，AI TIME已经邀请了900多位海内外讲者，举办了逾450场活动，超500万人次观看

我知道你

在看

哦

~

点击 阅读原文 查看回放！