Learning Invariant Graph Representations Through Redundant Information
This work addresses the problem of out-of-distribution generalization in graph learning for researchers and practitioners, presenting an incremental improvement by applying a new information-theoretic tool to an existing bottleneck.
The paper tackles the challenge of learning invariant graph representations for out-of-distribution generalization by introducing Partial Information Decomposition to precisely target redundant information between spurious and invariant subgraphs, proposing the RIG framework that maximizes this redundancy to improve generalization under distribution shifts, with experiments on synthetic and real-world datasets demonstrating its capabilities.
Learning invariant graph representations for out-of-distribution (OOD) generalization remains challenging because the learned representations often retain spurious components. To address this challenge, this work introduces a new tool from information theory called Partial Information Decomposition (PID) that goes beyond classical information-theoretic measures. We identify limitations in existing approaches for invariant representation learning that solely rely on classical information-theoretic measures, motivating the need to precisely focus on redundant information about the target $Y$ shared between spurious subgraphs $G_s$ and invariant subgraphs $G_c$ obtained via PID. Next, we propose a new multi-level optimization framework that we call -- Redundancy-guided Invariant Graph learning (RIG) -- that maximizes redundant information while isolating spurious and causal subgraphs, enabling OOD generalization under diverse distribution shifts. Our approach relies on alternating between estimating a lower bound of redundant information (which itself requires an optimization) and maximizing it along with additional objectives. Experiments on both synthetic and real-world graph datasets demonstrate the generalization capabilities of our proposed RIG framework.