吉林大学提出PGR-MOOD模型，通过扩散模型检测分布外的分子图

尽管分子表示学习最近取得了显著进展，但其有效性是建立在训练图和测试图来自相同分布的假设上的。而在实际应用中，测试数据集通常与分布外(OOD)样本混合在一起，使得部署的模型难以做出准确的预测。在药物筛选或设计中对分子性质的错误估计会导致湿实验室资源的巨大浪费，并延迟新疗法的发现。传统的检测方法需要权衡OOD检测和分布内(ID)分类性能，因为它们共享相同的表示学习模型，这就降低了模型检测OOD样本的性能。相比之下，由于生成模型偏向于重建分布内（in-distribution，ID）训练样本，因此，生成模型可以使得OOD分子与现有分子的相似度降低，从而便于检测出OOD分子。

虽然在概念上很简单，但是将这个普通框架扩展到实际的检测应用程序仍然受到两个重大挑战的限制。首先，流行的基于欧氏距离的相似度度量没有考虑复杂的图结构。其次，涉及迭代去噪步骤的生成模型非常耗时。

2024年8月，吉林大学王鑫老师团队在人工智能顶级会议KDD上发表文章Optimizing OOD Detection in Molecular Graphs: A Novel Approach with Diffusion Models。

作者提出通过采用基于扩散模型的框架来检测分布外（out-of-distribution，OOD）分子，构建了一种用于分子OOD检测的原型图重构方法，称为PGR-MOOD（Prototypical Graph Reconstruction for Molecular OOD detection）。具体而言，PGR-MOOD主要有三个创新：1)根据输入和重构分子的离散边缘和连续节点特征，综合量化输入和重构分子匹配程度的有效度量；2)通过基于扩散模型的图生成器，用于构建符合分布内（in-distribution，ID）但远离OOD分布的原型图列表；3)一个高效且可扩展的OOD检测器，用于比较测试样本和预先构建的原型图之间的相似性。实验结果表明，PGR-MOOD的性能超越了现有模型。

作者首先考虑一个朴素模型GR-MOOD，验证重建方法用于分子OOD检测的可行性。OOD检测问题可以看做对样本的二分类（ID或OOD）问题。然而，分子图具有非欧几里得结构，其固有的复杂性给分子图OOD检测带来了两个重大挑战。首先，分子图的这种性质使得传统的相似性度量(例如，欧几里得距离)在量化原始图和重建图之间的接近程度方面效果较差。同时，不同的分子经常有着分布变化，包括结构和特征的变化，进一步影响了相似性的评估。这就引出了挑战1：确定一个有效的度量来评估原始输入和重建之间的相似性。更重要的是，扩散模型需要数百或数千个采样步骤来从正态标准分布降噪到生成新图，这引入了额外的复杂性。这就引出了挑战2：解决重建所需的扩散模型的额外复杂性。

如图1所示，对于挑战1，在寻找一种评估原始输入与其重建之间相似性的有效度量方面，PGR-MOOD采用融合（fused）Gromov-Wasserstein (FGW)距离，利用分子图的结构信息和特征信息来增强对其匹配程度的度量。对于挑战2，PGR-MOOD建议创建一系列更接近ID样本、远离OOD样本的原型图，减少重建每个测试图的需要，只需将其与准备好的原型图的相似度进行比较。PGR-MOOD的创新有三个方面：强相似函数、原型图生成器和高效可扩展的OOD检测器。

图1 PGR-MOOD结构图

基于扩散模型的生成方法包括正向扩散过程和反向去噪过程。在正向过程中，模型逐步向原始数据中加入噪声，直到达到标准正态分布。在反向过程中，模型学习打分函数(即神经网络)，以相同的步长去除扰动噪声。给定一个图G，扩散模型可以用连续时间对扩散轨迹进行索引，通过不断加入噪声，使得G迭代到近似服从正态分布。这个过程可以通过正向随机微分方程(SDE)刻画。而在逆向过程中，为了生成符合需要的分布的图，从迭代到近似正态分布的G开始，利用反向SDE，对节点特征和图结构进行去噪。

受到生成扩散模型方法的启发，作者首先设计了一个基础的基于扩散模型的重构模型，GR-MOOD，用于分子图OOD检测。GR-MOOD在大规模复合数据集(如QM9或ZINC)上进行预训练和微调，在获取重构图后，利用在ID数据集上训练良好的GNN对的特征和结构信息进行编码，将二者之间的余弦相似度作为OOD判断评分。由于GR-MOOD被训练来重建与ID分布对齐的图，OOD样本由于其与ID分布的固有不相似性，在处理时通常会经历较差的重建。这种差异被量化为较低的判断分数，这表明存在一个良好的样本。这一机制突出了基于扩散模型的重构方法在识别不符合预期分布的图中的关键作用，从而为区分ID和OOD样本提供了定量依据。

尽管GR-MOOD具有直观的前景，但其在时间效率和重建质量测量方面存在不可忽视的局限性。首先，GR-MOOD的主要约束是由于分子图固有的结构复杂性，这种复杂性要求扩散模型采取大量去噪步骤来完成重建，以牺牲效率为代价提高模型性能。更糟糕的是，重复每个分子的生成过程使测试阶段的规模变得具有挑战性，因为测试阶段必须筛选大量候选分子。其次，另一个问题与模型中使用的相似函数的充分性有关，传统的余弦相似度等度量侧重于量化两个向量之间的角度，而不适用于图等非欧几里得数据。

实际上，测量图之间的相似度相当于计算图之间的匹配度，匹配度越高，图之间的相似度越高。FGW距离对图之间的相似度测量特别适用。它在以节点特征为代价的最优传输(OT)距离和拓扑结构之间的Gromov-Wasserstein (GW)距离之间实现了平衡。具体地说，FGW将与拓扑和节点特征相关联的图视为概率分布。它允许计算具有最优耦合的两个分布之间的代价，作为图之间的距离度量。对于用OT格式表示的两个图，在直接识别图之间的结构差异和特征差异方面表现出最佳性能。

由于GR-MOOD的朴素扩散模型重建的图更倾向于输入样本的分布，而不是遵循训练阶段学习到的分布，这将误导检测器对OOD样品的判断。为了解决这一挑战，我们提出了一个原型图生成器，它生成满足以下两个属性的原型图:（1）原型图与现有的ID图尽量相似；（2）原型图与现有的OOD图尽量不相似。因此，优化目标是生成一个原型图，使之接近ID图，远离OOD图，对应的损失函数分别用LID和LOOD表示。LID为指导生成器的去噪过程而制定，同时，为了符合属性映射，还需要引入损失函数LOOD，以增进OOD样本之间的距离。

注意到，OOD图在训练阶段无法直接获取，排除了直接制定，因此，需要通过对预训练的扩散模型添加扰动的参数权重，来生成OOD图。扰动可以诱导去噪轨迹偏离原始数据分布，从而使扩散模型生成在训练阶段生成OOD样本，并计算LOOD损失。最终的引导损失函数Lguide=LID+LOOD，被用来指导原型图生成器的训练。这样生成的原型图可以看作是ID图和OOD图的重构，但比GR-MOOD生成的重构具有更好的判别性。

扩散模型在测试阶段需要大量的时间和内存资源，因为它们需要为每个输入生成一个重构的图。为了减轻这种计算负担，PGR-MOOD通过在训练阶段准备原型图而消除了在测试阶段重构图的必要性。接着，PGR-MOOD对每一个新的测试样品进行相似度测量。采用最大化FGW相似度作为OOD检测的最终判断分数。这种策略能够有效降低扩散模型的时间复杂度和空间复杂度。

作者将PGR-MOOD与一些具有代表性的方法进行了比较，如表1所示。将AUROC, AUPR和FPR95分数作为对比指标。AUROC和AUPR越高越好。FPR95表示在TPR=95%时的FPR的值，这个值越低越好。实验采用了DrugOOD-IC50和DrugOOD-EC50数据集中在骨架（scaffold），大小（size），测定（assay）出现OOD的各3个数据集标签组，构建了共6个数据集。可见，PGR-MOOD在6个数据集的指标上超越了现有方法。表中加粗表示最好的结果，下划线表示第二好的结果。例如，在DrugOOD-IC50数据集的scaffold-OOD检测上，PGR-MOOD相对于现有方法，AUROC提升了7.22%，AUPR提升了3.43%，FPR95降低了9.89%。

表1 与其他方法对比

作者设计了消融实验来验证模型设计的有效性。作者在四个数据集（DrugOOD-IC50，DrugOOD-EC50，以及新数据集GOOD-HIV和GOOD-PCBA）针对骨架的OOD进行了实验，来验证去除LID，去除LOOD，和去除FGW对PGR-MOOD的影响。结果如表2所示。在四个数据集上，去除这三个部件中的任意一个，均使得AUROC和AUPR降低，FPR95升高。这表明PGR-MOOD的各个部件均增强了模型的性能。

表2 消融实验

作者还进行了案例分析。作者在多个数据集上，可视化了模型对OOD和ID样本的OOD分数分布。由于ID和OOD标签分别用1和0表示，因此OOD分数越低（越接近0），表明预测的分子越有可能是OOD。ID图和OOD图之间更显著的差距对应于更好的图OOD检测器。可见，ID和OOD被明显地分离为两个不同的分布，因此可以使用一个简单的阈值进行OOD检测并实现SOTA性能。

图2 案例分析

本文提出了分子图OOD检测模型，PGR-MOOD，通过使用扩散模型创建原型图来解决分子图OOD检测的局限性。为了克服重构测量和生成效率的挑战，PGR-MOOD使用了一种包含原型图生成器和基于FGW距离的相似函数的分子检测方法，在测试阶段，只需要测量原型图和当前输入之间的相似性，以识别具有较低值的OOD。实验结果证明了PGR-MOOD的有效性。由于药物分子数据集经常出现OOD的情况，对药物分子的OOD进行准确检测，有望在药物发现中发挥重要作用。

参考文献

Shen et al. Optimizing OOD Detection in Molecular Graphs: A Novel Approach with Diffusion Models. KDD. 2024

--------- End ---------