Discrete Neural Algorithmic Reasoning

## Abstract

Abstract Neural algorithmic reasoning aims to capture computations with neural networks via learning the models to imitate the execution of classic algorithms. While common architectures are expressive enough to contain the correct model in the weights space, current neural reasoners are struggling to generalize well on out-of-distribution data. On the other hand, classic computations are not affected by distributional shifts as they can be described as transitions between discrete computational states. In this work, we propose to force neural reasoners to maintain the execution trajectory as a combination of finite predefined states. To achieve that, we separate discrete and continuous data flows and describe the interaction between them. Trained with supervision on the algorithm’s state transitions, such models are able to perfectly align with the original algorithm. To show this, we evaluate our approach on multiple algorithmic problems and get perfect test scores both in single-task and multitask setups. Moreover, the proposed architectural choice allows us to prove the correctness of the learned algorithms for any test data.

## 1 Introduction

Learning to capture algorithmic dependencies in data and to perform algorithmic-like computations with neural networks are core problems in machine learning, studied for a long time using various approaches (Roni Khardon, ; Graves et al., ; Zaremba & Sutskever, ; Reed & De Freitas, ; Kaiser & Sutskever, ; Veličković et al., ).

Neural algorithmic reasoning (Veličković & Blundell, ) is a research area focusing on building models capable of executing classic algorithms. Relying on strong theoretical guarantees of algorithms to work correctly on any input of any size and distribution, this setting provides unlimited challenges for out-of-distribution generalization of neural networks. Prior work explored this setup using the CLRS-30 benchmark (Veličković et al., ), which covers classic algorithms from the Introduction to Algorithms textbook (Cormen et al., ) and uses graphs as a universal tool to encode data of various types. Importantly, CLRS-30 also provides the decomposition of classic algorithms into subroutines and simple transitions between consecutive execution steps, called hints, which can be used during training in various forms.

The core idea of the CLRS-30 benchmark is to understand how neural reasoners generalize well beyond the training distribution, namely on larger graphs. Classic algorithms possess strong generalization due to the guarantee that correct execution steps never encounter ‘out-of-distribution’ states, as all state transitions are predefined by the algorithm. In contrast, when encountering inputs from distributions that significantly differ from the train data, neural networks are usually not capable of robustly maintaining internal calculations in the desired domain. Consequently, due to the complex and diverse nature of all possible data that neural reasoners can be tested on, the generalization performance of such models can vary depending on particular test distribution (Mahdavi et al., ).

Given that, it is becoming important to interpret internal computations of neural reasoners to find errors or to prove the correctness of the learned algorithms (Georgiev et al., ).

Interpretation methods have been actively developing recently due to various real-world applications of neural networks and the need to debug and maintain systems based on them. Especially, the Transformer architecture (Vaswani et al., ) demonstrates state-of-the-art performance in natural language processing and other modalities, representing a field for the development of interpretability methods (Elhage et al., ; Weiss et al., ; Zhou et al., ; Lindner et al., ). Based on active research on a computational model behind the transformer architecture, recent works propose a way to learn models that are fully interpretable by design (Friedman et al., ).

We found the ability to design models that are interpretable in a simple and formalized way to be crucial for neural algorithmic reasoning as it is naturally related to the goal of learning to perform computations with neural networks.

In this paper, we propose to force neural reasoners to follow the execution trajectory as a combination of finite predefined states, which is important for both generalization ability and interpretability of neural reasoners. To achieve that, we start with an attention-based neural network and describe three building blocks to enhance its generalization abilities: feature discretization, hard attention and separating discrete and continuous data flows. In short, all mentioned blocks are connected:

- •
State discretization does not allow the model to use complex and redundant dependencies in data;
- •
Hard attention is needed to ensure that attention weights will not be annealed for larger graphs. Also, hard attention limits the set of possible messages that each node can receive;
- •
Separating discrete and continuous flows is needed to ensure that state discretization does not lose information about continuous data.

Then, we build fully discrete neural reasoners for different algorithmic tasks and demonstrate their ability to perfectly mimic ground-truth algorithm execution. As a result, we achieve perfect test scores on the multiple algorithmic tasks with guarantees of correctness on any test data. Moreover, we demonstrate that a single network is capable of executing all covered algorithms in a multitask manner with perfect generalization too.

In summary, we consider the proposed blocks as a crucial component for robust and interpretable neural reasoners and demonstrate that trained with hint supervision, discretized models perfectly capture the dynamic of the underlying algorithms and do not suffer from distributional shifts.

## 2 Background

### 2.1 Algorithmic reasoning

Performing algorithmic-like computations usually requires the execution of sequential steps and the number of such steps depends on the input size. To imitate such computations, neural networks are expected to be based on some form of recurrent unit, which can be applied to a particular problem instance several times (Kaiser & Sutskever, ; Zaremba & Sutskever, ; Vinyals et al., ; Veličković et al., ).

The CLRS Algorithmic Reasoning Benchmark (CLRS-30) (Veličković et al., ) defines a general paradigm of algorithmic modeling based on Graph Neural Networks (GNNs), as graphs can naturally represent different input types and manipulations over such inputs. Also, GNNs are proven to be well-suited for neural execution (Xu et al., ; Dudzik & Veličković, ).

The CLRS-30 benchmark covers different algorithms over various domains (arrays, strings, graphs) and formulates them as algorithms over graphs. Also, CLRS-30 proposes to utilize the decomposition of the algorithmic trajectory execution into simple logical steps, called hints. Using this decomposition is expected to better align the model to desired computations and prevent it from utilizing hidden non-generalizable dependencies of a particular train set. Prior work demonstrates a wide variety of additional inductive biases for models towards generalizing computations, including different forms of hint usage (Veličković et al., ; Bevilacqua et al., ), biases from standard data structures (Jürß et al., ; Jain et al., ), knowledge transfer and multitasking (Xhonneux et al., ; Ibarz et al., ; Numeroso et al., ), etc. Also, recent studies demonstrate several benefits of learning neural reasoners end-to-end without any hints at all (Mahdavi et al., ; Rodionov & Prokhorenkova, ).

The recently proposed SALSA-CLRS benchmark (Minder et al., ) enables a more thorough OOD evaluation compared to CLRS-30 with increased test sizes (up to 100-fold train-to-test scaling, compared to 4-fold for CLRS-30) and diverse test distributions. Despite significant gains in the performance of neural reasoners in recent work, current models still struggle to generalize to out-of-distribution (OOD) test data (Mahdavi et al., ; Georgiev et al., ; Minder et al., ). While de Luca & Fountoulakis () prove by construction the ability of the transformer-based neural reasoners to perfectly simulate graph algorithms (with minor limitations occurring from the finite precision), it is still unclear if generalizable and interpretable models can be obtained via learning. Importantly, the issues of OOD generalization are induced not only by the challenges of capturing the algorithmic dependencies in the data but also by the need to carefully operate with continuous inputs. For example, investigating the simplest scenario of learning to emulate the addition of real numbers, Klindt () demonstrates the failure of some models to exactly imitate the desired computations due to the nature of gradient-based optimization. This limitation can significantly affect the performance of neural reasoners on adversarial examples and larger input instances when small errors can be accumulated.

### 2.2 Transformer interpretability and computation model

Transformer (Vaswani et al., ) is a neural network architecture for processing sequential data. The input to the transformer is a sequence of tokens from a discrete vocabulary. The input layer maps each token to a high-dimensional embedding and aggregates it with the positional encoding.
The key components of each layer are attention blocks and MLP with residual connections. Providing a detailed description of mechanisms learned by transformer models (Elhage et al., ) is of great interest due to their widespread applications.

RASP (Weiss et al., ) is a programming language proposed as a high-level formalization of the computational model behind transformers. The main primitives of RASP are elementwise sequence functions, $select$ and $aggregate$ operations, which conceptually relate to computations performed by different blocks of the model. Later, Lindner et al. () presented Tracr, a compiler for converting RASP programs to the weights of the transformer model, which can be useful for evaluating interpretability methods.

While RASP might have limited expressibility, it supports arbitrary complex continuous functions which in theory can be represented by transformer architecture, but are difficult to learn.
Also, RASP is designed to formalize computations over sequences of fixed length. Motivated by that, Zhou et al. () proposed RASP-L, a restricted version of RASP, which aims to formalize the computations that are easy to learn with transformers in a size-generalized way. The authors also conjecture that the length-generalization of transformers on algorithmic problems is related to the ‘simplicity’ of solving these problems in RASP-L language.

Another recent work (Friedman et al., ) describes Transformer Programs: constrained transformers that can be trained using gradient-based optimization and then automatically converted into a discrete, human-readable program. Built on RASP, transformer programs are not designed to be size-invariant.

## 3 Discrete neural algorithmic reasoning

### 3.1 Encode-process-decode paradigm

Our work follows the encode-process-decode paradigm (Hamrick et al., ), which is usually employed for step-by-step neural execution.

All input data is represented as a graph $G$ with an adjacency matrix $A$ and node and edge features that are first mapped with a simple linear encoder to high-dimensional vectors of size $h$.
Let us denote node features at a time step $t$ ($1\leq t\leq T$) as $X^{t}=(x_{1}^{t},\dots,x_{n}^{t})$ and edge features as $E_{t}=(e_{1}^{t},\dots,e_{m}^{t})$. Then, the processor, usually a single-layer GNN, recurrently updates these features, producing node and edge features for the next step:

$$ $X^{t+1},E^{t+1}=\mathrm{Processor}(X^{t},E^{t},A).$ $$

The processor network can operate on the original graph defined by the task (for graph problems) or on the fully connected graph. For the latter option, the information about the original graph can be encoded into the edge features.

The number of processor steps $T$ can be defined automatically by the processor or externally (e.g., as the number of steps of the original algorithm). After the last step, the node and edge features are mapped with another linear layer, called the decoder, to the output predictions of the model.

If the model is trained with hint supervision, the changes of node and edge features at each step are expected to be related to the original algorithm execution. In this sense, the processor network is aimed to mimic the algorithm’s execution in the latent space.

### 3.2 Discrete neural algorithmic reasoners

In this section, we describe the constraints for the processor that allow us to achieve a fully interpretable neural reasoner. We start with Transformer Convolution (Shi et al., ) with a single attention head.

As mentioned above, at each computation step $t$ $(1\leq t\leq T),$ the processor takes the high-dimensional embedding vectors for node and edge features as inputs and then outputs the representations for the next execution step.

Each node feature vector $x_{i}$ is projected into $query$ $(Q_{i})$, $key$ $(K_{i})$, and $value$ $(V_{i})$ vectors via learnable parameter matrices $W_{Q}$, $W_{K}$, and $W_{V}$, respectively. Edge features $e_{ij}$ are projected into $key$ $(K_{ij})$ vector with a matrix $W_{K}^{E}$. Then, for each directed edge from node $j$ to node $i$ in the graph $G$, we compute the attention coefficient

$$ $\alpha_{ij}=\frac{\langle Q_{j},K_{i}+K_{ij}\rangle}{\sqrt{h}},$ $$

where $\langle a,b\rangle$ denotes the dot product. Then, each node $i$ normalizes all attention coefficients across its neighbors with the softmax function and temperature $\tau$ and receives the aggregated message:

$$ $\hat{\alpha}_{ij}=\frac{\exp(\alpha_{ij}/\tau)}{\sum_{k\in\textit{N}(i)}\exp( \alpha_{ik}/\tau)},\,\,\,\,\,\,M_{i}=\sum_{k\in\mathcal{N}(i)}\hat{\alpha}_{ik }V_{k},$ (1) $$

where $\mathcal{N}(i)$ denotes the set of all incoming neighbors of node $i$ and $M_{i}$ is the message sent to the $i$-th node.

For undirected graphs, we consider two separate edges in each direction. Also, for each node, we consider a self-loop connecting the node to itself. For multi-head attention, each head $l$ separately computes the messages $M_{i}^{l}$ which are then concatenated.

Similar to Transformer Programs, we enforce attention to be hard attention. We found this property important not only for interpretability but also for size generalization, as hard attention allows us to overcome the annealing of the attention weights for arbitrarily large graphs and strictly limits the set of messages that each vertex can receive.

After message computation, node and edge features are updated depending on the current values and sent messages using feed-forward MLP blocks:

$$ $\displaystyle\hat{x}_{i}^{t+1}$ $\displaystyle=\mathrm{FFN}_{nodes}([x_{i}^{t},M_{i}^{t}]),$ $\displaystyle\hat{e}_{ij}^{t+1}$ $\displaystyle=\mathrm{FFN}_{edges}([e_{ij}^{t},\hat{\alpha}_{ji}V_{i},\hat{ \alpha}_{ij}V_{j}]).$ $$

We also enforce all node and edge features to be from a fixed finite set, which we call *states*. We ensure such property by adding discrete bottlenecks at the end of the processor block:

$$ $\displaystyle x_{i}^{t+1}$ $\displaystyle=\mathrm{Discretize}_{nodes}(\hat{x}_{i}^{t+1}),$ $\displaystyle e_{ij}^{t+1}$ $\displaystyle=\mathrm{Discretize}_{edges}(\hat{e}_{ij}^{t+1}).$ $$

We implement discretization by projecting the features to the vectors of size $k$ which we force to be one-hot using annealing Gumbel-Softmax (Jang et al., ) during training and the argmax at the inference.

### 3.3 Continuous inputs

Clearly, most of the algorithmic problems operate with continuous or unbounded inputs (e.g., weights on edges). Usually, all input data is encoded into node and edge features and the processor operates over the resulting vectors. The proposed discretization of such vectors would lead to the loss of information necessary for performing correct execution steps. One possible option to operate with such inputs (we will call them *scalars*, meaning both continuous or size-dependent integer inputs, such as node indexes) is Neural Execution Engines (Yan et al., ), which allows one to operate with bit-wise representations of integer and (in theory) real numbers. Such representations are bounded by design, but fully discrete and interpretable.

We propose another option: to maintain scalar inputs (denote them by $S$) separately from the node and edge features and use them only as edge priorities $s_{ij}$ in the attention block. If scalars are related to the nodes, we assign them to edges depending on the scalar of the sender or receiver node. Now, we can consider the hard attention block as a selector which for each node selects the best edge based on an ordered pair of ‘states priority’ (attention weights described above, which depend only on states of the corresponding nodes and edges) and $s_{ij}$. We note that this selector is related to the theoretical primitive $select_best$ from RASP. We implement this simply by augmenting key vectors $K_{ij}$ of each edge with the indicator if the given edge has the “best” (min or max) scalar among the other edges to node $j$. Thus, scalars affect only the attention weights, not the messages and the node states.

For multiple different scalar inputs (e.g., weighted edges and node indexes), we use multi-head attention, where each head operates with separate scalars.

Given that, the interface of the proposed processor can be described as

$$ $X^{t+1},E^{t+1}=\mathrm{Processor}(X^{t},E^{t},A,S),$ $$

where $X^{t},X^{t+1}$ and $E^{t},E^{t+1}$ are from the fixed sets. State sets are independent of the execution step $t$ and the input graph (including scalar inputs $S$).

### 3.4 Manipulations over continuous inputs

Figure: Figure 1: An illustration of the proposed separation between discrete and continuous data flows. Scalars can only affect the attention weights (Green) and can be modified with actions via $\mathrm{ScalarUpdate}$ (Blue).

The proposed selector offers a read-only interface to scalar inputs, which is not expressive enough for most of the algorithms. However, we note that the algorithms can be described as *discrete* manipulations over input data. For example, the Dijkstra algorithm (Dijkstra, ) takes edge weights as inputs and uses them to find the shortest path distances. Computed distances can affect the consequent execution steps. We note that such distances can be described as the sum of the weights of the edges that form the shortest path to the given node. In other words, the produced scalars depend only on input scalars and discrete execution states.

To avoid the challenges of learning continuous updates with high precision (Klindt, ), we propose to learn discrete manipulations with scalars. The updated scalars can then be used with the described selector in the next steps.

In our experiments, we use a scalar updater capable of incrementing, moving, and adding scalars depending on discrete node/edge states:

$$ $\displaystyle s_{i}^{t+1}$ $\displaystyle=\mathrm{inc}(x_{i}^{t})+\mathrm{keep}(x_{i}^{t})\cdot s_{i}^{t}+ \sum_{j\in\mathcal{N}(i)}\mathrm{push}(e_{ji}^{t})\cdot s_{ji}^{t},$ $\displaystyle s_{ij}^{t+1}$ $\displaystyle=\mathrm{inc}(e_{ij}^{t})+\mathrm{keep}(e_{ij}^{t})\cdot s_{ij}^{ t}+\mathrm{push}(x_{i}^{t})\cdot s_{i}^{t},$ $$

where $s_{i}$ are node-related scalars, $s_{ij}$ are edge-related scalars, and $\mathrm{inc}$, $\mathrm{keep}$, $\mathrm{push}$ are 0-1 functions representing if scalar in each node/edge should be incremented, kept, or pushed to any of its neighbors. We implement these functions as simple linear projections of node/edge features with consecutive discretization.

Finally, our proposed method (see Figure 1) can be described as:

$$ $X^{t+1},E^{t+1}=\mathrm{Processor}(X^{t},E^{t},A,S^{t}),$ $$

$$ $S^{t+1}=\mathrm{ScalarUpdate}(X^{t+1},E^{t+1},A,S^{t}).$ $$

The proposed neural reasoners are fully discrete and can be interpreted by design. Moreover, the proposed selector block guarantees the predicted behavior of the message passing for any graph size, as it compares discrete state importance and uses continuous scalars only to break ties between equally important states.

## 4 Experiments

In this section, we perform experiments to evaluate how the proposed discretization affects the performance of neural reasoners on diverse algorithmic reasoning tasks.
Our main questions are:

- 1.
Can the proposed discrete neural reasoners capture the desired algorithmic dynamics with hint supervision?
- 2.
How does discretization affect OOD and size-generalization performance of neural reasoners?
- 3.
Is the proposed model capable of multi-task learning?

Also, we are interested if discrete neural reasoners can be learned without hints and how they tend to utilize the given amount of node and edge states to solve the problem. We discuss no-hint experiments separately in Section 6.

### 4.1 Datasets

We perform our experiments on the problems from the recently proposed SALSA-CLRS benchmark (Minder et al., ), namely BFS, DFS, Prim, Dijkstra, Maximum Independent Set (MIS), and Eccentricity. We believe that the proposed method is not limited by the covered problems, but we leave the implementation of the required data flows (e.g., edge-based reasoning (Ibarz et al., ), graph-level hints, interactions between different scalars) for future work.

The train dataset of SALSA-CLRS consists of random graphs with at most 16 nodes sampled from the Erdös-Renyi (ER) distribution with parameter $p$ chosen to be as low as possible while graphs remain connected with high probability. The test set consists of sparse graphs of sizes from 16 to 1600 nodes.

We slightly modify the hints from the benchmark without conceptual changes (e.g., we have modified the hints for the DFS problem to remove graph-level hints). Discrete states are fully described by the non-scalar hints and scalars are exactly the hints of the $scalar$ type (we refer to Veličković et al. () for the details on hint design).

### 4.2 Baselines and evaluation

We compare the performance of our proposed discrete model with two baseline sparse models, GIN (Xu et al., ) and Pointer Graph Network (Veličković et al., ). We report both node-level and graph-level metrics for the baselines and our model.
Also, we compare our model with Triplet-GMPNN (Ibarz et al., ) and two recent approaches, namely Hint-ReLIC (Bevilacqua et al., ) and G-ForgetNet (Bohde et al., ), which demonstrate state-of-the-art performance of hint-based neural algorithmic reasoning. However, as these methods are evaluated on the CLRS-30 benchmark and their code is not yet publicly available, we can only compare them on the corresponding tasks (BFS, DFS, Dijkstra, Prim) and CLRS-30 test data, namely ER graphs with $p=0.5$ of size 64. Note that this test data is more dense than that of SALSA-CLRS, meaning shorter roll-outs for given tasks. Also, only node-level metrics have been reported for these methods.

### 4.3 Model details

For our experiments, we use the model described in Section 3. We use one attention head for each scalar value. The number of processor steps is defined externally as the length of the ground truth algorithm trajectory, which is consistent with the prior work. We use one architecture (except the task-depended encoders/decoders), including the $\mathrm{ScalarUpdate}$ module for all the problems.

We recall that neural reasoners can operate either on the base graph (which is defined by the problem) or on more dense graphs with the original graph encoded into edge features. SALSA-CLRS proposes to enhance the size-generalization abilities of neural reasoners with increased test sizes (up to 100-fold train-to-test scaling), so we use the base graph for message-passing, similar to the SALSA-CLRS baselines. We also add a virtual node that communicates with all nodes of the graph.

### 4.4 Training details and hyperparameters

We train each model using the Adam optimizer, learning rate $\eta=0.001$, using teacher forcing, batch size of 32 graphs with 1000 optimization steps, and evaluate the resulting model. We anneal softmax temperatures in the discrete bottlenecks (attention weights, $\mathrm{ScalarUpdate}$ operations) geometrically from 3.0 to 0.01, decreasing the temperature at each training step. We report all hyperparameters in the source code.(^1^11[https://github.com/yandex-research/dnar](https://github.com/yandex-research/dnar))

**Table 1: Node $\backslash$ graph level test scores for the proposed discrete reasoner and the baselines on SALSA-CLRS test data. Scores are averaged across 5 different seeds, standard deviation is omitted. For the eccentricity problem, only the graph-level metric is applicable.**
| Task | Size | GIN | PGN | DNAR (ours) |
| --- | --- | --- | --- | --- |
| BFS | 16 | 98.8 $\backslash$ 92.5 | 100. $\backslash$ 100. | 100. $\backslash$ 100. |
|  | 80 | 95.3 $\backslash$ 59.4 | 99.8 $\backslash$ 88.1 | 100. $\backslash$ 100. |
|  | 160 | 95.1 $\backslash$ 37.8 | 99.6 $\backslash$ 66.3 | 100. $\backslash$ 100. |
|  | 800 | 86.9 $\backslash$ 0.9 | 98.7 $\backslash$ 0.2 | 100. $\backslash$ 100. |
|  | 1600 | 86.5 $\backslash$ 0.0 | 98.5 $\backslash$ 0.0 | 100. $\backslash$ 100. |
| DFS | 16 | 41.5 $\backslash$ 0.0 | 82.0 $\backslash$ 19.9 | 100. $\backslash$ 100. |
|  | 80 | 30.4 $\backslash$ 0.0 | 38.4 $\backslash$ 0.0 | 100. $\backslash$ 100. |
|  | 160 | 20.0 $\backslash$ 0.0 | 26.9 $\backslash$ 0.0 | 100. $\backslash$ 100. |
|  | 800 | 19.5 $\backslash$ 0.0 | 24.9 $\backslash$ 0.0 | 100. $\backslash$ 100. |
|  | 1600 | 17.8 $\backslash$ 0.0 | 23.1 $\backslash$ 0.0 | 100. $\backslash$ 100. |
| SP | 16 | 95.2 $\backslash$ 49.8 | 99.3 $\backslash$ 89.5 | 100. $\backslash$ 100. |
|  | 80 | 62.4 $\backslash$ 0.0 | 94.2 $\backslash$ 3.3 | 100. $\backslash$ 100. |
|  | 160 | 53.3 $\backslash$ 0.0 | 92.0 $\backslash$ 0.0 | 100. $\backslash$ 100. |
|  | 800 | 40.4 $\backslash$ 0.0 | 87.1 $\backslash$ 0.0 | 100. $\backslash$ 100. |
|  | 1600 | 36.9 $\backslash$ 0.0 | 84.5 $\backslash$ 0.0 | 100. $\backslash$ 100. |
| Prim | 16 | 89.6 $\backslash$ 29.7 | 96.4 $\backslash$ 69.9 | 100. $\backslash$ 100. |
|  | 80 | 51.6 $\backslash$ 0.0 | 79.7 $\backslash$ 0.0 | 100. $\backslash$ 100. |
|  | 160 | 49.5 $\backslash$ 0.0 | 75.6 $\backslash$ 0.0 | 100. $\backslash$ 100. |
|  | 800 | 45.0 $\backslash$ 0.0 | 69.5 $\backslash$ 0.0 | 100. $\backslash$ 100. |
|  | 1600 | 43.2 $\backslash$ 0.0 | 66.8 $\backslash$ 0.0 | 100. $\backslash$ 100. |
| MIS | 16 | 79.9 $\backslash$ 3.3 | 99.8 $\backslash$ 98.6 | 100. $\backslash$ 100. |
|  | 80 | 79.9 $\backslash$ 20.0 | 99.4 $\backslash$ 88.9 | 100. $\backslash$ 100. |
|  | 160 | 78.2 $\backslash$ 0.0 | 99.4 $\backslash$ 76.2 | 100. $\backslash$ 100. |
|  | 800 | 83.4 $\backslash$ 0.0 | 98.8 $\backslash$ 18.0 | 100. $\backslash$ 100. |
|  | 1600 | 79.2 $\backslash$ 0.0 | 98.9 $\backslash$ 5.2 | 100. $\backslash$ 100. |
| Ecc. | 16 | 25.3 | 100. | 100. |
|  | 80 | 23.8 | 100. | 100. |
|  | 160 | 26.1 | 100. | 100. |
|  | 800 | 17.1 | 100. | 100. |
|  | 1600 | 16.0 | 83.0 | 100. |

During training, we minimize the standard hints and output losses: scalar hints are optimized with MSE loss, and other types of hints are optimized with cross-entropy and categorized cross-entropy losses (Veličković et al., ; Ibarz et al., ). Note that we do not supervise any additional details in model behavior, e.g., selecting the most important neighbor in the attention block, the exact operations with scalars, etc.

For multitask experiments, we follow the setup proposed by Ibarz et al. () and train a single processor with task-dependent encoders/decoders. We make 10000 optimization steps on the accumulated gradients across each task and keep all hyperparameters the same as in the single task.

Our models are trained on a single A100 GPU, requiring less than 1 hour for single-task and 5-6 hours for multitask training.

### 4.5 Results

We found learning with teacher forcing suitable for discrete neural reasoners, as discretization blocks allow us to perform the exact transitions between the states. Trained with step-wise hint supervision, discrete neural reasoners are able to perfectly align with the original algorithm and generalize on larger test data without any performance loss. We report the evaluation results in Tables 1 and 2. Also, our multitask experiments show that the proposed discrete models are capable of multitask learning and demonstrate the perfect generalization scores in a multitask manner too.

To summarize, we have two key components of our contribution: separating discrete and continuous data flows and discrete bottlenecks. Note that using discrete bottlenecks without the separation between flows is not applicable, as would lead to the loss of the continuous information, as we mention in Section 3. On the other hand, we ablate the discretization block to verify its contribution to the performance of the proposed model. Our experiments show that the proposed processor without the discretization performs non-above the baselines level on all the tasks.

**Table 2: Node-level test scores for the proposed discrete reasoner and the baselines on CLRS-30 test data. Test graphs are of size 64. Scores are averaged across 5 different seeds.**
| Task | Triplet-GMPNN | Hint-ReLIC | G-ForgetNet | DNAR (ours) |
| --- | --- | --- | --- | --- |
| BFS | 99.73 $\pm$ 0.0 | 99.00 $\pm$ 0.2 | 99.96 $\pm$ 0.0 | 100. $\pm$ 0.0 |
| DFS | 47.79 $\pm$ 4.2 | 100. $\pm$ 0.0 | 74.31 $\pm$ 5.0 | 100. $\pm$ 0.0 |
| Dijkstra | 96.05 $\pm$ 0.6 | 97.74 $\pm$ 0.5 | 99.14 $\pm$ 0.1 | 100. $\pm$ 0.0 |
| MST-Prim | 86.39 $\pm$ 1.3 | 87.97 $\pm$ 2.9 | 95.19 $\pm$ 0.3 | 100. $\pm$ 0.0 |

## 5 Interpetability and testing

In addition to the empirical evaluation of the trained models on diverse test data, the discrete and size-independent design of the proposed models allows us to interpret and test them manually. The main idea is to show that the model will perform the exact discrete state transitions (including discrete operations with scalars) as the original algorithm.

First, we note that due to the hard attention, each node receives exactly one message. Also, the message depends only on the discrete states of the corresponding nodes and edges. Thereby, as each node and edge states after a single processor step depend only on the current states and received message, all the possible options can be directly enumerated and tested if all states change to the correct ones.

The only remaining part to fully interpret the whole model is the attention block. We note that our implementation of the $select_best$ selector (Section 3.3) does not necessarily produce the top-1 choice over the ordered pairs of ‘states priority’ and scalars $s_{ij}$ as it simply augments key vectors with indicators if the given edge has the best scalar among others. For example, it may happen that for some state, the maximum attention weight is achieved for an edge without the indicator. However, given the finite number of discrete states, we can manually check if the mentioned “best” indicator increases the attention weight between every pair of states. Combined with hard attention, this would imply that the attention block attends depending on the predefined states and uses scalar priorities only to break ties.

Given that, we can unit-test all possible state transitions and attention blocks. With full coverage of such tests, we can guarantee the correct execution of the desired algorithm for *any* test data. We tested our trained models from Section 4 manually verifying state transitions. As a result, we confirm that the attention block indeed operates as $select_best$ selector, as the model actually uses these indicators to increase the attention weights. Thus, we can guarantee that for any graph size, the model will mirror the desired algorithm, which is correct for any test size.

## 6 Towards no-hint discrete reasoners

In this section, we discuss the challenges of training discrete reasoners without hints, which can be useful when tackling new algorithmic problems.

Training deep discretized models is known to be challenging, therefore we focus only on the BFS algorithm, as it is well-aligned with the message-passing framework, has short roll-outs, and can be solved with small states count (note that for no-hint models node/edge states count is a hyperparameter due to the absence of the ground truth states trajectory).

We recall that the output of the BFS problem is the exploration tree pointing from each node to its parent. Each node chooses as a parent the neighbor from the previous distance layer with the smallest index.

We perform hyperparameters search over training graph sizes, discrete states count, softmax temperature annealing schedules, and training steps. The trained models never achieved the perfect validation scores. Then, we selected the best-performing models and tried to reverse-engineer how they utilize the given states. In short, we found that the model operating on $k$ node states tends to use them to assign the layers from the starting node, except for the nodes on a distance greater than $k$, where state updates do not mirror the expected BFS flow. Also, the major part of correctly predicted pointers is concentrated on the first $k$ layers. We found this behavior well-aligned to the BFS algorithm and indicating the possibility of achieving the perfect validation score with enough states count. However, this algorithm does not generalize since it fails at distances larger than those encountered during training.

On the other hand, one can demonstrate that for a small enough state count (for BFS, it is two node states and two edge states) and diverse enough validation data, the perfect validation performance implies that the learned solution will generalize to any graph size.

Therefore, we highlight the need to achieve perfect validation performance with models that use as few states as possible, which corresponds to the minimum description length (MDL) theory (Myung, ; Rissanen, ) and is related to the notion of Kolmogorov complexity (Kolmogorov, ).

## 7 Limitations and future work

#### Limitations

In this work, we propose a method to learn robust neural reasoners that demonstrate perfect generalization performance and are interpretable by design. In this section, we describe some limitations of our work and important directions for future research.

First, several proposed design choices strictly reduce the expressive power of the model. For example, due to the hard attention, the proposed model is unable to compute the mean value from all the neighbors in a single message-passing step, which is trivial for attention-based models (note that this can be computed in several message-passing steps). Thus, the model in its current form is unable to express transitions between hints for some algorithms from the CLRS-30 in a single processor step. However, we believe that the expressivity of the proposed model can be enhanced with additional architectural modifications (e.g., edge-based reasoning (Ibarz et al., ), global states, interactions between different scalars) that can be combined with the proposed discretization ideas.

Second, while we report the perfect scores for the covered tasks, we cannot guarantee that the training will converge to the correct model for any initialization/training data distributions. However, we empirically found the proposed method to be quite robust to various architecture/training hyperparameter choices.

#### Future work

Our method is based on the particular architectural choice and actively utilizes the attention mechanism. However, the graph deep learning field is rich in various architectures exploiting different inductive biases and computation flows. The proposed separation between discrete states and continuous inputs may apply to other models, however, any particular construction can require additional efforts.

Also, we provide only one example of the $\mathrm{ScalarUpdate}$ block. We believe that utilizing a general architecture (e.g., some form of discrete Neural Turing Machine (Graves et al., ; Gulcehre et al., )) capable of executing a wider range of manipulation is of interest for future work.

With the development of neural reasoners and their ability to execute classic algorithms on abstract data, it is becoming more important to investigate how such models can be applicable in real-world scenarios according to the Neural Algorithmic Reasoning blueprint (Veličković & Blundell, ) and transfer their knowledge to high-dimensional noisy data with intrinsic algorithmic dependencies. While there are several examples of NAR-based models tackling real-world problems (Beurer-Kellner et al., ; Numeroso et al., ), there are no established benchmarks for extensive evaluation and comparison of different approaches.

Lastly, we leave for future work a deeper investigation of learning interpretable neural reasoners without hints, which we consider essential from both theoretical perspective and practical applications, e.g., combinatorial optimization.

## 8 Conclusion

In this paper, we force neural reasoners to maintain the execution trajectory as a combination of finite predefined states. To achieve that, we separate discrete and continuous data flows and describe the interaction between them. The obtained discrete reasoners are interpretable by design. Moreover, trained with hint supervision, such models perfectly capture the dynamic of the underlying algorithms and do not suffer from distributional shifts. We consider discretization of hidden representations as a crucial component for robust neural reasoners.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
questions: 1. How does the proposed method handle algorithms that require more complex continuous manipulations or aggregate functions beyond simple increments and selections? Can the model be extended to support such algorithms without losing the benefits of discretization and interpretability?
2. Are there potential strategies to improve training without hint supervision? How does the model perform on tasks without hints when compared to continuous models?
3. The constraints improve generalization but reduce expressiveness. Is there a way to balance this trade-off by selectively relaxing some constraints?
4. Can the proposed separation between discrete and continuous data flows be applied to other neural network architectures beyond attention-based models? What challenges might arise in such adaptations?
soundness: 4
strengths: - Originality: The paper presents a novel method that enforces discrete state transitions in neural networks, which is a significant departure from traditional continuous representations. By integrating hard attention and separating discrete and continuous data flows, the authors address key challenges in neural algorithmic reasoning, particularly out-of-distribution generalization, and interpretability.
- Quality: The experimental results are strong, with DNAR achieving perfect test scores across multiple algorithmic tasks and graph sizes. The comparison with baseline models and state-of-the-art methods demonstrates the effectiveness of the proposed approach.
- Clarity: The paper is generally well-written and structured. The authors explain their methodology, including architectural choices and training procedures. The inclusion of diagrams and tables aids in understanding the proposed model and its performance.
- Significance: The work contributes to the field by showing that neural networks can be designed to mimic classical algorithms with perfect generalization and interpretability. This has implications for developing reliable and trustworthy AI systems that can be formally verified.
summary: This paper introduces a novel approach to neural algorithmic reasoning by enforcing neural networks to operate with discrete states and separating discrete and continuous data flows. The authors propose a model that integrates hard attention mechanisms, feature discretization, and a separation between discrete computations and continuous inputs (scalars). This design aims to align neural network computations closely with classical algorithms and thus improves out-of-distribution generalization and interpretability. The method is evaluated on several algorithmic tasks from the SALSA-CLRS benchmark, including BFS, DFS, Prim's algorithm, Dijkstra's algorithm, Maximum Independent Set (MIS), and Eccentricity calculations. The proposed Discrete Neural Algorithmic Reasoner (DNAR) achieves perfect test scores on these tasks, even on graphs significantly larger than those seen during training. The authors also discuss the limitations of their approach and potential directions for future work.
weaknesses: - Expressiveness Limitations: The enforced constraints, such as hard attention and discrete state transitions, limit the model's expressiveness. For instance, in a single message-passing step, the model cannot compute certain aggregate functions, like averaging over neighbors. This restricts the method's applicability to algorithms that fit within these constraints.
- Scope of Evaluation: The experimental evaluation focuses on specific algorithmic tasks where the proposed method aligns well. It remains unclear how the model would perform on more complex algorithms that require different computational primitives or continuous manipulations beyond simple increments.
- Training Without Hints: While the method achieves excellent results with hint supervision, training without hints is identified as challenging. This limitation reduces the method's applicability in scenarios where intermediate algorithmic steps (hints) are unavailable.
- Sensitivity to Hyperparameters: The paper mentions that certain hyperparameters, like the number of discrete states, significantly impact performance, especially when training without hints. However, there is limited discussion on how sensitive the model is to these hyperparameters and the implications for generalization.
- Presentation Details: While the paper is generally clear, some sections could benefit from additional explanations. For example, the scalar updater's mechanism could be elaborated further to better understand if someone is not in the NAR field.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. I'm very much not an expert in this area, so I did a bit of a literature search on related papers in the field. I found one such paper [1], used a text-based version of the CLRS-3o, called TextCLRS-text. Does it make sense to use that dataset here? Or to compare to the proposed TransNAR architecture?

2. As mentioned earlier, it would be useful to ablate the different algorithmic decisions made in this paper. In its current form, I can't ascertain which of the architectural decisions is responsible for which gains on these datasets.
soundness: 2
strengths: The approach seems to work much better than the baselines considered, and achieves perfect performance on the two datasets studied in this article.
summary: The authors make some modifications to Transformer training to make it more effective for algorithmic reasoning tasks. In particular, they constrain it to learn hard attention, and separate discrete and continuous "flows" to prevent information loss. Overall these modicfications result in a model than the GIN and PGN baselines they consider.
weaknesses: I found this paper to be difficult to read, but I'm not an expert in the area. It would be useful to see each of the design decisions in this paper ablated, with their corresponding effect on the datasets covered in this research.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. Results show perfect performance on SALSA-CLRS and CLRS-30 benchmarks. Could you elaborate on reported results and 0 variance between your seeds given that this is not the case in majority of comparable baselines or ablation experiments? Could you please report variance on results in table 1? 
2. Could you please explain the differences in computed discrete states used for SALSA-CLRS and CLRS-30 tasks as well as how the model performs on out-of-distribution data?
3. Could you provide a pseudocode behind Discretize_nodes and Discretize_edges functions shown on page 4?
4. Work mentions modifying hints from the benchmark, could you clarify if this was included for tested baselines as well as how the hints are modified?
5. Could you further elaborate on the addition of a virtual node in your model (section 4.3)?
6. Could you provide results of the hyperparameter search mentioned in section 6? 
7. Could you provide further detail on reverse-engineering (perhaps a toy example) mentioned in section 6?
soundness: 2
strengths: - Impressive performance achieved across selected tasks in SALSA-CLRS and CLRS-30 benchmarks
- Interpretability and simplicity of the proposed model which utilizes set of discrete states to capture continuous inputs using edge priorities
- Further perspectives on paving way for discrete neural algorithmic reasoners
summary: The paper proposes a use of a finite set of predefined discrete states and manipulation of continuous inputs to improve limitations of current paradigms in neural algorithmic reasoning such as redundant dependencies in algorithmic data. The method achieves perfect performance on chosen tasks in SALSA-CLRS and CLRS-30 benchmarks and introduces a starting perspective on use of discrete states to pave the way for future methods in discrete neural algorithmic reasoning and their interpretability.
weaknesses: - There is limited explanation of how the discrete states are computed, the paper extensively discusses related work in sections 1 and 2, perhaps it would be better if this space was used for providing further detail on proposed architecture
- Work is limited in providing mathematical / theoretical definitions or explanation of discrete state space and manipulation with continuous inputs which is the main novelty

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: Please see my comments in the weakness part.
soundness: 2
strengths: - The problem of learning neural algorithmic reasoning is important. I like the overall idea of learning discrete neural reasoners since the previous efforts of learning neural networks with continuous states and or continuous operators failed miserably.

&nbsp;

- The architecture design has some interesting components. For example, the separation of the data and the computational flow that manipulates the data is interesting. In particular, the scalars (continuous input) only affect the computation of attention weights and do not affect the node or edge states. Although designs with the same spirit have appeared before, e.g., in Neural Execution Engines, this part has its own merits.

&nbsp;


- The paper is easy to follow and well-written, except that a few technical details are sparse.
summary: This paper studies neural algorithmic reasoning and proposes a discrete transformer-based processor architecture. The authors show that the proposed model empirically achieves perfect size generalization on several algorithms picked from the CLRS benchmark.
weaknesses: - I am concerned about the significance of the contributions. If the claims of this paper are all correct, then we just obtain a recipe for learning a neural reasoner to perfectly mimic a known algorithm. However, to make the learning successful, we need to first run the algorithm to collect the full trace of the execution, i.e., the sequence of intermediate states generated by the algorithm (the so-called hints). In other words, we just perfectly fit the “correct algorithm” using a neural network. 

  Moreover, I do not see a theoretical guarantee of when this perfect fitting would happen. For example, would it happen on a specific class of algorithms or any algorithm? See my next comment on the claim about “the guarantee”.

  This is unsatisfying since the goal of neural algorithmic reasoning is to learn correct algorithms from data without knowing the algorithms. From this perspective, the really interesting and valuable part is the exploration under the no-hints setting. However, in Section 6, the authors did not provide any experimental results and just stated that their model never achieved perfect validation scores. In short, the authors should provide a thorough empirical study of their proposed model under the no-hints setting and compare it with other approaches on benchmarks like CLRS. 

&nbsp;

- In Section 5, the claim “we can guarantee that for any graph size, the model will mirror the desired algorithm, which is correct for any test size” is quite strong. In my opinion, such a strong claim needs rigorous theoretical proof. However, the authors only provide some vague arguments to support this claim. 
  In particular, can you elaborate on the following two questions?
  1) How do you confirm that the attention block indeed operates as “select_best” selector?
  I neither see any empirical investigation nor theoretical proof on this point.
  2) Even if the attention block operates as “select_best” selector, why does this condition lead to the above claim?
  You should at least demonstrate your detailed logic using an example algorithm like DFS.

&nbsp;

- I think the size of the discretized states would matter a lot to the expressivity of the proposed neural reasoner, i.e., what algorithms can be represented by the designed class of neural networks. However, I did not see the experimental study and discussion on its effect.

&nbsp;

- An interesting experiment is to check if there is some sort of phase transition in terms of problem size. Specifically, I would imagine the phenomenon of perfect fitting to the correct algorithm would disappear if we decreased the problem size in training. Then, what is the minimum problem size to ensure it fits perfectly for a particular algorithm?

&nbsp;

- The details of how to train task-dependent encoders and decoders are not provided, and how they affect the processor's training is not discussed at all. I would imagine the quality of the encoded embedding is quite important to the success of learning the processor, even with teacher forcing. 

&nbsp;

- In section 3.3, I get the high-level idea that the scalars (continuous input) only affect the computation of attention weights and do not affect the node or edge states. However, the description of the idea in the 2nd paragraph of Section 3.3 is not so clear that I could not figure out how exactly the computation is designed. It would be great to either write the equations and or illustrate the computational graph.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: 1. For the multi-task experiments, do the results in Table 2 correspond to these settings? It is unclear which algorithms were trained simultaneously and how this compares to the baselines’ configurations.

2. Without hints, DNAR struggles to generalize. Could you clarify the reverse-engineering issue described? Was any intermediate approach, such as noisy teacher forcing or teacher forcing decay, tested to examine the impact of hints on DNAR?

3. How were GIN and PGN used as baselines? Were they employed to treat this as an end-to-end node-level prediction task?
soundness: 2
strengths: 1. The writing is generally clear, though some areas could be further elaborated.
2. The paper is well-motivated, as both OOD generalization and interpretability in NAR are both important questions. The proposed architecture of separating the discrete and continuous data flows are novel and effective.
3. The perfect scores across algorithms are impressive, especially given the model’s capacity for size generalization on graphs 100 times larger, outperforming strong baselines.
4. The proposed architecture consists of three design components, each of which is validated through ablation studies to demonstrate its effectiveness.
summary: This paper addresses the challenges of generalization and interpretability in neural algorithmic reasoning (NAR) by introducing a novel architecture that guides the model to follow algorithmic execution as a series of finite, predefined states. The proposed architecture has three main components: feature discretization, hard attention, and separation of discrete and continuous data flows. The empirical results show that the model achieves perfect scores in both single-task and multi-task experiments across various algorithms. Additionally, the architecture enhances interpretability, allowing validation of the correct execution of desired algorithms.
weaknesses: 1. While the perfect scores achieved in the experiments are impressive, the paper could be strengthened by testing on a wider range of algorithms. Although both parallel (e.g., BFS) and sequential (e.g., Prim) algorithms are covered, most algorithms studied are graph-based, where previous NAR methods have already proven effective. The CLRS-30 dataset includes a broader variety of algorithms (e.g., sorting and search), where many NAR methods can struggle. Although SALSA-CLRS was chosen for its thorough OOD evaluation, testing the model on the CLRS-30 dataset with larger graph sizes would add valuable insights due to the dataset’s extensive algorithm coverage.

2. Another limitation is the lack of application to real-world datasets, as the authors note in the Future Work section. A significant advantage of NAR methods is their ability to operate on high-dimensional data by utilizing a pretrained model that mimic algorithms. This presents an additional OOD challenge with potential distribution shifts in real-world data. Evaluating the proposed architecture on real-world datasets would demonstrate its practical value; even a single real-world experiment, as seen in related works (e.g., Numeroso et al., 2023), could significantly strengthen the method's implications.

3. Although interpretability is a valuable strength of the proposed method, it is unclear how this model’s interpretability, achieved through analyzing state transitions and attention blocks, differs from other NAR approaches, which also allow interpretation of intermediate executions (e.g., using decoder outputs to indicate a node’s current predecessor in a sorting algorithm).

Meta-review: This paper presents an exciting approach to neural algorithmic reasoning which directly relies on constraining the model's conditioning to a finite, discrete set of states.

I believe the idea is important and definitely has promise. However, as it stands, even though the Authors' rebuttal resulted in some scores being raised, it appears the majority of the reviewers are leaning on the rejection side than the acceptance side.

From my read of the reviews and the paper, it appears that the Authors could definitely benefit from either (a) evaluating their method in more diverse experimental settings (beyond the specific algorithmic tasks considered here), or (b) having a more solid theoretical grounding of why the method is effective.

Upholding this opinion, I have decided to recommend this paper for rejection. I reiterate that the work is very interesting, but accepting it in its current form (without a more thorough investigation beyond the tightly controlled experimental environment studied so far) may do the work and its reach a disservice. The Reviewers did not oppose my decision.

Decision: Reject