Learning from Historical Activations in Graph Neural Networks

TL;DR

HISTOGRAPH is a novel attention-based pooling method for GNNs that leverages intermediate layer activations to improve graph classification, especially in deep architectures.

Contribution

Introduces HISTOGRAPH, a two-stage attention mechanism that utilizes historical activations across layers for enhanced graph pooling.

Findings

Consistently improves performance over traditional pooling methods.

Provides robustness in deep GNN architectures.

Enhances node representation by modeling activation evolution.

Abstract

Graph Neural Networks (GNNs) have demonstrated remarkable success in various domains such as social networks, molecular chemistry, and more. A crucial component of GNNs is the pooling procedure, in which the node features calculated by the model are combined to form an informative final descriptor to be used for the downstream task. However, previous graph pooling schemes rely on the last GNN layer features as an input to the pooling or classifier layers, potentially under-utilizing important activations of previous layers produced during the forward pass of the model, which we regard as historical graph activations. This gap is particularly pronounced in cases where a node's representation can shift significantly over the course of many graph neural layers, and worsened by graph-specific challenges such as over-smoothing in deep architectures. To bridge this gap, we introduce HISTOGRAPH, a novel two-stage attention-based final aggregation layer that first applies a unified layer-wise attention over intermediate activations, followed by node-wise attention. By modeling the evolution of node representations across layers, our HISTOGRAPH leverages both the activation history of nodes and the graph structure to refine features used for final prediction. Empirical results on multiple graph classification benchmarks demonstrate that HISTOGRAPH offers strong performance that consistently improves traditional techniques, with particularly strong robustness in deep GNNs.