Learning from Historical Activations in Graph Neural Networks
Yaniv Galron, Hadar Sinai, Haggai Maron, Moshe Eliasof
TL;DR
Pooling in graph neural networks often underutilizes intermediate activations, hindering multi-scale readouts. HISTOGRAPH introduces a two-stage attention framework that first attends over historical activations across layers and then across nodes to form a rich graph descriptor, with an option to train end-to-end or as a lightweight post-processing head on frozen backbones. The approach mitigates over-smoothing, supports deep GNNs, and delivers consistent improvements on TU and OGB benchmarks, often surpassing state-of-the-art pooling methods while incurring modest overhead. By explicitly leveraging the GNN computation trajectory, HISTOGRAPH provides a practical, general drop-in pooling layer that enhances graph classification, node classification, and link prediction tasks. This work demonstrates that activation history is a valuable signal for readout, enabling robust, scalable incorporation of multi-scale information in real-world graph learning settings.
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.
