KDD 2023 | 以灵活应万变，一种图分布外泛化框架

《易经》中的“不易之道”强调了“以不变应万变”的哲学思想，提出了“道”是宇宙的根本，而“道”是永恒不变的。因此，人们应该通过学习和理解“道”的不变性，来应对宇宙中的无尽变化。在机器学习模型训练中采用不变性原则，引导模型学习数据与标签之间的不变因果关系，使得模型在不同环境或数据分布下保持一致的性能，这是不变学习“以不变应万变”的工作机理。

然而，世界瞬息万变，唯一不变的是变化本身，宇宙的根本在当今世界的变化中可能很难找到。“以灵活应万变”是一种面对变化时更加可行的思路，能够及时调整和适应变化的策略或模型。遵循“以灵活应万变”的设计原则，本文提出了一种灵活的不变图表示学习框架FLOOD，可以根据目标分布灵活地更新图编码器参数以使模型获得最佳的泛化效果。

论文题目：

FLOOD: A Flexible Invariant Learning Framework for Out-of-Distribution Generalization on Graphs

论文地址：

https://dl.acm.org/doi/10.1145/3580305.3599355

作者主页：

http://yliu.site/

一、研究动机

现有的图数据分布外泛化解决方案基于不变学习（Invariant Learning）[2]设计，其核心目标是在多个环境中学习不变量。不变学习在面临分布偏移时具有显著的优势，能够保持稳定的泛化性能。在不变学习中，“环境”通常指的是具有特定特征分布的样本集合。例如，在图数据中可以将不同度分布、不同连边密度的节点集合视为不同的“环境”。每个环境都有其独特的数据分布，但所有环境共享相同的因果关系或生成机制。现有的基于不变学习的图分布外泛化解决方案存在如下挑战：

环境构造难：GIL[3]通过生成一个软掩码矩阵来识别不变子图和环境子图，应用于图分类任务难以适用于节点分类；EERM[4]用生成器生成多个训练环境，计算开销大且生成器输出难控制。
模型灵活度低：不变学习后模型参数是固定不变的，在测试分布发生偏移时，难以根据目标分布变化灵活地调整模型实现泛化。

二、研究方法

实现分布外泛化的图表示学习，要求模型在多种训练环境中，学习得到节点的不变表示。通过风险外推准则，在降低每个环境训练风险的同时，要求训练风险之间的相似性增大，当测试环境分布发生偏移时，测试风险的变化减小，模型泛化性增强。

为了捕捉目标分布的变化，需要设计一种不依赖于真实标签的自监督学习模块感知分布变化，因为分布偏移后的样本标签往往是不可获得的。图对比学习是常用的图自监督学习方法，但其中负样本采样往往需要大量计算资源，不利于在发生概念漂移时迅速调整模型，因此拟采用不同训练环境间的自助表示学习，避免负采样带来的计算开销。

三、实验分析

在图分布外泛化基准数据集GOOD[7]上进行实验，四个数据集在不同偏移域上的节点分类的分布内和分布外准确率如下表所示。

可以得出以下结论：

测试时训练提升模型泛化能力。由于自助表示学习模块的灵活性，模型可以通过共享编码器根据目标分布变化自适应调整参数，使得FLOOD取得最佳的分布外泛化效果。
图数据增强提升模型泛化能力。Mixup在多个数据集上取得次优表现，说明图数据增强有利于增加训练数据分布的多样性。得益于此，FLOOD通过图数据增强的方式获得了多样性较高的训练环境。

测试时训练前后的节点表示可视化如下图所示，不同颜色代表不同的节点类别，可以看出，根据目标分布微调后的表示具有更加明显的边界，更容易被分类器区分开从而取得更好的泛化效果。

四、总结与展望

本文提出一种灵活的不变图表示学习框架FLOOD，包括环境无关的不变表示学习、分布感知的自助表示学习和分布偏移自适应的测试时训练三个模块。环境无关的不变表示学习通过图数据增强构建多个训练环境，并在风险外推的准则下学习节点的不变表示；分布感知的自助表示学习以一种自监督的方式与不变表示学习共享图编码器进行训练；在测试阶段，共享编码器可以根据目标分布灵活地更新参数以使模型获得最佳的泛化效果。

参考文献

[1] Liu, Yang, Xiang Ao, Fuli Feng, Yunshan Ma, Kuan Li, Tat-Seng Chua, and Qing He. "FLOOD: A flexible invariant learning framework for out-of-distribution generalization on graphs." In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 1548-1558. 2023.

[2] Arjovsky, Martin, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. "Invariant risk minimization." arXiv preprint arXiv:1907.02893 (2019).

[3] Li, Haoyang, Ziwei Zhang, Xin Wang, and Wenwu Zhu. "Learning invariant graph representations for out-of-distribution generalization." Advances in Neural Information Processing Systems 35 (2022): 11828-11841.

[4] Wu, Qitian, Hengrui Zhang, Junchi Yan, and David Wipf. "Handling distribution shifts on graphs: An invariance perspective." In International Conference on Learning Representations, 2022.

[5] Krueger, David, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. "Out-of-distribution generalization via risk extrapolation (rex)." In International Conference on Machine Learning, pp. 5815-5826. PMLR, 2021.

[6] Wu, Ying-Xin, Xiang Wang, An Zhang, Xiangnan He, and Tat-Seng Chua. "Discovering invariant rationales for graph neural networks." In International Conference on Learning Representations, 2022.

[7] Gui, Shurui, Xiner Li, Limei Wang, and Shuiwang Ji. "Good: A graph out-of-distribution benchmark." Advances in Neural Information Processing Systems 35 (2022): 2059-2073.

Illustration From IconScout By Delesign Graphics