Learning Invariant Graph Representations Through Redundant Information

TL;DR

This paper introduces a novel information-theoretic framework using Partial Information Decomposition to improve invariant graph representations, enhancing out-of-distribution generalization by isolating spurious from causal features.

Contribution

It proposes a new multi-level optimization framework called RIG that maximizes redundant information to better separate spurious and invariant graph components for OOD generalization.

Findings

RIG outperforms existing methods on synthetic datasets.

RIG demonstrates improved OOD generalization on real-world graph datasets.

The approach effectively isolates causal features from spurious correlations.

Abstract

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.