TLDR: Graph Neural Networks (GNNs) suffer from “over-squashing,” an information bottleneck. This paper redefines over-squashing into sensitivity and storage capacity issues, focusing on the latter. It introduces a new task, Neighbor Associative Recall (NAR), to measure storage capacity and proposes gLSTM, a GNN architecture inspired by xLSTM and associative memory. gLSTM significantly improves information storage and retrieval, demonstrating superior performance on NAR and real-world benchmarks by mitigating capacity over-squashing.
Graph Neural Networks, or GNNs, are powerful tools for understanding and learning from data structured as graphs, like social networks or molecular structures. These models typically work by passing messages between connected nodes, allowing them to gather information from their neighbors and update their own representations. While GNNs have found widespread use, they often encounter a significant challenge known as “over-squashing.”
Over-squashing occurs when a large amount of information from a node’s extensive network of connections, called its receptive field, is forced into a single, fixed-size data packet or vector. This compression creates an information bottleneck, leading to a loss of crucial details as the information travels through the network. Imagine trying to fit a detailed map of a city into a tiny sticky note – much of the important information would be lost.
Traditionally, over-squashing has been viewed through two main lenses: reduced sensitivity and saturating storage capacity. Reduced sensitivity refers to how well a node’s output can respond to changes in its input, often linked to issues like vanishing gradients. This paper, however, brings the focus back to the second aspect: saturating storage capacity. This is defined as the amount of information a node’s representation can hold for later use. When this capacity is saturated, the node simply cannot store any more incoming information.
To better understand and measure this capacity over-squashing in isolation, the researchers introduced a new synthetic task called Neighbor Associative Recall (NAR). Unlike previous tasks that often conflated sensitivity and capacity issues, NAR is specifically designed to test a GNN’s ability to store and retrieve information from its immediate neighbors in a shallow network, avoiding the complexities of deep networks and their associated sensitivity problems.
Inspired by advancements in sequence modeling, particularly the xLSTM model, the authors developed a novel GNN architecture named gLSTM. This new model aims to directly address the capacity limitations of traditional GNNs by incorporating associative memory. In essence, gLSTM equips each node with not just a standard vector hidden state, but also a matrix hidden state. This matrix acts like a memory bank, storing “key-value” pairs derived from the node’s previous states and its neighbors. When a node needs to retrieve information, it “queries” this memory bank, allowing it to selectively recall relevant data.
A key feature of gLSTM is its K-hop aggregation scheme. Instead of just gathering information from immediate neighbors, gLSTM can aggregate information from nodes up to ‘K’ hops away. This creates a highly connected computational graph, which not only helps mitigate sensitivity over-squashing but also provides a beneficial structure for gLSTM’s recall mechanism. Information stored in the associative memory is not repeatedly passed through subsequent layers, allowing later nodes to query this memory independently.
Experiments on the NAR task demonstrated that gLSTM significantly outperforms traditional GNNs like GCN in recall abilities. gLSTM maintained perfect recall until the number of neighbors exceeded its memory dimension, showing a graceful degradation thereafter. In contrast, GCN models experienced capacity over-squashing much earlier. Furthermore, the study empirically showed that capacity over-squashing can occur independently of sensitivity over-squashing, as gLSTM’s sensitivity often continued to increase even after its performance started to decline due to memory saturation.
Beyond synthetic tasks, gLSTM also showed strong performance on real-world benchmarks designed to test long-range interactions, such as the Graph Property Prediction (GPP) tasks and the Long Range Graph Benchmark (LRGB). It achieved state-of-the-art results on Diameter and Eccentricity GPP tasks and strong performance on SSSP, highlighting its ability to overcome information bottlenecks in practical applications.
Also Read:
- NHA: A Unified Approach to Long and Short Context in AI
- Bridging Efficiency and Detail in Long-Context AI with Artificial Hippocampus Networks
This research redefines over-squashing by separating it into distinct issues of sensitivity and storage capacity. By focusing on and enhancing the storage capacity of GNNs through associative memory, the gLSTM architecture offers a promising new direction for developing more robust and capable graph learning models. For more technical details, you can read the full paper here.
