[论文解读] Unleashing the Power of Graph Data Augmentation on Covariate Distribution Shift

Yongduo Sui, Wu, Qitian+6 more|arXiv (Cornell University)|Nov 5, 2022

Advanced Graph Neural Networks被引用 7

一句话总结

本文提出对抗性不变增强（AIA），一种新颖的图数据增强方法，通过显式建模稳定特征与环境特征，以应对图表示学习中的协变量偏移问题。通过对抗性地生成多样化环境的同时保持稳定特征，AIA 提升了分布外泛化能力，在多个基准测试中取得最先进性能，并具备坚实的分布鲁棒性理论基础。

ABSTRACT

The issue of distribution shifts is emerging as a critical concern in graph representation learning. From the perspective of invariant learning and stable learning, a recently well-established paradigm for out-of-distribution generalization, stable features of the graph are assumed to causally determine labels, while environmental features tend to be unstable and can lead to the two primary types of distribution shifts. The correlation shift is often caused by the spurious correlation between environmental features and labels that differs between the training and test data; the covariate shift often stems from the presence of new environmental features in test data. However, most strategies, such as invariant learning or graph augmentation, typically struggle with limited training environments or perturbed stable features, thus exposing limitations in handling the problem of covariate shift. To address this challenge, we propose a simple-yet-effective data augmentation strategy, Adversarial Invariant Augmentation (AIA), to handle the covariate shift on graphs. Specifically, given the training data, AIA aims to extrapolate and generate new environments, while concurrently preserving the original stable features during the augmentation process. Such a design equips the graph classification model with an enhanced capability to identify stable features in new environments, thereby effectively tackling the covariate shift in data. Extensive experiments with in-depth empirical analysis demonstrate the superiority of our approach. The implementation codes are publicly available at https://github.com/yongduosui/AIA.

研究动机与目标

• 为解决图表示学习中协变量偏移这一未被充分探索的问题，即测试数据包含训练数据中未出现过的环境特征。
• 开发一种数据增强策略，显式分离并建模图数据中的稳定特征与环境特征。
• 通过生成多样化、受控且环境上不同的增强环境，提升分布外泛化能力。
• 提供一个可度量的环境特征差异指标，并利用该指标指导增强过程。
• 在理论鲁棒半径内实现稳健泛化，确保在分布偏移下性能的稳定性。

提出的方法

• AIA 采用双掩码架构：一个掩码用于识别稳定特征，另一个用于隔离环境特征。
• 通过对抗性增强器扰动环境特征以生成新环境，同时利用一致性损失保持稳定特征不变。
• 通过最大化一个新提出的度量 GCS(𝑃̃,𝑃)，显式测量并鼓励增强分布与原始分布之间的环境特征差异。
• 通过充分性与独立性损失强制保持稳定特征的一致性，确保模型依赖稳定特征进行预测。
• 训练框架在对抗性增强器与图神经网络主干之间交替优化，受分布鲁棒优化（DRO）原则指导。
• AIA 在理论基础上建立于 DRO，保证在鲁棒半径 ρ 内的泛化能力，而基于 IRM 的方法则需已知测试环境。

实验结果

研究问题

• RQ1 显式分离稳定特征与环境特征是否能提升图数据在协变量偏移下的泛化能力？
• RQ2 如何设计数据增强方法，以在保留稳定特征的同时生成多样化且未见过的环境特征？
• RQ3 可度量的环境特征差异指标是否能提升图增强的有效性？
• RQ4 通过保留稳定特征的对抗性增强是否能实现优于现有方法的分布外泛化性能？
• RQ5 基于 DRO 的训练目标是否能在协变量偏移下确保在定义分布半径内的鲁棒性？

主要发现

• AIA 在多种图分类基准测试中，在协变量偏移设置下取得最先进性能，优于现有的不变性学习与数据增强方法。
• 所提出的 GCS(𝑃̃,𝑃) 度量能有效衡量环境特征差异，并与泛化性能的提升呈正相关。
• 消融实验表明，保留稳定特征并显式建模环境变化是实现鲁棒性的关键因素。
• AIA 在测试数据包含训练期间未见过的全新环境特征时，仍表现出优越的泛化能力。
• 该方法在多种数据集（包括分子图与推荐系统）中，在受控协变量偏移设置下均实现一致的性能提升。
• 尽管训练过程更复杂，但性能稳定且方差低，表明双目标优化被有效实现。