KDD 24 | 基于特权结构识别的分子不变学习

深度学习与图网络 · 公众号 · · 2024-07-24 14:39

正文

题目：Advancing Molecule Invariant Representation via Privileged Substructure Identification

会议：KDD 2024

图神经网络通过将分子建模为图，即视原子为节点，化学键为边，革新了分子表示学习。尽管取得了一定突破，它们在分布外场景中仍面临挑战。例如，具有相同性质的分子在大小或骨架变化时，可能给出错误的预测结果。一些研究尝试利用图不变学习来缓解该问题，即通过惩罚不同环境下的预测差异来学习不变表示。然而，分子有其特殊的领域知识： 形成特权子结构的核心官能团主导了分子的性质，并在分布变化中保持不变 。因此，怎样融合这一先验知识？并确保其与传统图不变学习兼容？

1. 背景

分子作为物质的基本组成部分，在科学探索和药物发现中发挥着关键作用。为此，分子表示学习（Molecule Representation Learning，MRL）成为一个重要的研究领域，即将复杂的分子嵌入至向量表示中。最近，图神经网络（Graph Neural Networks，GNNs）利用分子图来学习这些表示，在分子性质预测[1]和潜在药物识别[2]等任务中展现出了优越性能。

尽管取得了一定突破，它们通常依赖于一个基本假设，即分子是从一致环境中独立且同分布地采样的。然而，真实场景中经常出现环境变化和分布偏移。例如，在药物再利用中，在某些条件下筛选的分子常需要重新评估以对抗全新疾病或生物靶点。现有基于GNN的MRL方法在这些分布外（Out-Of-Distribution，OOD）场景中性能显著下降[3]，迫切需要提高其泛化能力。

最近缓解GNN OOD问题的主流方法是图不变学习（Graph Invariant Learning，GIL）[4,5]：假设因果子图在各个环境中不变，而环境子图可变；通过惩罚不同环境下的预测差异，模型可以捕捉因果因素而非虚假相关因素。在这里，一个开放问题是如何有效地确定环境划分。现有方法已探索了各种策略，包括预定义划分[6]、图增强[7]以及设计专门用于环境划分的附加模型[8]。

然而，将GIL应用于分子表示学习涉及三个关键点。首先，整合领域知识是必不可少的。特权子结构[9]是决定分子活性的核心官能团，例如阿司匹林的镇痛特性可以归因于其酯官能团。这表明MRL中的不变子图应是这些具有化学含义的特权子结构，但大多数现有方法学习任意子图，忽视了这一先验知识。其次，探索相关理论指导GIL的环境划分是必要的。使用图增强可能会生成无意义的分子，难以为领域专家提供新洞见。此外，现有方法中预定义或学习的划分与下游GIL无关，无法在理论上确保其与GIL兼容。最后，设计环境划分和下游预测的统一模型也至关重要。现有方法通常将环境划分和下游预测视为两阶段过程，导致其取得次优性能。

为解决上述问题，本文提出了一个名为MILI的全新框架，通过特权子结构识别来推进分子不变学习（Molecule Invariant Learning）。为整合领域知识，首先基于特权子结构识别形式化分子不变学习，并引入子结构不变约束（Substructure Invariant Constraint，SIC）。然后，从理论上建立了两个环境划分标准，以确保其增强分子不变学习：环境应基于ground truth和可变结构下游预测的一致性进行划分，并（1）最大程度违反SIC和（2）保持类分布公平。为满足这些标准，本文进一步设计了一个双头图神经网络。其中，共享识别器识别特权子结构，随后任务头和环境头使用特权子结构和可变结构进行下游预测。根据所提标准，环境划分通过最大化不变风险来违反SIC，并重加权经验风险来增强类分布公平性。最终，这一统一框架允许环境划分和分子不变学习相互增强。在不同数据集上的广泛实验验证了所提MILI的有效性。

2. 分子不变学习

本节定义分子上的OOD泛化问题，然后基于特权子结构识别扩展不变学习框架。

分子的OOD泛化

代表分子图的随机变量为，其中节点为原子，边为化学键。记是分子图空间，是标签空间。考虑数据集，其中且。在实际应用中，数据集通常来源于多个环境。这里，表示来自环境的数据集，而表示训练数据中的环境空间。

定义 : 是所有可能环境的空间，是分子表示空间。假设预测器可以分解为，其中为编码器，将分子映射到表示空间，而为分类器，通过线性映射将表示映射到的对数空间。分子的OOD泛化目标是找到一个最优预测器，使其在所有环境中都表现良好：

这里，表示在环境上的经验风险，是损失函数。

分子图及其对应标签的联合分布表示为。分布偏移指训练数据中的联合分布与测试数据中的联合分布不同。

分子不变学习

分子图由决定其性质的特权子结构特征化，这表明这些特权子结构与相应标签间的关系在所有环境中都是不变的。的补集记为，代表随环境变化的结构。遵循不变学习相关文献[10]，本节为分子不变学习定义子结构不变约束（Substructure Invariance Constraint，SIC）。

定义 (子结构不变约束) 假设所学的最优识别器能够识别出分子图中的特权子结构。那么，分子不变表示需满足以下约束：

为避免平凡解，此约束被整合为训练目标中的一个正则项。类似于不变风险最小化（Invariant Risk Minimization，IRM）理论，将设为每个输出维度的常数标量乘子1.0。分子不变学习的目标函数可写为如下形式：

显然，该公式需要预定义的环境标签，而其获取在实际场景中存在挑战。此外，即使有可用的环境标签，并不意味着它们适用于分子不变学习。

3. 分子不变学习的环境划分准则

在上述形式化的基础上，本节旨在推导出有益于分子不变学习的环境划分准则。直观地，这些环境应揭示可变特征的变化。因此，如果环境划分仅基于可变结构，它可以实现任意变化。这里，使用标识符表示所识别的特权子结构的补集。

假设存在另一个环境预测器使用可变结构预测标签。与遵循SIC的分子不变表示不同，环境预测器的结果会违反SIC。

定理一 对于仅依赖可变结构的最优环境预测器，将其预测记为，真实标签记为。如果环境通过以下方式划分：

那么子结构不变约束将被最大程度地违反。其中，函数用于判断两随机变量的等价性。

然而，这一理想结果要求环境预测器只利用可变结构。在没有先验知识以供准确提取可变结构的情形下，此要求难以满足。因此，在实际实现中，所学环境划分应尽可能地违反SIC。

准则一
环境根据真实标签与可学习环境预测器预测值间的一致性进行划分，即。这种环境划分应被优化至最大程度违反子结构不变约束。

此外，如果环境划分由最优环境预测器和真实标签确定，有分子不变表示。因此，可以推导出以下定理。

定理二 对于由最优环境预测器和真实标签确定出的环境划分，以下等式

对任何和任何都成立。

此定理建立了环境划分与类分布间的关系，即类分布对环境划分是公平的。为达到理想结果，进一步引入环境划分的第二个优化准则。

准则二
环境基于真实标签和可学习环境预测器的预测值来划分，这种环境划分应被优化至不同环境中类分布公平。

4. 模型设计

基于所提环境划分准则，本文提出了基于特权子结构识别的分子不变学习模型MILI，本节将介绍其神经网络架构。

回顾之前提出的分子不变学习框架：分子性质预测器由两部分组成，识别器用于识别特权子结构，而用于预测分子性质。在第三节中，环境划分依赖于环境预测器，其中是所识别特权子结构的补集，而为基于可变结构的预测器，这自然而然形成了一个双头图神经网络。具体地，共享部分为特权子结构识别器，和分别为任务头和环境头。整体框架如上图所示，每个模块的实现介绍如下。