TLDR: This paper systematically studies superposition in Graph Neural Networks (GNNs), a phenomenon where multiple features share the same latent space directions, hindering interpretability. Using controlled experiments and novel metrics, the researchers found that GNN width, graph topology, and pooling mechanisms significantly influence superposition. Sharper pooling (like max pooling) can reduce feature sharing and encourage axis alignment, while shallow models can suffer from “rank collapse” where representations occupy a very low-dimensional subspace. The study provides practical guidance for designing more interpretable GNNs.
Graph Neural Networks (GNNs) are powerful tools for analyzing complex data structures like social networks, molecular graphs, and more. However, understanding exactly *how* they arrive at their decisions – their internal workings – has remained a significant challenge. This difficulty largely stems from a phenomenon known as ‘superposition,’ where multiple distinct features within the network’s latent space end up sharing the same internal directions or channels, making interpretation incredibly difficult.
Decoding Superposition: A Deep Dive into GNN Interpretability
A new research paper, titled “Superposition in Graph Neural Networks,” by Lukas Pertl, Han Xuanyuan, and Pietro Liò, tackles this interpretability hurdle head-on. Unlike previous work that often focuses on identifying which nodes or edges are important, this study delves into the fundamental question of how features are represented and arranged within the GNN’s internal ‘mind.’ The authors investigate how superposition arises in GNNs, how architectural choices and graph structures influence it, and what the practical consequences are for making these models more understandable.
Investigating the Mechanisms
To systematically study superposition, the researchers developed small, controlled datasets where the underlying graph concepts were clear and unambiguous. They trained standard GNN architectures like GCN, GIN, and GAT, and then directly examined the geometry of their internal representations at both the node level (before information is aggregated for the whole graph) and the graph level (after aggregation). They used two main methods to extract features: ‘class-conditional centroids’ for task-aligned concepts and ‘linear-probe directions’ for model-decodable concepts. To quantify superposition, they introduced simple, basis-invariant metrics such as Effective Rank (how many distinct axes are used), Superposition Index (features per effective axis), and Welch-Normalized Overlap (how tightly packed directions are). They also used an Alignment Index to measure how much features prefer to align with coordinate axes.
Key Findings: Width and Topology
The study revealed several fascinating insights into how GNNs handle features:
-
Width’s Influence: As the width (number of channels in the final hidden layer) of a GNN increases, the overlap of features follows a distinct three-phase pattern. Initially, overlap decreases, then briefly increases around the point where the network’s capacity matches the number of concepts it needs to learn, and finally decreases again. Interestingly, while graph-level features can achieve near-ideal packing at high widths, node-level concepts often remain entangled. This suggests that the final ‘readout’ layers of a GNN play a crucial role in disentangling what the message-passing layers might keep mixed.
-
Topology’s Imprint: Even when GNNs have ample capacity, the underlying graph topology significantly influences how features are represented. For instance, detectors for cycles of similar lengths (e.g., triangles and squares) tend to align more closely in the node-level latent space. However, pooling mechanisms can re-mix these topological features to better align with the specific task, sometimes even causing distant concepts to align or oppose each other to help with classification.
The Role of Pooling and Rank Collapse
The research also highlighted the critical role of pooling and identified a phenomenon called ‘rank collapse’:
-
Sharper Pooling for Alignment: The way information is aggregated from nodes to form a graph-level representation (pooling) has a profound impact. Making pooling ‘sharper’ – for example, using max pooling instead of mean pooling – encourages features to align more closely with the coordinate axes and reduces the sharing of features at the node level. This happens because sharper pooling biases gradient updates towards larger magnitude coordinates, effectively pushing features to occupy distinct, axis-aligned channels.
-
Rank Collapse: The authors observed that in some cases, especially in shallow GNNs, the numerical rank of the pooled representation can remain surprisingly low, even when the model achieves high accuracy. This ‘rank collapse’ means that the learned features are confined to a very thin, low-dimensional subspace, forcing superposition. Two main factors contribute to this: global channel gating by ReLU activations (which can zero out entire channels) and a loss-driven preference for ‘mutually obtuse’ class directions that resist activating new dimensions.
Also Read:
- Making AI Decisions Clearer: The Signed Quadratic Shrink Activation
- Unveiling the Silent Thought Processes of Large Language Models
Practical Implications for GNN Design
The findings offer valuable practical guidance for designing more interpretable GNNs. Modest increases in network width can help, and using LeakyReLU instead of standard ReLU in the final MLP layers can mitigate rank collapse by preventing channels from being completely zeroed out. Furthermore, employing sharper but stable pooling mechanisms can encourage axis alignment and reduce feature sharing, leading to more disentangled and understandable representations. The paper also points to future work connecting these superposition dynamics to other GNN challenges like over-smoothing and over-squashing.
For more in-depth details, you can read the full research paper here.
