Table of Contents
Fetching ...

Massive Activations in Graph Neural Networks: Decoding Attention for Domain-Dependent Interpretability

Lorenzo Bini, Marco Sorbi, Stephane Marchand-Maillet

TL;DR

The study addresses how attention in edge-featured GNNs yields Massive Activations (MAs) and how these outliers relate to domain-specific edge information. By establishing a baseline with untrained models, normalizing activations by edge medians, and modeling the base distribution with a Gamma law, the authors detect training-induced MA distribution shifts and validate their findings across GraphTransformer, GraphiT, and SAN on ZINC, TOX21, and PROTEINS. They introduce an Explicit Bias Term (EBT) to stabilize activation magnitudes and demonstrate that MAs align with edge informativeness, notably concentrating on common bond types in molecular graphs, while enabling post-hoc interpretability through edge-type heatmaps and ablations. The work offers actionable insights for leveraging MA patterns to improve model transparency, guide augmentations, and relate GNN internals to domain knowledge, with broad implications for downstream science tasks.

Abstract

Graph Neural Networks (GNNs) have become increasingly popular for effectively modeling graph-structured data, and attention mechanisms have been pivotal in enabling these models to capture complex patterns. In our study, we reveal a critical yet underexplored consequence of integrating attention into edge-featured GNNs: the emergence of Massive Activations (MAs) within attention layers. By developing a novel method for detecting MAs on edge features, we show that these extreme activations are not only activation anomalies but encode domain-relevant signals. Our post-hoc interpretability analysis demonstrates that, in molecular graphs, MAs aggregate predominantly on common bond types (e.g., single and double bonds) while sparing more informative ones (e.g., triple bonds). Furthermore, our ablation studies confirm that MAs can serve as natural attribution indicators, reallocating to less informative edges. Our study assesses various edge-featured attention-based GNN models using benchmark datasets, including ZINC, TOX21, and PROTEINS. Key contributions include (1) establishing the direct link between attention mechanisms and MAs generation in edge-featured GNNs, (2) developing a robust definition and detection method for MAs enabling reliable post-hoc interpretability. Overall, our study reveals the complex interplay between attention mechanisms, edge-featured GNNs model, and MAs emergence, providing crucial insights for relating GNNs internals to domain knowledge.

Massive Activations in Graph Neural Networks: Decoding Attention for Domain-Dependent Interpretability

TL;DR

The study addresses how attention in edge-featured GNNs yields Massive Activations (MAs) and how these outliers relate to domain-specific edge information. By establishing a baseline with untrained models, normalizing activations by edge medians, and modeling the base distribution with a Gamma law, the authors detect training-induced MA distribution shifts and validate their findings across GraphTransformer, GraphiT, and SAN on ZINC, TOX21, and PROTEINS. They introduce an Explicit Bias Term (EBT) to stabilize activation magnitudes and demonstrate that MAs align with edge informativeness, notably concentrating on common bond types in molecular graphs, while enabling post-hoc interpretability through edge-type heatmaps and ablations. The work offers actionable insights for leveraging MA patterns to improve model transparency, guide augmentations, and relate GNN internals to domain knowledge, with broad implications for downstream science tasks.

Abstract

Graph Neural Networks (GNNs) have become increasingly popular for effectively modeling graph-structured data, and attention mechanisms have been pivotal in enabling these models to capture complex patterns. In our study, we reveal a critical yet underexplored consequence of integrating attention into edge-featured GNNs: the emergence of Massive Activations (MAs) within attention layers. By developing a novel method for detecting MAs on edge features, we show that these extreme activations are not only activation anomalies but encode domain-relevant signals. Our post-hoc interpretability analysis demonstrates that, in molecular graphs, MAs aggregate predominantly on common bond types (e.g., single and double bonds) while sparing more informative ones (e.g., triple bonds). Furthermore, our ablation studies confirm that MAs can serve as natural attribution indicators, reallocating to less informative edges. Our study assesses various edge-featured attention-based GNN models using benchmark datasets, including ZINC, TOX21, and PROTEINS. Key contributions include (1) establishing the direct link between attention mechanisms and MAs generation in edge-featured GNNs, (2) developing a robust definition and detection method for MAs enabling reliable post-hoc interpretability. Overall, our study reveals the complex interplay between attention mechanisms, edge-featured GNNs model, and MAs emergence, providing crucial insights for relating GNNs internals to domain knowledge.
Paper Structure (21 sections, 9 equations, 7 figures, 1 table)

This paper contains 21 sections, 9 equations, 7 figures, 1 table.

Figures (7)

  • Figure 1: Comparison of activation distributions between base and trained model. The trained model presents two long tails representing MAs values. Notice the $y$-axis is log-scaled, making the base model activations distribution clustered around zero.
  • Figure 2: Comparison of MAs on trained against base models. Represented ratios have been sorted increasingly for each layer independently. Orange box represents the range of normal ratios obtained in the base model, while ratios exceeding the base come from MAs.
  • Figure 3: Activation distributions for base and trained (with MAs) models. In \ref{['fig:gammad']} we clearly distinguish a spike on the left of the distribution, corresponding to a ratio of 1000 (-$\log(\text{ratio}) = -3$), which identifies the separation between the basic and massive regimes. The approximation pdf is rescaled to match the histogram scale.
  • Figure 4: Heatmaps showing MAs concentration across the three edge types in the ZINC dataset. Each heatmap visualizes the percentage of edges with MAs per attention head and hidden feature dimension. Notably, MAs predominantly aggregate on edge types 1 and 2, while being absent on type 3, indicating the model's tendency to allocate MAs to more frequent edges. Type $3$ consists of $0.26\%$ of edges in the dataset.
  • Figure 5: Ablation studies show the reallocation of MAs on the chemically meaningless edges (types 4 and 5), designed to carry low intrinsic information. This supports the hypothesis that MAs can serve as indicators of edge importance.
  • ...and 2 more figures