GSINA: Improving Subgraph Extraction for Graph Invariant Learning via Graph Sinkhorn Attention
Fangyu Ding, Haiyang Wang, Zhixuan Chu, Tianming Li, Zhaoping Hu, Junchi Yan
2024-02-12
Abstract:Graph invariant learning (GIL) has been an effective approach to discovering
the invariant relationships between graph data and its labels for different
graph learning tasks under various distribution shifts. Many recent endeavors
of GIL focus on extracting the invariant subgraph from the input graph for
prediction as a regularization strategy to improve the generalization
performance of graph learning. Despite their success, such methods also have
various limitations in obtaining their invariant subgraphs. In this paper, we
provide in-depth analyses of the drawbacks of existing works and propose
corresponding principles of our invariant subgraph extraction: 1) the sparsity,
to filter out the variant features, 2) the softness, for a broader solution
space, and 3) the differentiability, for a soundly end-to-end optimization. To
meet these principles in one shot, we leverage the Optimal Transport (OT)
theory and propose a novel graph attention mechanism called Graph Sinkhorn
Attention (GSINA). This novel approach serves as a powerful regularization
method for GIL tasks. By GSINA, we are able to obtain meaningful,
differentiable invariant subgraphs with controllable sparsity and softness.
Moreover, GSINA is a general graph learning framework that could handle GIL
tasks of multiple data grain levels. Extensive experiments on both synthetic
and real-world datasets validate the superiority of our GSINA, which
outperforms the state-of-the-art GIL methods by large margins on both
graph-level tasks and node-level tasks. Our code is publicly available at
\url{https://github.com/dingfangyu/GSINA}.
Machine Learning,Artificial Intelligence