NeurIPS 2024 | 用于药物相互作用预测的定制化子图选择与编码

今天介绍的是发表在 NeurIPS 2024 的论文Customized Subgraph Selection and Encoding for Drug-drug Interaction Prediction，该论文提出了一种结合子图选择空间松弛机制和单路径NAS优化的定制化子图选择与编码方法，用于药物相互作用预测。

一、摘要

基于子图的方法已被证明在预测药物-药物相互作用（DDI）中有效且具有解释性，这对于医学实践和药物开发至关重要。然而，在这些方法中定制子图选择和编码的组件仍未得到充分开发，主要由于人工调整的高成本。在本研究中，受神经架构搜索（NAS）成功的启发，作者提出了一种方法，用于在基于子图的框架内搜索特定数据的组件。具体而言，作者引入了广泛的子图选择和编码空间，以考虑DDI预测中不同上下文的药物相互作用。为了解决大搜索空间和高采样成本的挑战，作者设计了一种使用近似策略的松弛机制，能够高效地探索所有可能的子图配置。该方法允许对搜索空间进行稳健的探索。大量实验表明，该方法在发现子图和编码函数方面的有效性和优越性，展示了模型的适应性。

二、子图方法与NAS在药物相互作用预测中的结合探索

2.1 子图方法在图神经网络中的应用及其局限性

近年来，图神经网络（Graph Neural Network, GNN）因其在复杂关系建模方面的出色表现，成为DDI预测的重要工具。传统的GNN方法通过对整个药物交互网络进行信息传播来学习全局特征。然而，这种方法往往难以捕捉特定药物对之间的局部交互信息，导致预测准确性和解释性受限。

子图方法（Subgraph-based Methods）通过提取目标药物对周围的局部子图，将复杂的多关系预测问题转化为局部子图分类问题，为提高预测精度提供了新思路。然而，现有的子图方法通常采用固定范围的子图提取策略，无法针对不同药物对的特定交互语义动态调整，容易引入无关信息或遗漏关键特征。此外，这些方法的子图编码函数大多为人工设计，难以适应数据的多样性与复杂性。

2.2 NAS技术在自动化模型设计中的成功应用启发

神经网络架构搜索（Neural Architecture Search, NAS）技术近年来在深度学习领域取得了显著进展，尤其是在图神经网络和知识图谱学习中展现出强大的潜力。通过自动化搜索最优模型架构，NAS技术能够在数据驱动的基础上设计出性能优越且高效的深度学习模型。NAS的成功应用启发作者将这一技术引入到DDI预测中，通过自动化搜索子图选择和编码组件，克服人工设计方法的局限性，从而实现更加精准和高效的药物相互作用预测。

三、方法介绍

3.1 定制化子图选择空间

在药物相互作用预测中，不同的药物对可能具有不同的语义上下文，因此需要针对性地选择子图进行特征学习。为了实现这一目标，本文提出了一个灵活的子图选择空间。具体地，定义每对药物的子图选择空间为：

其中，表示药物和的跳和跳邻域的联合子图。公式可以表示为：

其中和分别表示节点和的跳和跳邻居。

上述定义允许模型在多种子图范围内灵活选择适合特定药物对的子图，有效避免了固定范围选择可能带来的噪声或信息遗漏问题。同时，为了控制子图选择空间的复杂度，设置了最大跳数阈值𝜂，从而限制搜索空间的规模。

3.2 子图编码空间

子图的语义编码需要能够适应药物交互的多样性，例如代谢层面的交互通常是非对称的，而表型层面的交互则是对称的。为此，本文基于以下通用的消息传递框架设计了灵活的子图编码空间：其中：

和分别表示节点和交互关系的嵌入；
是消息计算函数（如乘法、相关性等）；
是聚合函数（如求和、取平均）；
是组合函数（如拼接、全连接网络）；
是激活函数（如 ReLU、Tanh）。

以 DrugBank 数据集为例，其代谢层面交互具有非对称特性。通过上述设计，模型能够自适应选择最佳的编码函数，精准建模多样的药物交互语义。

3.3 高效搜索策略

为了有效探索上述子图选择和编码空间，本文设计了连续化搜索机制和子图表示近似方法，显著降低了计算成本。首先，将离散的子图选择空间转化为连续形式。采用 Gumbel-Softmax 技术，计算子图选择的概率：其中：

是子图编码表示；
是子图得分函数；
是 Gumbel 随机变量；
是控制分布锐度的温度参数。

此外，通过子图表示近似方法，将全图编码结果用于近似子图表示，避免了显式采样大规模子图的高昂开销。公式如下：其中表示在全图上由节点生成的跳隐藏表示。

通过以上设计，模型在性能和效率之间取得了良好平衡，为药物相互作用预测提供了一个高效且可解释的解决方案。

3.4 鲁棒的搜索算法

论文通过提出的子图选择空间松弛机制，将离散的高维搜索空间转化为连续形式，使得单路径NAS（One-shot NAS）优化成为可能，并利用子图表示近似策略有效降低了搜索成本。此外，为了减少超级网络训练的计算开销，本文采用了单路径一阶训练策略（SPOS），同时设计了一种消息感知的分区训练方法，将超级网络分为子超级网络，以提高网络一致性和鲁棒性，从而使搜索算法更加高效稳定。算法的具体流程如下：

Algorithm 1: CSSE-DDI的搜索算法

输入: 超级网络，基于消息计算函数类别的分区数（例如），子超级网络。

超级网络训练

持续对超级网络 ( S ) 进行单路径采样并训练直至收敛。

将分割为个子超级网络。

阶子超级网络训练

对于每个，执行以下步骤：

从中传递权重初始化。
持续对进行单路径采样并训练直至收敛。

搜索阶段

a. 通过自然梯度下降法，从子超级网络中搜索最优的编码函数。

b. 根据验证数据中具有最高概率的子图保留操作，选择中的最优子图。

3.5 方法优势

高效性：单路径优化与连续化搜索结合，显著降低搜索成本。
鲁棒性：分区训练提高了模型对不同消息算子的适应能力。
准确性：优化子图选择和编码过程，实现更高的预测性能。

四、实验与结果

4.1 数据集与任务设置

实验采用 DrugBank 和 TWOSIDES 两个基准药物相互作用数据集进行评估。针对不同的预测需求，实验设置了两种场景：

S0 设置：药物对存在于已知的 DDI 图中，适合验证模型对已知数据的预测能力。
S1 设置：包含新药物，适合评估模型在未知场景中的泛化能力。

评价指标包括 F1 分数、准确率Accuracy、Cohen’s （DrugBank 数据集），以及 ROC-AUC、PR-AUC 和 AP@50（TWOSIDES 数据集）。

4.2 S0 设置下的性能比较

实验结果表明，CSSE-DDI 在 S0 设置下表现优异，显著优于基线方法。具体数据如下：

DrugBank 数据集：
CSSE-DDI 的 F1 分数为 92.08，显著高于基于 GNN 方法的 TransFOL（89.97）。
Cohen’s 为 95.56，展示了对多分类任务的强大适应性。
TWOSIDES 数据集：
ROC-AUC 和 PR-AUC 分别达到 95.47 和 94.21，同样领先其他方法。

此外，实验分析了不同基线的子图选择策略和编码功能对性能的影响，CSSE-DDI 在设计灵活性和子图语义捕获上均表现出优势。

4.3 S1 设置下的性能比较

S1 设置中包含了新药物，使预测任务更加具有挑战性。在这种场景下，CSSE-DDI 仍然保持了较强的鲁棒性：

DrugBank 数据集：

F1 分数达到 37.24，高于 SumGNN（26.57）。

Cohen’s 为 49.97，展示了对未知药物对预测的适应能力。

TWOSIDES 数据集：

PR-AUC 为 86.47，同样领先其他基线。

这表明，CSSE-DDI 能够在新药物预测任务中有效学习相关特征，并保持较高的预测性能。

4.4 案例研究

4.4.1 细粒度子图选择

实验可视化了 DrugBank 数据集中某些药物对的子图范围，例如：

阿司匹林（DB00945）与华法林（DB00682）：子图范围为。该子图揭示了药物间的药代动力学和代谢关系，为药物相互作用提供了生物学解释。
DB09580 与 DB00652：子图范围为，同样展示了语义相关的药物交互特征。

4.4.2 数据特定的编码函数

通过 DrugBank 和 TWOSIDES 数据集上的编码函数可视化，可以看出：

DrugBank 数据集中非对称的 CORR 操作更常见，适合建模代谢相关的非对称关系。
TWOSIDES 数据集中以对称的 MULT 操作为主，更适合建模表型交互。

CSSE-DDI 能够动态适配不同数据集的编码需求，从而提升预测性能。

4.5 搜索策略影响

为验证搜索策略的有效性，实验分别对比了以下两种变体：

CSSE-DDI w/o MAP：移除基于消息感知的分区训练策略。
CSSE-DDI w/o SPOS：移除单路径优化策略。

结果表明，完整的 CSSE-DDI 相比变体在所有评价指标上均有显著提升，例如：

在 DrugBank 数据集上，完整模型的 F1 分数为 92.08，去除 MAP 或 SPOS 后分别下降至 90.17 和 90.97。

4.6 超参数分析

对跳数阈值进行敏感性分析显示：

当时，模型性能接近最佳。
进一步增加并未带来显著提升，反而可能引入噪声并增加计算成本。

实验结果表明，CSSE-DDI 结合了灵活的子图选择和数据特定的编码功能，不仅在多分类任务中表现出色，还能适应更复杂的多标签预测任务。通过对药物相互作用的细粒度建模，该方法为真实世界的药物研发提供了强大的支持。

五、总结

本文提出了一种用于药物相互作用预测的定制化子图选择与编码方法，通过灵活设计子图选择空间和编码函数，提高了模型对药物交互语义的捕捉能力。通过引入子图选择空间的松弛机制和搜索空间连续化，支持单路径神经架构搜索（NAS），显著优化了搜索效率和预测性能。此外，设计的消息感知分区训练策略有效缓解了消息计算算子之间的耦合效应，提升了模型的稳定性和鲁棒性。实验结果表明，本文方法在多个基准数据集上均超越了现有的GNN、子图方法和NAS方法，展现了其在多分类和多标签任务中的优越性。通过案例分析，进一步验证了模型在细粒度子图选择和数据特定编码上的适应性，为药物研发和精准医疗提供了强有力的工具支持。