PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning from Historical Activations in Graph Neural Networks

Yaniv Galron, Hadar Sinai, Haggai Maron, Moshe Eliasof

arXiv:2601.01123v2cs.LGcs.AI

TL;DR

Graph pooling commonly uses final-layer features, leaving intermediate activations underused even though depth can overwrite early representations and cause over-smoothing. HISTOGRAPH pools layer histories with layer-wise and node-wise attention; the paper reports consistent improvements across graph-level, node-classification, and link-prediction settings. It can also be used as a post-processing tool with models trained using standard pooling.

  • Problem

    Existing pooling layers underuse intermediate activations, which carry multi-scale information that may be lost as depth overwrites early features or causes over-smoothing.

    Source: §1 · pp. 1–2

  • Method

    HISTOGRAPH applies layer-wise attention over stacked node embeddings from GNN layers, followed by node-wise self-attention at readout.

    Source: §1 · p. 2, §3 · p. 4

  • Results

    Across graph-level classification, node classification, and link prediction, HISTOGRAPH consistently improves over strong pooling baselines and matches or surpasses leading methods.

    Source: §1 · p. 3, §6 · p. 10

  • Takeaways & Limitations

    HISTOGRAPH can serve as a post-processing tool to enhance models trained with standard graph pooling layers.

    Source: §1 · p. 3

    The current design is best suited to small-to-medium-sized graphs.

    Source: §3 · p. 5

Abstract

from arXiv · show

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.

1 INTRODUCTION

Existing pooling methods underuse intermediate GNN activations, despite the multi-scale information layers encode and the risk that depth overwrites early features. HISTOGRAPH attends over layer histories and node interactions, and the paper reports consistent gains on TU and OGB benchmarks.

1 INTRODUCTION

  • Pooling layers underuse intermediate message-passing activations, limiting their ability to capture long-range dependencies and hierarchical patterns.The paper frames this as a gap in existing graph pooling schemes.Source: §1 · p. 1
  • Early GNN layers capture local motifs while deeper layers encode global patterns, but depth can overwrite early information and cause over-smoothing.Over-smoothing can make node representations indistinguishable.Source: §1 · p. 1
  • HISTOGRAPH uses layer-wise attention over each node’s historical embeddings followed by node-wise attention across the graph.The layer-wise module selects informative representations; the node-wise module aggregates global context without changing the underlying GNN architecture.Source: §1 · p. 2
  • HISTOGRAPH consistently outperforms strong GNN and pooling baselines on TU and OGB benchmarks, supporting computational history as a general inductive bias.The paper also presents HISTOGRAPH as a post-processing head for frozen pretrained GNNs.Source: §1 · pp. 2–3

2 RELATED WORKS

Prior graph pooling methods range from simple invariant readouts to learned selection, clustering, and attention, while residual approaches preserve information through node updates or hierarchical coarsening. HISTOGRAPH instead pools layer histories at readout and combines them with spatial information.

Pooling in Graph Learning

  • Table 1 compares pooling methods by intermediate-representation use, structural considerations, and layer-node modeling; the text says HISTOGRAPH alone combines intermediate representations with structural information.Its readout fuses layer-wise features with node-wise spatial dependencies without clustering or dropping nodes.Source: Table 1 · p. 3, §2 · p. 3
  • Hierarchical pooling captures multi-scale structure but can add complexity and risk information loss; HISTOGRAPH instead directly pools historical activations.Earlier pooling approaches include learned node selection and soft clustering.Source: §2 · p. 3

Residual Connections

  • Unlike residual connections within node updates or hierarchical coarsening, HISTOGRAPH creates a readout-level shortcut that revisits and integrates node histories across layers.The related-work discussion contrasts this with residual and skip-connection approaches.Source: §2 · pp. 3–4

3 LEARNING FROM HISTORICAL GRAPH ACTIVATIONS

HISTOGRAPH stacks node embeddings from GNN layers as historical activations, then applies layer-wise attention followed by node-wise self-attention at readout. Its design also supports frozen-backbone training, while the authors identify graph size as a scope constraint.

3 LEARNING FROM HISTORICAL GRAPH ACTIVATIONS

  • HISTOGRAPH treats hidden representations across GNN layers as historical activations and pools them to produce graph representations.When layer outputs have different dimensions, per-layer projections map them to a common hidden dimension before stacking.Source: §3 · p. 4
  • Fixed sinusoidal positional encodings are added to projected activations to represent layer order.The positional encodings are broadcast across nodes.Source: §3 · p. 4
  • Layer-wise attention uses each node’s final-layer embedding as a query over its historical states to form an aggregated embedding.The method applies scaled dot-product attention and averages across nodes to obtain layer weights.Source: §3 · p. 4
  • Signed normalization permits negative layer weights, so the aggregation can express additive or subtractive relationships between layers.The paper contrasts this with softmax, which enforces non-negative weights.Source: §3 · pp. 4–5
  • Node-wise multi-head self-attention refines the aggregated node embeddings once at readout, and averaging its outputs yields a graph representation.Applying spatial attention only at readout adds a single O(N^2D) cost rather than repeating it at every GNN depth.Source: §3 · p. 5
  • With a frozen pretrained backbone, HISTOGRAPH caches layer activations and trains only its head, avoiding backpropagation through the GNN layers.The authors describe this mode as useful in low-resource or few-shot regimes.Source: §3 · p. 5
  • The current design is best suited to small-to-medium-sized graphs; scaling strategies for larger graphs are discussed in Appendix G.Source: §3 · p. 5

4 PROPERTIES OF HISTOGRAPH

HISTOGRAPH’s layer-wise attention acts as an adaptive filter over node representations across layers, while node-wise attention aggregates spatial context. The section describes how these choices support signal filtering and mitigate over-smoothing.

Properties of HISTOGRAPH

  • The section identifies node-wise attention at readout as another beneficial property of HISTOGRAPH.The stated properties also include layer-wise mitigation of over-smoothing and signed-normalization filtering.Source: §4 · p. 5

Mitigating Over-smoothing

  • Nonzero attention to early layers lets final embeddings retain discriminative representations and remain distinct despite over-smoothing in deeper layers.The passage links this mechanism to robustness in deep GNNs and reports corroborating depth and feature-distance diagnostics.Source: §4 · p. 6

Interpretability of Learned Attention Weights

  • Learned attention places substantial weight on early pre-over-smoothing layers and the final layer, balancing local and global information.Figure 2 visualizes attention patterns across training regimes and embedding evolution during training.Source: §4 · p. 6, Figure 2 · p. 6

Adaptive Trajectory Filter

  • Signed normalization permits negative layer weights, enabling low-pass averaging, high-pass first differences, and general FIR filtering unavailable to softmax attention.The filter balances shallow-layer local information against deep-layer global information at readout.Source: §4 · p. 7
  • On a barbell graph, HISTOGRAPH preserves a sharp gradient discontinuity through high-pass filtering, whereas a standard GCN fails to do so.Figure 3 contrasts the input, gradient target, GCN output, and HISTOGRAPH output.Source: §4 · p. 7, Figure 3 · p. 7

5 EXPERIMENTS

The experiments evaluate HISTOGRAPH on graph classification and node classification, both end-to-end and as a head on pretrained GNNs, and assess the contributions of its components.

5 EXPERIMENTS

  • The evaluation spans TU and OGB benchmarks, node-classification datasets, and comparisons with diverse pooling and graph-representation baselines.The experiments include 5-layer GINs on TU datasets, 3-layer GCNs on OGB, and depth-varying node-classification tests.Source: §5 · pp. 7–8, §5.1 · p. 8

5.1 END-TO-END ACTIVATION AGGREGATION WITH HISTOGRAPH

  • HISTOGRAPH achieves state-of-the-art performance on 5 of 7 TU datasets and top ROC-AUC on 3 of 4 OGB datasets, while remaining competitive through 64 GNN layers on node classification.The reported node-classification pattern contrasts with standard GCN accuracy degrading as depth increases.Source: §5.1 · pp. 8–9, Table 2 · p. 8, Table 3 · p. 8

5.2 POST-PROCESSING OF TRAINED GNNS WITH HISTOGRAPH

  • Frozen-head HISTOGRAPH improves over MeanPool, raising IMDB-M accuracy from 54.7% to 67.3%; on IMDB-B, FT and Full FT both reach 94.0%.Full FT attains the best PTC score at 97.1%.Source: §5.2 · p. 9, Table 5 · p. 9
  • For 32-layer GCNs, frozen-head training is faster than MeanPool on both evaluated datasets, and HISTOGRAPH is faster than GMT in almost all cases.Figure 4 compares average training time per epoch for 3- and 32-layer GCNs on MOLHIV and TOXCAST.Source: §5.2 · p. 9, Figure 4 · p. 10

5.3 ABLATION STUDY

  • On PROTEINS, HISTOGRAPH reaches 97.80% ± 0.40 accuracy, while removing division-by-sum normalization yields 74.45% ± 6.28; removing either attention component also reduces accuracy.The authors attribute instability and overshadowing of early discriminative layers to unbounded attention weights when layer-wise attention is disabled.Source: §5.3 · p. 10, Table 6 · p. 10

6 CONCLUSION

HISTOGRAPH uses historical activations in a two-stage attention pooling layer, and the conclusion reports consistent improvements across several graph-learning settings. It also identifies very large graphs as a boundary for the current design.

6 CONCLUSION

  • Across TU and OGB graph-level benchmarks, node classification, and OGBL-COLLAB link prediction, HISTOGRAPH consistently improves over strong pooling baselines and matches or surpasses leading methods on multiple datasets.The conclusion reports these cross-task results without specifying individual benchmark values here.Source: §6 · p. 10
  • As a post-processing head on frozen backbones, HISTOGRAPH delivers additional gains without retraining the encoder.Source: §6 · p. 10
  • The current design is best suited to small-to-medium graphs; sparse attention or hierarchical coarsening are proposed directions for extending it to very large graphs.Source: §6 · p. 10

A DATASET STATISTICS

The appendix reports dataset statistics for molecular-property and graph-classification benchmarks. Together, the datasets vary in graph size and label space.

A DATASET STATISTICS

  • Tables 7 and 8 report dataset statistics for OGB molecular-property datasets and TU graph-classification datasets, respectively.The reported statistics include number of graphs, number of classes, and average nodes per graph.Source: A DATASET STATISTICS · p. 16, Table 7 · p. 16, Table 8 · p. 16
  • The OGB statistics cover MOLHIV, MOLBBBP, MOLTOX21, and TOXCAST; the TU statistics cover IMDB-B, IMDB-M, MUTAG, PTC, PROTEINS, RDT-B, and NCI1.Source: A DATASET STATISTICS · p. 16
  • The datasets span small, medium, and large graphs and diverse class distributions, providing an evaluation setting across varied graph sizes and label spaces.Source: A DATASET STATISTICS · p. 16

B IMPLEMENTATION DETAILS OF HISTOGRAPH

HISTOGRAPH aggregates node representations across GNN layers and then models spatial dependencies among nodes to produce a graph-level representation. A mean-pooling signal can be blended into the output to guide early optimization.

B IMPLEMENTATION DETAILS OF HISTOGRAPH

  • The forward pass projects historical node embeddings to a shared dimension, adds layer-position encodings, attends over layers, applies node-wise self-attention, and averages nodes into a graph representation.The final-layer embeddings query the layer-wise attention; all layer embeddings serve as keys and values.Source: B IMPLEMENTATION DETAILS OF HISTOGRAPH · p. 16
  • A learnable convex weight combines HISTOGRAPH’s output with mean pooling of final-layer node embeddings; the mean signal helps guide optimization early in training.The authors experimented with fixed and learnable values of the weight.Source: B IMPLEMENTATION DETAILS OF HISTOGRAPH · p. 16

C EXPERIMENTAL DETAILS

The supplied material describes HISTOGRAPH’s forward pass, hyperparameter selection, and two properties: node-wise attention can preserve a singleton node’s importance, while retaining early-layer representations mitigates over-smoothing.

C EXPERIMENTAL DETAILS

  • Experiments used NVIDIA L40, A100, and GeForce RTX 4090 GPUs, with W&B for logging, hyperparameter tuning, and model selection.Source: C EXPERIMENTAL DETAILS · p. 16
  • The forward-pass algorithm uses last-layer embeddings as the query and all layers as keys and values, normalizes attention over layers, then applies node-wise attention and averages nodes.Source: C EXPERIMENTAL DETAILS · p. 17

C.1 HYPERPARAMETERS

  • Hyperparameters were selected by grid search and Bayesian optimization using validation performance, across optimization, architecture, attention, and learning-rate-schedule choices.The tuned parameters include batch size, hidden dimension, learning rate, weight decay, attention heads, dropout, and attention-coefficient parameterization.Source: §C.1 · p. 17

D ADDITIONAL PROPERTIES OF HISTOGRAPH

  • For permutation-invariant tasks that depend on non-uniform node importance or inter-node interactions, mean pooling alone cannot approximate the target without additional inter-node operations.In the illustrated distribution shift, node-wise attention can selectively attend to a singleton node despite growth of the larger connected subgraph.Source: D ADDITIONAL PROPERTIES OF HISTOGRAPH · p. 17, Figure 5 · p. 18

E.1 HISTOGRAPH MITIGATES OVERSMOOTHING

  • Retaining a discriminative early-layer representation in the final embedding mitigates over-smoothing, which would otherwise make node embeddings indistinguishable.The argument relies on at least one early layer receiving nonzero attention mass.Source: §E.1 · p. 18

F ADDITIONAL EXPERIMENTS

HISTOGRAPH improves results across node classification, link prediction, and molecular property prediction settings. Its benefits also appear in oversmoothing measurements and in post-processing, where a small trained head can use cached intermediate representations.

F ADDITIONAL EXPERIMENTS

  • HISTOGRAPH+GCN consistently outperforms standard GCN and JKNet across heterophilic and homophilic node-classification datasets.Improvements are particularly pronounced on Actor, Squirrel, and Chameleon, and smaller but consistent on Cora, Citeseer, and Pubmed.Source: F ADDITIONAL EXPERIMENTS · pp. 18–19, Table 11 · p. 20
  • Feature distances are higher with HISTOGRAPH than standard GCN across layers, with the largest improvement at the pre-classifier layer.The authors present this as evidence of reduced oversmoothing.Source: F ADDITIONAL EXPERIMENTS · pp. 19–20, Table 12 · p. 20
  • HISTOGRAPH improves OGBL-COLLAB link prediction and GraphGPS molecular-property results; deeper retained histories and learned aggregation outperform alternatives across the reported ablations.The historical-layer ablation reports the best PTC results with more layers retained; HISTOGRAPH also outperforms mean aggregation and randomized attention.Source: F ADDITIONAL EXPERIMENTS · p. 20

F.1 SCALABLE POST-PROCESSING WITH HISTOGRAPH

  • Fine-tuning a post-processing head often matches or exceeds full end-to-end tuning despite using far fewer trainable parameters.At 16 layers, the reported FT accuracy is 64.7%.Source: §F.1 · p. 20
  • At extreme depths such as 64 layers, known GNN optimization challenges remain despite HISTOGRAPH’s mitigation of over-smoothing.The authors nevertheless report that HISTOGRAPH outperforms baseline pooling methods in depth-varying experiments on Cora, Citeseer, and Pubmed.Source: §F.1 · p. 20

G SCALABILITY CONSIDERATIONS

HISTOGRAPH’s dense node-wise attention and cached activations make its current design best suited to small-to-medium graphs. The paper also reports interpretable layer-attention patterns and a particularly strong PROTEINS result.

G SCALABILITY CONSIDERATIONS

  • The current design is best suited to small-to-medium graphs: node-wise attention costs O(N^2D) and activation storage costs O(NLD).For substantially larger graphs, the authors identify neighbor sampling, sparse attention, and hierarchical coarsening as possible extensions, leaving their integration to future work.Source: G SCALABILITY CONSIDERATIONS · pp. 20–21

H INTERPRETABILITY OF LEARNED ATTENTION WEIGHTS

  • Layer-wise attention places substantial weight on layers 0–5 and the final layer, while the difference between final and intermediate embeddings grows with depth.The authors interpret these patterns as task-adapted weighting of distinct early- and late-layer information.Source: H INTERPRETABILITY OF LEARNED… · p. 21

I DISCUSSION OF PROTEINS PERFORMANCE

  • HISTOGRAPH achieves a +16.6% gain over the next-best baseline on PROTEINS.The authors attribute the result to combining local motifs captured early with global fold topology captured deeper in the network.Source: I DISCUSSION OF PROTEINS PERFORMANCE · p. 21
  • On PROTEINS, removing signed normalization, layer-wise attention, or node-wise attention reduces accuracy to 74.45%, 78.61%, or 80.78%, respectively.Uniform averaging reaches 70.08% and randomized attention 80.05%, compared with HISTOGRAPH’s reported result.Source: I DISCUSSION OF PROTEINS PERFORMANCE · p. 21, Table 16 · p. 22

J INFERENCE TIME ANALYSIS

The inference-time comparison reports that HISTOGRAPH is faster than GMT on the tested 32-layer molecular datasets, while adding modest overhead relative to MeanPool.

J INFERENCE TIME ANALYSIS

  • At 32 layers on MOLHIV, HISTOGRAPH’s 3.174 ms runtime is modestly above MeanPool’s 2.130 ms.The authors attribute the added overhead to the HISTOGRAPH head, which remains small for the molecular-scale graphs evaluated.Source: J INFERENCE TIME ANALYSIS · p. 22, Table 18 · p. 23
  • At 32 layers on MOLHIV, HISTOGRAPH takes 3.174 ms versus GMT’s 15.158 ms, making it approximately 4.8× faster.The table reports average inference time for GCN backbones on MOLHIV and TOXCAST.Source: J INFERENCE TIME ANALYSIS · p. 22, Table 18 · p. 23
Loading 2601.01123v2…