Table of Contents
Fetching ...

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.

Learning from Historical Activations in Graph Neural Networks

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.
Paper Structure (21 sections, 1 theorem, 18 equations, 5 figures, 17 tables, 1 algorithm)

This paper contains 21 sections, 1 theorem, 18 equations, 5 figures, 17 tables, 1 algorithm.

Key Result

Proposition 1

Let $\mathbf{x}_v^{(l)} \in \mathbb{R}^D$ denote the embedding of node $v$ at layer $l$ of a GNN. Suppose the GNN exhibits over-smoothing, i.e., there exists some layer $L_0$ sufficiently large such that for all layers $l_1, l_2 > L_0$ and all nodes $u,v$, Let HistoGraph compute the final node embedding as where $\alpha_l$ are learned attention weights. Then, for distinct nodes $u$ and $v$, ther

Figures (5)

  • Figure 1: Overview of HistoGraph. (1) Given input node features $\mathbf{X}_0$ and adjacency $\mathbf{A}$, a backbone GNN produces historical graph activations$\mathbf{X}_{1}, .., \mathbf{X}_{L-1}$. (2) The Layer-wise attention module uses the final‐layer embedding as a query to attend over all historical states while averaging across nodes, yielding per‐node aggregated embeddings $\mathbf{H}$. (3) A Node-wise self‐attention module refines $\mathbf{H}$ by modeling interactions across nodes, producing $\mathbf{Z}$, then averaged if graph embeddings $\mathbf{G}$ is wanted.
  • Figure 2: Visualizations on the imdb-b dataset with 64-layer HistoGraph. (left) Attention patterns across layers under different training regimes. (right) Embedding evolution throughout training, measured by the normed difference between final and intermediate representations.
  • Figure 3: Graph and signal transformations: (a) input node features; (b) prediction target, the node-feature gradient; (c) GCN output trained to approximate (b) from (a); (d) HistoGraph output. The gap between GCN and HistoGraph underscores the importance of adaptive trajectory filtering. Node colors: red, blue, and green denote values $-1,0,1$.
  • Figure 4: Average training time per epoch (in log scale) for GCN backbones with 3 and 32 layers, evaluated on the molhiv and toxcast datasets. Each configuration is compared across four post-processing methods: GMTbaek2021accurate, MeanPool, HistoGraph, and HistoGraph-FT.
  • Figure 5: Barbell graph illustrating a distribution shift: a singleton node (right) is connected to a larger subgraph (left) whose size increases at test time (blue) compared to training (green). Node-wise attention helps preserve the importance of the singleton node despite the dominance of the larger subgraph.

Theorems & Definitions (1)

  • Proposition 1: Mitigating Over-smoothing with HistoGraph