Early-Exit Graph Neural Networks

TL;DR

This paper introduces Early-Exit GNNs with symmetry-based inductive biases that enable adaptive inference, reducing computation on simple graphs while maintaining accuracy on complex ones, thus improving efficiency in graph learning tasks.

Contribution

We propose Symmetric-Anti-Symmetric GNNs and attach confidence-aware exit heads to create adaptive Early-Exit GNNs that optimize inference efficiency and accuracy.

Findings

EEGNNs learn task-driven exit strategies.

Achieve favorable accuracy-efficiency trade-offs.

Perform well on heterophilic and long-range graph tasks.

Abstract

Early-exit mechanisms allow deep neural networks to stop inference once prediction confidence is high, reducing latency and energy on easy inputs while retaining full-depth accuracy on harder ones. Similarly, adding early exit mechanisms to Graph Neural Networks (GNNs), the go-to models for graph-structured data, allows for dynamic trading depth for confidence on simple graphs while maintaining full-depth accuracy on harder ones to capture intricate relationships. Yet, their potential in deep GNNs, where over-smoothing, over-squashing or more generally vanishing gradients prevent these model to properly learn, remains largely unexplored. To address this, we introduce Symmetric-Anti-Symmetric GNNs (SAS-GNN), whose symmetry-based inductive biases yield stable intermediate representations that support safe early exits. Building on this backbone, we propose Early-Exit GNNs (EEGNNs), which attach confidence-aware exit neural heads which are trainable end-to-end based on the task objective, enabling on-the-fly termination at node or graph level. Experiments show that EEGNNs learn task-driven exit strategies, while achieving competitive results on heterophilic graphs and long-range tasks. Even when not outperforming the strongest baselines, EEGNNs consistently deliver favorable accuracy-efficiency trade-offs thanks to their adaptive and parameter-efficient design. We plan to release the code to reproduce our experiments.