[论文解读] Does Invariant Graph Learning via Environment Augmentation Learn Invariance?

Yongqiang Chen, Yatao Bian+4 more|arXiv (Cornell University)|Oct 29, 2023

Advanced Graph Neural Networks被引用 5

一句话总结

该论文提出GALA，一种新颖的不变图表示学习框架，通过使用代理环境助手模型在正确和错误预测图上进行对比学习，以识别不变子图，从而无需依赖真实环境标签。在变化充分性和一致性等最小假设下，GALA可证明地隔离出不变子图，在包括DrugOOD在内的12个数据集上实现了最先进（SOTA）的分布外（OOD）泛化性能，同时具备优异的训练效率。

ABSTRACT

Invariant graph representation learning aims to learn the invariance among data from different environments for out-of-distribution generalization on graphs. As the graph environment partitions are usually expensive to obtain, augmenting the environment information has become the de facto approach. However, the usefulness of the augmented environment information has never been verified. In this work, we find that it is fundamentally impossible to learn invariant graph representations via environment augmentation without additional assumptions. Therefore, we develop a set of minimal assumptions, including variation sufficiency and variation consistency, for feasible invariant graph learning. We then propose a new framework Graph invAriant Learning Assistant (GALA). GALA incorporates an assistant model that needs to be sensitive to graph environment changes or distribution shifts. The correctness of the proxy predictions by the assistant model hence can differentiate the variations in spurious subgraphs. We show that extracting the maximally invariant subgraph to the proxy predictions provably identifies the underlying invariant subgraph for successful OOD generalization under the established minimal assumptions. Extensive experiments on datasets including DrugOOD with various graph distribution shifts confirm the effectiveness of GALA.

研究动机与目标

• 探究在图表示学习中，环境增强是否能在无真实环境标签的情况下真正学习到不变性。
• 识别现有环境增强方法在捕捉虚假子图变化方面的根本局限性。
• 确立在无环境标签时，不变图学习可行的最小假设——变化充分性和变化一致性。
• 提出GALA框架，利用代理环境助手通过基于预测正确性的对比学习提取最大不变子图。
• 在12个具有不同分布偏移的多样化图数据集上，验证GALA的有效性和效率。

提出的方法

• GALA引入一个通过经验风险最小化（ERM）训练的环境助手模型，用于预测图标签，其预测结果作为环境变化的代理。
• 该框架根据助手模型的输出，将训练图划分为正例（正确预测）和负例（错误预测）两组。
• 应用单边对比损失，将标签相同但预测正确性不同的图拉近，从而学习到最大不变子图。
• 通过鼓励模型关注在正确和错误预测组中均保持一致的子图，确保不变性。
• 环境助手可依据数据特征，实现为ERM模型或基于聚类的可解释GNN代理预测模型。
• GALA在应用对比正则化前，先使用ERM对主GNN进行10个周期的预训练，必要时通过过采样实现类别平衡。

Figure 1: An illustration of GALA with the task of classifying graphs according to whether there exists a “House” or “Cycle” motif. Given the training data where the “House” subgraph often co-occurs with a “Grid” and the “Cycle” subgraph often co-occurs with a “Hexagon”. An ERM trained environment a

实验结果

研究问题

• RQ1 能否在无真实环境标签的情况下，通过环境增强实现不变图表示学习？
• RQ2 当前环境增强方法在图OOD泛化中存在哪些根本局限性？
• RQ3 在无环境标签时，何种最小假设下不变图学习是可行的？
• RQ4 在无环境标签的情况下，如何利用代理模型的预测来识别真正的不变子图？
• RQ5 GALA是否在OOD泛化性能上优于现有方法，同时保持有竞争力的训练效率？

主要发现

• GALA在12个多样化图数据集（包括DrugOOD）上实现了最先进性能，相较于CIGA、GREA和DisC等基线模型表现出一致的性能提升。
• 该框架保持了有竞争力的训练时间，仅比CIGA略有增加，远低于DisC和GIL等方法的开销。
• 平均而言，与GIL和DisC相比，GALA将训练时间减少了15%-25%，后者即使采用早停策略也表现出缓慢收敛。
• 在DrugOOD数据集中，基于聚类的代理预测相比ERM预测在实验中表现更优。
• 理论分析表明，在变化充分性和一致性的假设下，GALA可证明地识别出底层不变子图。
• 大量消融实验表明，环境助手预测的正确性对于隔离不变结构至关重要，验证了核心设计原则。

(a) Correlation strengths