Neural Priority Queues for Graph Neural Networks
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.
Rishabh Jain, Petar Veličković, Pietro Liò
# Neural Priority Queues for Graph Neural Networks (GNNs)

###### 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.

Machine Learning, Neural Networks, Graph Neural Networks, Graph Neural Networks, Graph Neural Networks

## 1 Introduction

Algorithms and Deep Learning methods possess very fundamentally different properties. Training deep learning models to mimic algorithms would allow us to get neural models that show generalisation ability similar to the algorithms, while retaining the robustness to noise of deep learning systems. This building and training of neural networks to execute algorithmic computations is referred to as Neural Algorithmic Reasoning (Velickovic and Blundell, 2021).

Architectures that align more with the underlying algorithm for the reasoning task, tend to generalize better (Xu et al., 2019). Previous works have drawn inspiration from the external memory and data structure use of programmes and algorithms, and have found success in improving the algorithmic reasoning capabilities of recurrent neural networks (RNNs) by extending them with differentiable variants for these memory and data structures (Graves et al., 2014; Grefenstette et al., 2015).

Recently, graph neural networks (GNNs) have found immense success with algorithmic tasks (Chen et al., 2020; Velickovic et al., 2022). There have been works attempting to augment GNNs with memory, with majority using gates to do so. However, gated memory leads to very limited persistence. Furthermore, these works have solely focused on dynamic graphs, and extending these to non-dynamic graphs would involve significant effort.

In this paper, we propose the extension of the message passing framework of GNNs with external memory modules. We focus on adding a differentiable analogue to priority queues, as priority queues are a general data structure used by different algorithms and can be reduced to other data structures like stacks and queues. We name the thus formed framework for differentiable priority queues as 'Neural PQs'. We describe NPQ, an implementation under this framework, and also explore various variants for this. NPQ shows various properties that were lacking in previous works with GNNs, which we believe enable NPQ to help GNNs with algorithmic reasoning.

We summarize the contributions of this paper below:

* We propose the 'Neural PQ' framework, an extension of the message-passing GNN framework to allow use of memory modules, with particular inspiration from priority queues.
* (1) Memory-Persistence, (2) Permutation-Equivariance, (3) Reducibility to Priority Queues, and (4) No dependence on intermediate supervision. Past works have already expressed some subsets of these as desirables.
* We propose NPQs, an implementation within the Neural PQ framework, that exhibit all the above mentioned properties. This is the first differentiable analogue to priority queues, and the first memory modules for GNNs to exhibit all the above desiderata, to the best of our knowledge.
* We perform extensive quantitative analysis, via a variety of experiments and find:
* NPQs, when training to reason Dijkstra's shortest path algorithm, close the gap between the baseline test performance and ground truth by over \(40\%\).
* The various Neural PQs outperform the baseline on 26 out of 30 algorithms from the CLRS-30 dataset (Velickovic et al., 2022). The performance gains are not restricted to algorithms that actually use a priority queue.
* Neural PQs also help with long-range reasoning. These help local message-passing networks to capture long-range interaction. Thus, the benefits of using the Neural PQs are not limited to algorithmic reasoning, and these can be used on a variety of other tasks.

## 2 Background

CLRS Benchmark (Velickovic et al., 2022)Various prior works have shown the efficiency of GNNs for algorithmic tasks. However, many of these works tend to be disconnected in terms of the algorithms they target, data processing and evaluation, making direct comparisons difficult. To take the first steps in solving this issue, Velickovic et al. (2022) propose the CLRS Algorithmic Reasoning Benchmark which consists of 30 algorithms from the 'Introduction to Algorithms' textbook by Cormen et al. (2022). They name this dataset as CLRS-30.

The authors employ the encode-process-decode paradigm (Hamrick et al., 2018) and compare different processor networks (which are different GNNs) choices. Below we provide some more details on this encode-process-decode setup. Since we focus on the CLRS Benchmark for evaluation, this forms as the baseline architectural structure.

Let us take a graph \(\mathcal{G}=(\mathcal{V},\mathcal{E})\), with Let \(\mathcal{N}_{i}\) as the one-hop neighbourhood of node \(i\). Let \(\mathbf{x}_{i}\in\mathbb{R}^{d_{k}}\) be the node features for node \(i\in\mathcal{V}\), \(\mathbf{e}_{ji}\in\mathbb{R}^{d_{e}}\) the edge features for edge \((j,i)\in\mathcal{E}\) and \(\mathbf{g}\in\mathbb{R}^{d_{g}}\) the graph features. The _encode_ step involves encoding these inputs using linear layers \(f_{n}:\mathbb{R}^{d_{k}}\rightarrow\mathbb{R}^{d_{h}}\), \(f_{e}:\mathbb{R}^{d_{e}}\rightarrow\mathbb{R}^{d_{h}}\) and \(f_{g}:\mathbb{R}^{d_{g}}\rightarrow\mathbb{R}^{d_{h}}\):

\[\mathbf{h}_{i}=f_{n}(\mathbf{x}_{i})\qquad\mathbf{h}_{ij}=f_{e}(\mathbf{e}_{ij })\qquad\mathbf{h}_{g}=f_{g}(\mathbf{g}) \tag{1}\]

These are then used in a processor network during the _process_ step. The previous latent features \(\mathbf{h}_{i}^{(t-1)}\) are used along with the current node feature \(\mathbf{h}_{i}\) encoding to get a recurrent encoded input \(\mathbf{z}_{i}^{(t)}\) using a recurrent encoding function \(f_{A}\). This recurrent cell update is line with the work of Velickovic et al. (2019). A message from node \(i\) to node \(j\), \(\mathbf{m}_{ij}\) is computed for each pair of nodes using a message function \(f_{m}\). These messages are aggregated using a permutation-invariant aggregation function \(\bigoplus\). Finally, a readout function \(f_{r}\) transforms the aggregated messages and node encodings into processed node latent features.

\[\mathbf{z}_{i}^{(t)} =f_{A}(\mathbf{h}_{i},\mathbf{h}_{i}^{(t-1)}) \tag{2}\] \[\mathbf{m}_{ij} =f_{m}(\mathbf{z}_{i}^{(t)},\mathbf{z}_{j}^{(t)},\mathbf{h}_{ij },\mathbf{h}_{g})\] (3) \[\mathbf{m}_{i} =\bigoplus_{j\in\mathcal{N}_{i}}\mathbf{m}_{ji}\] (4) \[\mathbf{h}_{i}^{(t)} =f_{r}(\mathbf{z}_{i}^{(t)},\mathbf{m}_{i}) \tag{5}\]

For different processors, \(f_{A}\), \(f_{m}\) and \(f_{r}\) may differ.

The last _decode_ step consists of using relevant decoding functions to get the required prediction. This might be the predicted hints or predicted output.

## 3 Related Work

The ability of RNNs to work in a sequential manner led to their popularity in previous works for reasoning about algorithms, as algorithms tend to be iterative in nature. Noting that most computer programmes make use of external memory, Graves et al. (2014) proposed addition of an external memory module to RNNs, which makes reasoning about algorithms easier. Subsequent methods have worked upon this idea and have found success with memory modules inspired from different data structures. This includes Stack-Augmented RNNs by (Joulin and Mikolov, 2015) and Neural DeQues by Grefenstette et al. (2015). A key limitation of all these proposals is that they are only defined for

Figure 1: **Left: A GNN processor based on the message passing framework. At each timestep, pair-wise messages are formed using the node features. These messages are aggregated, and then used to update the node features for the next timestep. Right: The memory module framework we propose. The memory module takes the node features and previous memory state, to output the next memory state and messages for the GNN processor. These messages are aggregated alongside the traditional node-to-node messages. In this project, we focus on memory modules inspired from priority queues.**
use by RNNs. Unlike GNNs, RNNs are unable to use the structured information about the algorithms' input spaces.

Early explorations on augmenting GNNs with memory focused on the use of internal memory in the form of gates, such as Gated Graph Sequence Networks by Li et al. (2015), and Temporal Graph Networks by Rossi et al. (2020). However, the use of such RNN-like gate mechanisms limits the persistence of the graph/node histories. Persistent Message Passing (PMP) by Strathmann et al. (2021) is a noteworthy GNN that makes use of non-gated persistent external memory, by persisting some of the nodes at each timestep. However, PMPs cannot be applied to non-dynamic graphs without significant effort. Furthermore, they require intermediate-supervision.

## 4 Neural PQ Framework

Previous works have proposed Neural Stacks, Queues and DeQues (Grefenstette et al., 2015; Joulin and Mikolov, 2015) that have a RNN controller. In this project, we propose the use of memory modules, with focus on differentiable PQs (or Neural PQs), with a GNN model acting as the controller. Furthermore, we propose integration of such memory modules with message-passing framework by allowing the Neural PQ to send messages to each node. The setup for this is shown in Figure 2.

### Desiderata

We form the framework with the following desiderata in mind - (1) Memory-Persistence, (2) Permutation-Equivariance, (3) Reducibility to Priority Queues, and (4) No dependence on intermediate supervision. We motivate the need for these below.

Memory-PersistenceMemory-Persistence is necessary to make full use of the extended capacity provided by the external memory modules. This is especially true for models running over several timestep, with a long temporal interaction, where we would want to access memory added at a much earlier timestep. Furthermore, memory persistence helps avoid over-smoothing. Node embeddings of GNNs tend to start off varied. As more messages are passed, the embeddings tend to converge to each-other, thus making the nodes indistinguishable, as the number of layers increases. Memory-persistence would allow GNNs to remember older states, when the embeddings were more distinguished, and use these to promote more varied embeddings, despite the depth of the model.

Permutation-EquivarianceA GNN layer is said to be equivariant to permutation of the nodes if and only if any permutation of the node IDs, while maintaining the overall graph structure, leads to the same permutation of the node features. This reflects one of the most basic symmetries of graph structures, and is necessary to ensure that isomorphic graphs receive the same representation, up to certain permutations and transformation. This makes permutation-equivariance an essential property for GNN layers. Thus, we want our Neural PQs to also be equivariant to permutations.

Priority Queue AlignmentPriority queues are a general data structure used by various algorithms. Furthermore, different data structures can be modelled using priority queues, like stacks and queues are priority queues with the time of insertion as the priority. A memory module that aligns well with priority queues would lead to the overall model, that uses the memory module, to align with the related algorithms better. Algorithmically aligned models lead to greater generalisation (Xu et al., 2019). Thus, this is essential for greater algorithmic reasoning.

No intermediate supervision requirementBy not requiring any intermediate supervision, memory modules become easier to apply directly to different algorithmic tasks. This allows us to even use the Neural PQ for reasoning over algorithms that may not use a priority queue themselves. General Neural PQs can be helpful with such tasks due to additional properties, like memory-persistence.

Figure 2: Neural PQ controlled using a GNN processor. At each timestep, the node features are used to pop values from the priority queue. These values are used to form the messages that are sent to the different nodes. The node features and previous state are also used to determine what values to push, and update the priority queue. This uses Neural PQ as the memory module in Figure 1.
### Framework

We present the framework for a Neural PQ controlled by a message-passing GNN below. We use the equations for the baseline message-passing GNN as described in Section 2. Let us suppose we have the same setup as the baseline. The encode and decode part remain the same, but now the processor uses an implementation of the Neural PQ Framework. Let the graph in consideration be \(\mathcal{G}=(\mathcal{V},\mathcal{E})\). Let the previous hidden state of the GNN be \(\mathbf{h}_{i}^{(t-1)}\), and the previous state of the Neural PQ be \(\mathbf{h}_{pq}^{(t-1)}\).

In the Neural PQ framework, we calculate the set of values \(V_{i}\) to be popped for each node \(i\in\mathcal{V}\), using a pop function \(f_{pop}\). Messages are formed from these popped values using a message encoding function \(f_{M}^{(pq)}\). Each node aggregates these messages along with the traditional node-to-node pair-wise messages. Lastly, a push function \(f_{push}\) updates the state of the Neural PQ, to obtain the next state \(\mathbf{h}_{pq}^{(t)}\). Formally, we can define the following equations for the framework:

\[\mathbf{z}_{i}^{(t)} =f_{A}(\mathbf{h}_{i},\mathbf{h}_{i}^{(t-1)}) \tag{6}\] \[\mathbf{m}_{ij} =f_{m}(\mathbf{z}_{i}^{(t)},\mathbf{z}_{i}^{(t)},\mathbf{h}_{ij},\mathbf{h}_{g})\] (7) \[V_{i} =f_{pop}(\mathbf{z}_{i}^{(t)},\mathbf{z}^{(t)},\mathbf{h}_{pq}^ {(t-1)})\] (8) \[M_{i} =f_{M}^{(pq)}(V_{i},\mathbf{z}_{i}^{(t)})\,\cup\,\{\mathbf{m}_{ji} \,|\,j\in\mathcal{N}_{i}\}\] (9) \[\mathbf{m}_{i} =\bigoplus_{\mathbf{m}\in M_{i}}\mathbf{m}\] (10) \[\mathbf{h}_{i}^{(t)} =f_{r}(\mathbf{z}_{i}^{(t)},\mathbf{m}_{i})\] (11) \[\mathbf{h}_{pq}^{(t)} =f_{push}(\mathbf{h}_{pq}^{(t-1)},\mathbf{z}^{(t)}) \tag{12}\]

where \(\mathbf{z}^{(t)}\) is a multi-set of all the encoded inputs \(\mathbf{z}_{i}^{(t)}\), i.e. \(\mathbf{z}^{(t)}=\{\{\mathbf{z}_{i}^{(t)}\,|\,i\in\mathcal{V}\}\}\). \(f_{A},f_{m}\) and \(f_{r}\) depend on which message-passing GNN processor we choose, while \(f_{pop}\), \(f_{M}^{(pq)}\) and \(f_{push}\) depend on the Neural PQ implementation.

Note that, in the above proposed framework, we choose to delay the update of the priority queue due to pop operation until the \(f_{push}\). This is done to keep the Neural PQ framework general and to segregate the queue read and update operations. This also allows us to prove the permutation-equivariance properties of the framework, as discussed below.

Even though the presented framework is inspired from priority queues, we can implement various other data structures, like queues and stacks, by appropriate \(f_{pop}\) and \(f_{push}\) definitions. The Neural PQ framework exhibits and promotes various properties from the desiderata. By design, these do not require any additional supervision. Furthermore, since the push, pop and message encoding functions only depend on the destination node's features, the multi-set of all node features and the Neural PQ state, all implementations are also equivariant to permutations of the nodes, under certain assumptions. For a detailed proof, refer to Appendix A.

## 5 Npq

We propose NPQ, an implementation following the foremtioned Neural PQ framework that exhibits all the proposed desiderata. We divide the overall definition of NPQ into 4 sub-sections - (1) State, (2) Pop Function, (3) Message Encoding Function, and (4) Push Function. Taking inspiration from Neural DeQues (Grefenstette et al., 2015), NPQs consist of continuous push and pop operations.

### State

The state of the NPQ must hold all the memory values pushed into the queue. Alongside these values, since we define continuous push and pop operations, we also need to keep track of the strengths/proportions of each queue element still present. We can represent this state \(\mathbf{h}_{pq}^{(t)}\) as a tuple of the list of memory values \(\mathbf{v}^{(t)}=[\mathbf{v}_{1}^{(t)},\ldots,\mathbf{v}_{i}^{(t)},\ldots]\) and the list of strengths of these memory values \(\mathbf{s}^{(t)}=[\mathbf{s}_{1}^{(t)},\ldots,\mathbf{s}_{i}^{(t)},\ldots]\). The \(i\)th element of \(\mathbf{v}^{(t)}\) is \(\mathbf{v}_{i}^{(t)}\), the value of the \(i\)th element of the priority queue, and the \(i\)th element of \(\mathbf{s}^{(t)}\) is \(\mathbf{s}_{i}^{(t)}\in(0,1)\), the strength of the \(i\)th element of the priority queue.

\[\mathbf{h}_{pq}^{(t)}=\langle\mathbf{v}^{(t)},\mathbf{s}^{(t)}\rangle \tag{13}\]

where \(\langle\cdot\rangle\) is a tuple. Figure 3 shows a sample state for the NPQ.

### Pop function

We propose a continuos pop function, i.e. we pop a fractional proportions of the values in the queue. This fraction, \(s_{pop}^{(i)}\in(0,1)\) for node \(i\), is computed as noted in Equation 14. This equation is similar to the ones used for Neural DeQues by Grefenstette et al. (2015).

\[s_{pop}^{(i)}=\text{ sigmoid}\left(f_{s}^{(pop)}(\mathbf{z}_{i}^{(t)})\right) \tag{14}\]

We use a request-grant framework to maintain the constraint that no value can be popped more than it is present, i.e. one cannot pop \(0.7\) of a value \(v_{j}\) that may be present in the PQ with only a strength of \(s_{j}=0.4\). Each node \(i\in\mathcal{V}\) requests to pop fraction \(p_{j}^{(i)}\in(0,1)\) of PQ element \(j\). NPQ takes

Figure 3: **State of the NPQ.** It consists of two lists of the same length, representing the values \(\mathbf{v}\) in the queue and their respective strengths \(\mathbf{s}\). In the above example, the NPQ consists of three elements, \(\mathbf{v}_{1}\), \(\mathbf{v}_{2}\) and \(\mathbf{v}_{3}\), with strengths \(0.6\), \(0.8\) and \(0.3\) respectively.
all the \(p_{j}^{(i)}\) values into consideration and grants a fraction \(q_{j}^{(i)}\in(0,1)\) of PQ element \(j\) to node \(i\), which may or may not be the same as the requested \(p_{j}^{(i)}\). Equation 15 shows the calculation of this granted fraction given the requested fractions \(p_{j}^{(i)}\) and the PQ element strengths \(\mathbf{s}_{j}^{(t-1)}\).

\[q_{j}^{(i)}=\left\{\begin{array}{cl}p_{j}^{(i)}&\text{, if }\sum_{k\in\mathcal{V }}p_{j}^{(k)}\leq\mathbf{s}_{j}^{(t-1)}\\ \frac{p_{j}^{(i)}}{\sum_{k\in\mathcal{V}}p_{j}^{(k)}}\cdot\mathbf{s}_{j}^{(t- 1)}&\text{, else}\end{array}\right. \tag{15}\]

The main idea behind this equation is that ideally we would want to satisfy each request \(p_{j}^{(i)}\) for popping. We cannot do that when the sum of pop requests is greater than the strength with which the element is present in the PQ. In this case, we completely pop the element and return to each node fraction of this value in proportion to the strength each node requested. It maintains the requirement that \(q_{j}^{(i)}\leq p_{j}^{(i)}\) and \(\sum_{i\in\mathcal{V}}q_{j}^{(i)}\leq\mathbf{s}_{j}^{(t-1)}\). These granted proportions are used to calculate the final value popped.

\[\mathbf{v}=\sum_{j\in\mathcal{I}_{pq}^{(t)}}q_{j}^{(i)}\cdot \mathbf{v}_{j}^{(t-1)} \tag{16}\] \[f_{pop}(\mathbf{z}_{i}^{(t)},\mathbf{z}^{(t)},\left\langle \mathbf{v}^{(t-1)},\mathbf{s}^{(t-1)}\right\rangle)=\{\mathbf{v}\} \tag{17}\]

where \(\mathcal{I}_{pq}^{(t-1)}=\left[1,\ldots,\left\lfloor\mathbf{v}^{(t-1)}\right \rfloor\right]\) is the set of indices for the NPQ. Note that since we are only popping a single value from the NPQ, we are returning a single element set. The requested pop proportions \(p_{j}^{(i)}\) are calculated using the continuous pop strength value \(s_{pop}^{(i)}\), and attention coefficients \(c_{j}^{(i)}\in(0,1)\), denoting the coefficient for the \(j\)th element of the queue with respect to the \(i\)th node. These are calculated using a multi-head additive attention mechanism (Bahdanau et al., 2014; Vaswani et al., 2017). This is done with inspiration from GATs by Velickovic et al. (2018).

\[e_{j}^{(i,h)}= \text{ LeakyReLU}\left(f_{a_{1}}^{(h)}(\mathbf{z}_{i}^{(t)})+f_{a_{2}}^{(h) }(\mathbf{v}_{j}^{(t-1)})\right) \tag{18}\] \[\alpha_{j}^{(i,h)}= \text{ softmax}_{j}\left(e_{j}^{(i,h)}\right)\] (19) \[c_{j}^{(i)}= \text{ softmax}_{j}\left(f_{a}([\alpha_{j}^{(i,1)},\ldots,\alpha_ {j}^{(i,h)},\ldots])\right) \tag{20}\]

where \(\alpha_{j}^{(i,h)}\in(0,1)\) is the attention coefficients for \(j\)th element of the queue with respect to node \(i\) via attention-head \(h\), and \(f_{a_{1}}^{(h)}\), \(f_{a_{2}}^{(h)}\) and \(f_{a}\) are linear layers.

Using these coefficients, we propose two ways of popping elements from the queue - Max Popping and Weighted Popping. We refer to NPQ using max popping and weighted popping as NPQ\({}_{\text{M}}\) and NPQ\({}_{\text{W}}\), respectively.

#### Max Popping

The element \(j\) of the queue with the highest attention coefficient \(c_{j}^{(i)}\) is requested to be popped for the node \(i\).

\[k =\operatorname*{argmax}_{k\in\mathcal{I}_{pq}^{(t-1)}}c_{k}^{(i)} \tag{21}\] \[p_{j}^{(i)} =s_{pop}^{(i)}\cdot\mathbb{I}_{\{k\}}\left(j\right) \tag{22}\]

where \(\mathbb{I}_{A}(\cdot)\) is the indicator function for set \(A\), i.e. \(\mathbb{I}_{A}(a)=1\iff a\in A\) and \(\mathbb{I}_{A}(a)=0\iff a\notin A\).

### Weighted Popping

The attention coefficients are treated as soft-weights with which each element in the PQ is requested to be popped.

\[p_{j}^{(i)}=s_{pop}^{(i)}\cdot c_{j}^{(i)} \tag{23}\]

Figure 4 shows sample pop operation for NPQ.

### Priority Queue Message Function

We use a simple message encoding function, where each output is passed through a linear layer \(f_{m}^{(pq)}\).

\[f_{M}^{(pq)}(V_{i},\mathbf{z}_{i}^{(t)})=\left\{f_{m}^{(pq)}(\mathbf{v})\,|\, \mathbf{v}\in V_{i}\right\} \tag{24}\]

### Push function

As mentioned earlier, the push function is actually the state update function. Here we first delete the popped proportions from the NPQ. Let \(\mathbf{h}_{pq}^{(t-1)}=\left\langle\mathbf{v}^{(t-1)},\mathbf{s}^{(t-1)}\right\rangle\) be the previous

Figure 4: Sample **pop operation** for a single node \(i\) in NPQ\({}_{\text{M}}\). **Left:** Pop-request \(p_{j}^{(i)}\) generation. First we compute the attention coefficients \(c_{j}^{(i)}\) using the node features \(\mathbf{z}_{i}\) and the memory values \(\mathbf{v}_{j}^{(t-1)}\). We also calculate the pop-strength \(s_{pop}^{(i)}\). The coefficients and the pop strength are together used to determine the pop fractions to request. NPQ\({}_{\text{M}}\) uses Max Popping, which only pops the element with the highest attention coefficient, in this case \(\mathbf{v}_{2}\). Thus, we request popping of only this element with strength \(s_{pop}^{(i)}\). **Right:** Using pop requests from all the nodes and the strengths of each value currently in the queue, NPQ grants certain fraction \(q_{j}^{(i)}\) to node \(i\) to pop element \(j\). We use this granted proportion to determine the value popped.
NPQ state. Then, we can define the NPQ state with the popped proportions deleted as \(\langle\mathbf{v^{\prime}},\mathbf{s^{\prime}}\rangle\), which are calculated as below.

\[\mathbf{s^{\prime}}_{i} =\mathbf{s}_{i}^{(t-1)}-\sum_{k\in\mathcal{V}}q_{i}^{(k)} \tag{25}\] \[\mathbf{s^{\prime}} =\text{ nonzero}_{i}\left(\mathbf{s^{\prime}_{i}}\right)\] (26) \[\mathbf{v^{\prime}} =\mathbf{v^{(t)}}[\text{ arg-nonzero}_{i}\left(\mathbf{s^{ \prime}_{i}}\right)] \tag{27}\]

where \(q_{i}^{(k)}\in(0,1)\) is the proportion NPQ element \(i\) granted to be popped for node \(k\) as defined in Equation 15, \(\text{nonzero}_{i}(\mathbf{s^{\prime}_{i}})\) is sequence of \(\mathbf{s^{\prime}_{i}}\) with all zero \(\mathbf{s^{\prime}_{i}}\) removed, and similarly, \(\text{arg-nonzero}_{i}\) is the relevant indices of the sequence.

We push a single value \(\mathbf{v}\) for the whole graph. To determine this value, we pass each node embedding through a linear layer \(f_{v}\) and sum the formed values across all the nodes. In line with Neural DeQues by Grefenstette et al. (2015), this values is activated using a tanh function to get the final value to be pushed.

The push function is continuous and so requires calculation of the push strength \(\mathbf{s}_{push}\). This is done in a similar manner to the push values calculation, using a linear layer \(f_{s}^{(push)}\). We use a logistic sigmoid activation here instead of tanh, akin to Neural DeQues.

\[\mathbf{v}=\text{ tanh}\left(\sum_{i\in\mathcal{V}}f_{v}(\mathbf{z }_{i}^{(t)})\right) \tag{28}\] \[\mathbf{s}=\text{ sigmoid}\left(\sum_{i\in\mathcal{V}}f_{s}^{( push)}(\mathbf{z}_{i}^{(t)})\right)\] (29) \[f_{push}(\langle\mathbf{v}^{(t-1)},\mathbf{s}^{(t-1)}\rangle, \mathbf{z}^{(t)})=\langle\mathbf{v^{\prime}}||[\mathbf{v}],\mathbf{s^{\prime} }||[\mathbf{s}_{push}]\rangle. \tag{30}\]

Note that in the above equations, we use \(q_{i}^{(k)}\) and \(\mathbf{z}_{i}^{(t)}\), which are not actually inputs to the \(f_{push}\) function. This is done mainly to maintain readability of the functions. These equations can be easily reformulated to only use \(\textbf{h}_{pq}^{(t-1)}\) and \(\mathbf{z}^{(t)}\), in order to follow the general Neural PQ framework. Refer to Appendix B for the reformulation.

### Properties

Simply by virtue of following the Neural PQ framework, NPQ exhibits two of the desiderata - Permutation-Equivariance, and no dependence on intermediate supervision. We do not update or replace the previously stored NPQ elements, but rather persist them as long as possible, and only delete their proportions when we pop them. This allows NPQ to achieve much greater memory-persistence than done using gated memories.

Lastly, the push and pop operations of the NPQ are defined to be aligned close to the push and pop operations of the traditional priority queue. In fact, under some assumptions, we can prove that \(\text{NPQ}_{\text{M}}\) can be reduced to a traditional priority queue. This can be done by taking the push and pop functions to be encoding the key-value pairs for the priority queue elements. For a detailed proof, refer to Appendix C.

Thus, NPQ satisfies the four stated desiderata.

### Variants

We also explore some variations on the proposed NPQ. One such variation involves consideration of greater memory-persistence by not deleting the popped elements. We refer to this variation as NPQ-P.

Notably, NPQ treats popping as a node-wise activity. We can instead treat popping as a graph operation, i.e. each node receives the same set of popped values. This can be done by either sending all the node-wise popped values to all the nodes, or by popping a single value for all the nodes. We refer to these two variants as NPQ-SA and NPQ-SV, respectively.

Empirically, we found these latter two variants more useful when combined with the first one. We refer to these combined variations as NPQ-P-SA and NPQ-P-SV, respectively. For the exact equations for the variants, refer to Appendix D-F.

Figure 5: Sample **push operation**. The overall push operation can be divided into two main steps. **Step 1:** In the first step, we update the queue strengths to reflect the removal of popped fractions of each element. This is done by summing the granted pop fractions \(q_{j}^{(i)}\) from all nodes. The granted pop-fractions \(q_{j}^{(i)}\) can be calculated from the previous NPQ state and the node embeddings \(\mathbf{z}\). These aggregated pop-grants are then subtracted from the strengths of the respective queue elements. Some elements might end up with \(0\) strength, and these are then removed from the queue, as shown here with the greyed-out value \(\mathbf{v}_{2}\) here. **Step 2:** The next step is to actually push a value into the queue. The value to be pushed and its strength are determined by the node embeddings \(\mathbf{z}\).
## 6 Evaluation

The main hypothesis we test is whether the Neural PQ implementations are useful for algorithmic reasoning by using the CLRS-30 dataset (Velickovic et al., 2022). To do so, we undertake multiple experiments - (1) We first focus on a single algorithm, Dijkstra's Shortest Path algorithm, evaluating the performance of the Neural PQs with a MLP MPNN as the base GNN, comparing them with the MPNN baseline as well as an MLP MPNN with an oracle priority queue. (2) We also evaluate the performance of the Neural PQs on rest of the algorithms from the CLRS benchmark. (3) Lastly, we also test whether the Neural PQs are useful for long-range reasoning, by evaluating their performance on a dataset from the Long Range Graph Benchmark (Dwivedi et al., 2022). Appendix G shows some more experiments performed.

### Dijkstra's Algorithm - MPNN Base

We train the models on Dijkstra's algorithm from CLRS-30, and test for out-of-distribution (OOD) generalisation, i.e. the models are trained on smaller input graphs, containing 16 nodes, and tested on larger graphs, containing 256 nodes. The training data consists of 1000 samples, while the testing and validation data consist of 32 samples each. We test the models on larger graph sizes than done by Velickovic et al. (2022) (they use graphs with 64 nodes) to better test the generalisation ability, and because baseline MPNN model already gets around \(91.5\%\) test performance with 64 nodes.

To test the limit of attainable performance from Neural PQs, we test an MPNN with access to an Oracle PQ, where apart from the standard input features, we also take information about the values pushed and popped from the priority queue as input. The Oracle Neural PQ forces the push and pop operation to be determined by the algorithmic PQ. This information about the actual PQ is used in training, validation as well as testing.

Table 1 shows the test performance of the different models. We see that the last model performs much better than the early-stopped model for the baseline and Oracle PQ. Notably, the last and early-stopped model perform similarly for NPQ\({}_{\text{W}}\). NPQ\({}_{\text{W}}\) outperforms the baseline as well as the Oracle PQ. In fact, we see that it **closes the gap** between the test performance of baseline MPNN and true solution, i.e. 100% test performance, by over **40\(\%\)**.

### Different Algorithms from CLRS-30

We train and test five models for each algorithm from CLRS-30 dataset - 'Baseline' (no memory module), NPQ\({}_{\text{M}}\)-P-SA, NPQ\({}_{\text{W}}\)-P-SV, NPQ\({}_{\text{W}}\) and NPQ\({}_{\text{M}}\). We train each model on graphs with 16 nodes, and test them on graphs with 128 nodes, and consider only the early-stopped models.

Figure 6 shows the comparison for best performing Neural PQ and the baseline MPNN for each algorithm. We see that for **26 out of the 30 algorithms, at least one of the Neural PQs outperforms the baseline MPNN**. Interestingly, the optimal Neural PQ version depends on the algorithm of choice. Notably, the performance gain of using a Neural PQ does not seem to be limited to algorithms that use a traditional priority queue. This supports our belief that the Neural PQ implementations are quite general, and these can perform various roles, such as acting as a traditional data structure, or a persistent-memory for accessing past overall graph states. We provide the table with algorithm-wise performance of each Neural PQ in Appendix G. Focussing on NPQ\({}_{\text{M}}\), we found that it outperforms the baseline for \(17\) algorithms (more than half of the algorithms). We see that for \(12\) algorithms, it improves the performance or closes the gap to true prediction by at least \(10\%\). For \(4\) algorithms, it improves performance/reduces the gap by at least \(50\%\).

### Long-Range Reasoning

Message-passing based GNNs exchange information between 1-hop neighbours to build node representations at each layer. Past works have shown that such information propagation leads to over-squashing when the path of information traversal is long, and so such models perform poorly on tasks requiring long-range interaction (Alon and Yahav, 2021; Dwivedi et al., 2022). Dwivedi et al. (2022) have proposed a collection of graph learning datasets to form 'Long Range Graph Benchmark' (LRGB), each of which arguably require long-range interaction reasoning to achieve strong performance. In these experiments, we test the performance of using Neural PQs on Peptides-struct dataset from the LRGB benchmark.

Figure 7 shows the test MAE results for the different Neural PQs and the baseline. Notably, all Neural PQs outperform the baseline for GATv2 processor, while only NPQ\({}_{\text{W}}\)-P

\begin{table}
\begin{tabular}{l l l} \hline
**Method** & **Best** & **Last** \\ \hline Baseline & \(68.58\%\pm 9.71\) & \(76.97\%\pm 4.38\) \\ NPQ\({}_{\text{W}}\) & \(\mathbf{85.48\%\pm 3.50}\) & \(\mathbf{86.22\%\pm 2.20}\) \\ NPQ\({}_{\text{M}}\) & \(74.54\%\pm 8.37\) & \(74.68\%\pm 3.06\) \\ NPQ\({}_{\text{W}}\)-SA & \(79.74\%\pm 2.99\) & \(69.36\%\pm 12.02\) \\ NPQ\({}_{\text{W}}\)-SV & \(77.04\%\pm 2.99\) & \(79.89\%\pm 6.28\) \\ NPQ\({}_{\text{M}}\)-P-SA & \(79.19\%\pm 5.17\) & \(78.26\%\pm 6.19\) \\ NPQ\({}_{\text{W}}\)-P-SV & \(71.46\%\pm 6.75\) & \(79.44\%\pm 5.90\) \\ \hline \hline Oracle PQ & \(75.85\%\pm 3.65\) & \(85.37\%\pm 3.72\) \\ \hline \end{tabular}
\end{table}
Table 1: Test performance (Mean \(\pm\) Standard Deviation) of models with MLP MPNN base on learning Dijkstra’s Algorithm with 256 node graphs, run with 3 different seeds. The table shows the results for the _Best_ validation score model (early-stopped model) and the _Last_ model in training.
SV and \(\text{NPQ}_{\text{W}}\) outperform the baseline on the other two processors. The success of \(\text{NPQ}_{\text{W}}\)-P-SV and \(\text{NPQ}_{\text{W}}\) means that these Neural PQs are empirically helping the models with long-range reasoning. Notably, we see that Weighted popping seems more useful for long-range reasoning.

## 7 Conclusion and Future Works

External memory modules have helped traditional RNNs improve their algorithmic reasoning capabilities. A natural hypothesis would be that external memory modules can also help graph neural networks (GNNs) with algorithmic reasoning. However, this remains a largely unexplored domain. In this paper, we proposed Neural PQs, a general framework for adding memory modules to GNNs, with inspirations from traditional priority queues. We proposed and motivated a desiderata for memory modules, and presented NPQ, an implementation with the Neural PQ framework that exhibited the desiderata.

We empirically show that NPQs indeed help with algorithmic reasoning, and without any extra supervision, matches the performance of the baseline model that has access to true priority queue operations on Dijkstra's algorithm. The performance gains are not limited to algorithms using priority queues. Furthermore, we show that the Neural PQs help with capturing long-range interaction, by demonstrating their prowess on the Peptides-struct dataset from the Long-Range Graph Benchmark.

The success of the Neural PQs has a wide effect on the field of representational learning. It opens up a research domain exploring the use of memory modules with GNNs, especially their interfacing with the message-passing framework. The Neural PQs take crucial steps towards advancing the neural algorithmic reasoning field. These also hold potential with various other fields and tasks, as seen by their performance on the long-range reasoning task.

We have limited our focus on simple memory module operations. Potential future works could involve exploration of more complicated definitions. These definitions might be formed by analysing the reasons behind the greater success of Neural PQs on some algorithms as opposed to others. Neural PQs can also be used for various other graph tasks, and it would be interesting to explore their uses for these.

Figure 6: Evaluation results for best performing Neural PQ and the baseline MPNN model for the 30 algorithms from CLRS-30, sorted by the relative improvement in performance.

Figure 7: Test MAE (Mean \(\pm\) Standard Deviation) of different Neural PQs with different base processors on Peptides-struct dataset, run with 3 different seeds. **Lower** the test MAE, **better** is the performance.
## Acknowledgements

We thank Adria Puigdomenech and Karl Tuyls for reviewing the paper prior to the submission.
