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.
