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

题目：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，本节将介绍其神经网络架构。

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

分子分割

将以SMILES格式提供的分子分割成一系列化学子结构。这种分割通过BRICS方法执行，该方法能够有效分割出复杂分子中的关键子结构。

特权子结构识别

使用完整分子表示作为查询，子结构表示作为关键词，基于注意力机制识别特权子结构。具体地，采用GIN编码器来学习分子的表示：

对化学子结构，利用另一个GIN子编码器获得其表示：

使用完整分子表示作为查询，子结构表示作为关键词，基于注意力机制识别特权子结构。具体地，采用GIN编码器来学习分子的表示：

其中，矩阵是可学习的线性变换，用于增强表达能力，是表示维度。考虑到任务头需要特权子结构作为输入，因而可以直接使用其表示：

其中，为子结构表示堆叠而成的矩阵。另一方面，环境头需要可变结构。为实现这一点，可以翻转注意力：

这里，是可变结构表示，作为后续环境头的输入。

双头

任务头和环境头分别利用特权子结构和可变结构的表示进行分类。具体地，任务头使用多层感知机（Multi-Layer Perceptron，MLP）生成预测：

同样，环境头获取软预测和最终预测：

根据定理一，环境的划分原则是真实标签和预测间的一致性。因此，和。

优化目标

给定环境划分，可以定义每个环境的风险：

其中，是用于分类的交叉熵。为满足准则一，即环境划分最大程度违反子结构不变约束，可以固定分子性质预测器，并基于软分配最大化不变风险来优化环境头：

这里，是与相对应维度的元素。请注意，该公式中不包括经验风险。为防止平凡解，停止环境头与识别器间的梯度传播，仅优化以使其专注于可变信息。 为确保准则二中需要的类分布公平性，基于环境中的类比例重加权风险：

与定理二一致，的理想值等于1。实际中，它用于动态平衡不同环境中的类分布。通过对较大值施加更大的惩罚，增强特权子结构识别，从而调整类分布。因此，分子性质预测器的损失函数是：

其中，第一项代表重加权的经验风险，第二项是不变风险，为超参数。

5. 实验

本文在八个基准数据集上评估了所提MILI的有效性。具体地，来自Open Graph Benchmark（OGB）的BACE和BBBP数据集分别关注与人类-分泌酶1的结合亲和力及脑血屏障穿透性，并根据骨架差异划分训练/验证/测试集。其他六个数据集由DrugOOD提供，预测药物靶点结合亲和力。DrugOOD在IC50和EC50测量上利用测定、骨架和大小三种划分策略，创建了六个数据集 IC50/EC50-assay/scaffold/size。对于所有数据集，使用标准训练/验证/测试集划分。需要注意的是，只有DrugOOD的六个数据集具有环境标签。由于上述性质预测任务都为分类问题，遵循以往工作计算ROC-AUC分数。

OOD泛化

下表展示了在不同随机种子下5次独立试验的平均值和标准差结果。

从这些结果中，可以得出以下结论：（1）MILI在所有数据集上均超越所有基线，证明了所提的统一分子不变学习框架显著增强了分子表示学习的泛化能力。（2）MILI相对于MoleOOD的性能优势可以归因于环境划分准则和相应的双头图神经网络，确保环境划分和分子不变学习相互加强。（3）OOD基线相对于ERM的提升突显了在分子表示学习中考虑OOD场景的重要性。没有额外的设计，神经网络易利用虚假相关信息。此外，相对欧几里得数据的基线方法，图OOD基线方法的性能提升说明了图学习适用于捕捉分子中的复杂模式。

特权子结构识别

为加深对MILI的理解，本小节展示了在BBBP数据集上识别出的三种特权子结构，如下图所示。这些子结构根据学习到的权重排序，分别以红色、蓝色和紫色高亮显示。

聚焦于脑血屏障穿透，BBBP数据集根据分子渗透脑细胞膜的能力对分子进行分类。值得注意的是，被标记为红色的最重要的功能团（哌嗪），对于帮助分子通过脑血屏障至关重要。哌嗪的存在增强了脑血屏障的穿越能力，有利于中枢神经系统的进入。此外，以蓝色高亮的（乙酰胺）在一定程度上支持脑血屏障穿透，利用其酰胺基团进行氢键结合。尽管如此，由于其富含氮的环状结构，哌嗪在脑血屏障穿透方面优于乙酰胺，在排名中获得更高的重要性也与这一先验知识相符。这些发现说明了MILI在识别和分析对药物功效和安全性至关重要的子结构方面的优秀能力。同时，MILI能够适应分布偏移以确保药物是稳健和可靠的，有望成为现代药物研发中的有效工具。 更多实验，请参考原文。

参考文献

[1] Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George EDahl. 2017. Neural message passing for quantum chemistry. In ICML. 1263–1272. [2] Tianyi Zhao, Yang Hu, Linda R Valsdottir, Tianyi Zang, and Jiajie Peng. 2021. Identifying drug–target interactions based on graph convolutional network and deep neural network. Briefings in bioinformatics 22, 2 (2021), 2141–2150. [3] Yuanfeng Ji, Lu Zhang, Jiaxiang Wu, Bingzhe Wu, Lanqing Li, Long-Kai Huang, Tingyang Xu, Yu Rong, Jie Ren, Ding Xue, et al . 2023. Drugood: Out-of-distribution dataset curator and benchmark for ai-aided drug discovery–a focus on affinity prediction problems with noise annotations. In AAAI, Vol. 37. 8023–8031. [4] Yongqiang Chen, Yonggang Zhang, Yatao Bian, Han Yang, MA Kaili, Binghui Xie, Tongliang Liu, Bo Han, and James Cheng. 2022. Learning causally invariant representations for out-of-distribution generalization on graphs. In NeurIPS, Vol. 35. 22131–22148. [5] Shaohua Fan, Xiao Wang, Yanhu Mo, Chuan Shi, and Jian Tang. 2022. Debiasing graph neural networks via learning disentangled causal substructure. In NeurIPS, Vol. 35. 24934–24946. [6] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. 2019. Invariant risk minimization. arXiv preprint arXiv:1907.02893 (2019). [7] Yingxin Wu, Xiang Wang, An Zhang, Xiangnan He, and Tat-Seng Chua. 2021. Discovering Invariant Rationales for Graph Neural Networks. In ICLR. [8] Haoyang Li, Ziwei Zhang, Xin Wang, and Wenwu Zhu. 2022. Learning invariant graph representations for out-of-distribution generalization. In NeurIPS, Vol. 35. 11828–11841. [9] Chuleeporn Phanus-Umporn, Watshara Shoombuatong, Veda Prachayasittikul, Nuttapat Anuwongcharoen, and Chanin Nantasenamat. 2018. Privileged sub- structures for anti-sickling activity via cheminformatic analysis. RSC advances 8, 11 (2018), 5920–5935. [10] Elliot Creager, Jörn-Henrik Jacobsen, and Richard Zemel. 2021. Environment inference for invariant learning. In ICML. 2189–2200.