Enhancing Size Generalization in Graph Neural Networks through Disentangled Representation Learning

📄 arXiv: 2406.04601v3 📥 PDF

作者: Zheng Huang, Qihui Yang, Dawei Zhou, Yujun Yan

分类: cs.LG

发布日期: 2024-06-07 (更新: 2024-06-11)

🔗 代码/项目: GITHUB

💡 一句话要点

DISGEN：通过解耦表示学习增强图神经网络的尺寸泛化能力

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 图神经网络 尺寸泛化 解耦表示学习 图表示 互信息最小化

📋 核心要点

  1. 现有GNN在处理大于训练集的图时性能下降，主要原因是图表示中尺寸信息未被有效移除。
  2. DISGEN通过解耦图表示中的尺寸因素，采用尺寸和任务不变的增强方法，并引入解耦损失。
  3. 实验结果表明，DISGEN在真实数据集上超越了现有最佳模型，性能提升高达6%。

📝 摘要（中文）

尽管大多数图神经网络（GNN）可以在任意大小的图上运行，但其分类性能通常在大于训练期间遇到的图上会下降。现有方法未能充分解决从图表示中移除尺寸信息的问题，导致次优性能并依赖于骨干模型。为此，我们提出DISGEN，一种新颖且模型无关的框架，旨在从图表示中解耦尺寸因素。DISGEN采用尺寸和任务不变的增强方法，并引入了解耦损失，以最大限度地减少隐藏表示中的共享信息，并为该方法的有效性提供理论保证。我们的实验结果表明，DISGEN在真实世界数据集上的性能优于最先进的模型高达6%，突显了其在增强GNN尺寸泛化能力方面的有效性。我们的代码可在https://github.com/GraphmindDartmouth/DISGEN获取。

🔬 方法详解

问题定义：现有图神经网络（GNNs）虽然可以处理不同大小的图，但在测试时如果遇到比训练集更大的图，性能会显著下降。这是因为GNN学习到的图表示中包含了图的尺寸信息，导致模型对训练集中未见过的大尺寸图泛化能力不足。现有方法难以有效去除图表示中的尺寸信息，并且效果依赖于特定的骨干网络模型。

核心思路：DISGEN的核心思路是通过解耦表示学习，将图的表示分解为与任务相关但与尺寸无关的部分，以及与尺寸相关但与任务无关的部分。通过这种方式，模型可以专注于学习与图结构和节点特征相关的本质信息，而忽略图的尺寸带来的干扰。这样，即使在面对更大尺寸的图时，模型也能保持较好的泛化能力。

技术框架：DISGEN框架主要包含以下几个关键模块：1) 尺寸和任务不变的增强：对图进行增强，生成不同尺寸和任务相关的变体，用于训练解耦器。2) 解耦器：负责将图的表示分解为尺寸相关和尺寸无关两部分。3) 解耦损失：通过最小化尺寸相关和尺寸无关表示之间的互信息，鼓励模型学习到真正解耦的表示。4) 分类器：利用尺寸无关的表示进行图分类。整个框架是模型无关的，可以与各种GNN骨干网络结合使用。

关键创新：DISGEN的关键创新在于其解耦表示学习的方法，它显式地将图的尺寸信息从图表示中分离出来。与现有方法相比，DISGEN不依赖于特定的GNN结构，而是通过一种通用的解耦框架来提升尺寸泛化能力。此外，DISGEN还提供了理论保证，证明了解耦损失的有效性。

关键设计：DISGEN的关键设计包括：1) 尺寸和任务不变的增强策略：例如，随机删除节点或边，或者改变节点特征的某些属性，以生成不同尺寸和任务相关的图变体。2) 解耦损失函数：采用互信息最小化（Mutual Information Minimization）的策略，鼓励尺寸相关和尺寸无关的表示尽可能独立。3) 模型无关性：DISGEN可以与各种GNN骨干网络（如GCN、GAT等）结合使用，只需将GNN的输出作为DISGEN的输入即可。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，DISGEN在多个真实世界数据集上显著优于现有最先进的方法。例如，在某些数据集上，DISGEN的性能提升高达6%。此外，实验还验证了DISGEN与不同GNN骨干网络的兼容性，证明了其模型无关性。消融实验也表明，解耦损失和尺寸不变增强是提升性能的关键因素。

🎯 应用场景

DISGEN具有广泛的应用前景，例如在药物发现、社交网络分析、推荐系统等领域。在这些领域中，图的尺寸可能变化很大，而模型的泛化能力至关重要。通过DISGEN，可以提升GNN在这些场景下的性能，从而更好地进行药物性质预测、社交关系分析和个性化推荐。未来，该方法可以进一步扩展到其他图相关的任务中，例如图生成、图编辑等。

📄 摘要（原文）

Although most graph neural networks (GNNs) can operate on graphs of any size, their classification performance often declines on graphs larger than those encountered during training. Existing methods insufficiently address the removal of size information from graph representations, resulting in sub-optimal performance and reliance on backbone models. In response, we propose DISGEN, a novel and model-agnostic framework designed to disentangle size factors from graph representations. DISGEN employs size- and task-invariant augmentations and introduces a decoupling loss that minimizes shared information in hidden representations, with theoretical guarantees for its effectiveness. Our empirical results show that DISGEN outperforms the state-of-the-art models by up to 6% on real-world datasets, underscoring its effectiveness in enhancing the size generalizability of GNNs. Our codes are available at: https://github.com/GraphmindDartmouth/DISGEN.