Neural Priority Queues for Graph Neural Networks
This work addresses the problem of enhancing GNNs with explicit memory for neural algorithmic reasoning, which is incremental as it builds on existing GNN frameworks.
The paper tackles the limited exploration of augmenting Graph Neural Networks (GNNs) with external memory by introducing Neural Priority Queues, a differentiable analogue to algorithmic priority queues, and demonstrates empirical results on the CLRS-30 dataset and improved long-range interactions on the Long-Range Graph Benchmark.
Graph Neural Networks (GNNs) have shown considerable success in neural algorithmic reasoning. Many traditional algorithms make use of an explicit memory in the form of a data structure. However, there has been limited exploration on augmenting GNNs with external memory. In this paper, we present Neural Priority Queues, a differentiable analogue to algorithmic priority queues, for GNNs. We propose and motivate a desiderata for memory modules, and show that Neural PQs exhibit the desiderata, and reason about their use with algorithmic reasoning. This is further demonstrated by empirical results on the CLRS-30 dataset. Furthermore, we find the Neural PQs useful in capturing long-range interactions, as empirically shown on a dataset from the Long-Range Graph Benchmark.