CIGA 的实现:在实践中,估计两个子图的互信息通常比较困难,而监督式的对比学习 [11] 则提供了一种可行的解法:
其中对应着公式 (4) 中的正样本,而则是对应于的图表示。当时,公式 (5) 提供了对于的一种基于 von Mises-Fisher kernel density 的非参数再代入熵估计(Nonparameteric Resubstitution Entropy Estimator )[13,14]。最终 CIGA 核心部分的实现如图 5 所示,即通过在隐表示空间拉近同个类别不变子图的图表示,同时最大化不同类别不变子图的图表示,以最大化。此外,对于公式 (4) 中的另一个约束,我们则可以通过铰链损失(hinge loss)的思路进行实现,即,只优化预测时经验损失大于对应的不变子图的虚假子图。
实验与讨论
在实验中,我们使用 16 个合成或来自真实世界的数据集,对 CIGA 在不同图分布偏移下进行了充分的验证。在实验中,我们使用可解释 GNN 框架 [9] 实现了 CIGA 的原型,而实际上 CIGA 有更多实现的方式。具体的数据集以及实验细节详见文中实验部分。
合成数据集上图结构分布偏移以及混合分布偏移的表现
我们首先基于 SPMotif 数据集 [9] 构造了 SPMotif-Struc 以及 SPMotif-Mixed 数据集,其中 SPMotif-Struc 包含了特定子图与图中其他子图结构的虚假关联,以及图大小的分布偏移;而 SPMotif-Mixed 则在 SPMotif-Struc 的基础上新增了图节点属性层级的分布偏移。表中第一栏为 ERM 以及可解释 GNN 的基线,第二栏则为欧式空间最先进的分布外泛化算法。从结果中可以发现,不论是更好的 GNN 框架还是欧式空间的分布外泛化算法,都受制于图上的分布偏移,且当更多的分布偏移出现时,性能损失(更小的平均分类性能或更大的方差)将进一步增强。相对的,CIGA 则能在不同强度的分布偏移下保持良好的性能,并极大超越最好的基线表现。
真实数据集上各类图分布偏移的表现
我们接着在真实数据集和各种真实数据中存在的图分布偏移进一步测试了 CIGA 的表现,包括来自 AI 辅助制药中药物分子属性预测的 DrugOOD 中三种不同环境划分(实验环境 Assay,分子骨架 Scaffold,分子大小 Size)的三个数据集,包含了各种真实应用场景的图分布偏移;基于欧式空间中经典的图像数据集 ColoredMNIST [10] 转换得到的 CMNIST-SP,主要包含图节点属性的 PIIF 类型分布偏移;基于自然语言情感分类数据集 SST5 以及 Twitter 转化得到的 Graph-SST5 以及 Twitter [15],并且额外添加了图度数的分布偏移。此外,我们还使用了先前研究较多的 4 个分子图大小分布偏移数据集 [7],
测试结果如上表所示,可以发现,在真实数据中,由于任务难度增加,使用更好架构的 GNN 或者欧式空间的分布外泛化优化目标训练得到的模型性能甚至弱于使用 ERM 训练得到的普通 GNN 模型。这一现象也与欧式空间中更难任务下的分布外泛化实验观察得到的现象类似 [16],反应了真实数据上的分布外泛化难度以及现有方法的不足。与之相对地,CIGA 则能在所有的真实数据和图分布偏移上获得提升,甚至在某些数据集如 Twitter、PROTEINS 中达到经验最优的 Oracle 水准。在最新的图分布外泛化测试基准 GOOD 上图分类数据集的初步测试也显示了 CIGA 是目前最好且能应对各种个样的图分布偏移的图分布外泛化算法。
由于使用了可解释 GNN 作为 CIGA 的原型实现架构,我们也对模型识别得到的 DrugOOD 中的进行了可视化,发现 CIGA 确实发现了一些比较一致的分子基团用于分子属性预测。这可以为后续 AI 辅助制药提供更好的依据。
图 6. DrugOOD 中 CIGA 识别得到的部分不变子图。
总结及展望
本文通过因果推断的角度,首次将因果不变性引入至多种图分布偏移下的图分布外泛化问题中,并提出了一个全新的具有理论保证的解决框架 CIGA。大量实验也充分验证了 CIGA 优秀的分布外泛化性能。放眼未来,基于 CIGA,我们可以进一步探索更好的实现框架 [17],或为 CIGA 引入更好的具有理论保障的数据增强方法 [3,18],并在理论上建模纳入图上的协变量偏移(Covariate Shift)[19],以进一步提升 CIGA 识别不变子图的能力,促进图神经网络在 AI 辅助制药等真实应用场景的真实落地使用。
部分参考文献
[1] Zhang et al., CausalAdv: Adversarial Robustness Through the Lens of Causality, ICLR 2022.
[2] Zhu et al., Shift-Robust GNNs: Overcoming the Limitations of Localized Graph Training Data, NeurIPS 2021.
[3] Wu et al., Handling Distribution Shifts on Graphs: An Invariance Perspective, ICLR 2022.
[4] Chen et al., Bias and Debias in Recommender System: A Survey and Future Directions, TOIS 2022.
[5] Boris et al., Understanding Attention and Generalization in Graph Neural Networks, NeurIPS 2019.
[6] Gilad et al., From Local Structures to Size Generalization in Graph Neural Networks, ICML 2021.
[7] Bevilacqua et al., Size-Invariant Graph Representations for Graph Classification Extrapolations, ICML 2021.
[8] Ji et al., DrugOOD: OOD Dataset Curator and Benchmark for AI-aided Drug Discovery, arXiv 2022.
[9] Wu et al., Discovering Invariant Rationales for Graph Neural Networks, ICLR 2022.
[10] Arjovsky et al., Invariant Risk Minimization, arXiv 2020.
[11] Morris, et al., Weisfeiler and leman go neural: Higher-order graph neural networks, AAAI 2019.
[12] Khosla et al., Supervised contrastive learning, NeurIPS 2020.
[13] Wang and Isola, Understanding contrastive representation learning through alignment and uniformity on the hypersphere, ICML 2020.
[14] Ahmad and Lin, A nonparametric estimation of the entropy for absolutely continuous distributions (corresp.), IEEE Transactions on Information Theory, 22 (3):372–375, 1976.
[15] Hao et al., Explainability in Graph Neural Networks: A Taxonomic Survey, arXiv 2021.
[16] Gulrajani and Lopez-Paz, In Search of Lost Domain Generalization, ICLR 2021.
[17] Miao et al., Interpretable and Generalizable Graph Learning via Stochastic Attention Mechanism, ICML 2022.
[18] Yu et al., Finding Diverse and Predictable Subgraphs for Graph Domain Generalization, arXiv 2022.
[19] Gui et al., GOOD: A Graph Out-of-Distribution Benchmark, NeurIPS 2022 Datasets and Benchmarks.