Neural Priority Queues for Graph Neural Networks

TL;DR

This paper introduces Neural Priority Queues, a differentiable memory module for Graph Neural Networks, enhancing their ability to perform algorithmic reasoning and capture long-range interactions, demonstrated through empirical results on multiple datasets.

Contribution

The paper proposes Neural Priority Queues as a novel differentiable memory structure for GNNs, enabling improved algorithmic reasoning and long-range interaction modeling.

Findings

Neural PQs satisfy key desiderata for memory modules.

Neural PQs improve reasoning on CLRS-30 dataset.

Neural PQs effectively model long-range interactions on LRGB dataset.

Abstract

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.