SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Gleb Rodionov
rodionovgleb@yandex-team.ru
Neural algorithmic reasoning is an emerging area of machine learning focusing on building models that can imitate the execution of classic algorithms, such as sorting, shortest paths, etc. One of the main challenges is to learn algorithms that are able to generalize to out-of-distribution data, in particular with significantly larger input sizes. Recent work on this problem has demonstrated the advantages of learning algorithms step-by-step, giving models access to all intermediate steps of the original algorithm. In this work, we instead focus on learning neural algorithmic reasoning only from the input-output pairs without appealing to the intermediate supervision. We propose simple but effective architectural improvements and also build a self-supervised objective that can regularise intermediate computations of the model without access to the algorithm trajectory. We demonstrate that our approach is competitive to its trajectory-supervised counterpart on tasks from the CLRS Algorithmic Reasoning Benchmark and achieves new state-of-the-art results for several problems, including sorting, where we obtain significant improvements. Thus, learning without intermediate supervision is a promising direction for further research on neural reasoners.
Building neural networks that can reason has been a desired goal for decades of research. While the ability to reason can be defined and measured in various ways, most of them represent a challenge for current models, solving which can allow the community to use such models for complex real-world problems. Significant progress has been achieved in the reasoning abilities of language models (Touvron et al., 2023, Hoffmann et al., 2022), where reasoning can be described as the ability to answer questions or solve logical and mathematical problems in formal (Polu and Sutskever, 2020) or natural language (Lewkowycz et al., 2022, Jiang et al., 2023). However, the essential part of this progress is inseparable from the large scale of such models, which is needed to abstract over complex language structures (Wei et al., 2022).
Recent research considers classic algorithms as a universal tool to define reasoning on simple instances like sets of objects and relations between them, making algorithmic modeling a popular domain for testing neural networks and highlighting current reasoning limitations of existing architectures (Zaremba and Sutskever, 2014, Kaiser and Sutskever, 2015, Trask et al., 2018, Kool et al., 2019, Dwivedi et al., 2023, Yan et al., 2020, Chen et al., 2020b).
Neural algorithmic reasoning is an area of machine learning focusing on building models that can execute classic algorithms (Veličković and Blundell, 2021). The motivation is to combine the advantages of neural networks (processing raw and noisy input data) with the appealing properties of algorithms (theoretical guarantees and strong generalization). Assuming that we have a neural network capable of solving a classic algorithmic task, we can incorporate it into a more complex
pipeline and train end-to-end. For instance, if we have a neural solver aligned to the shortest path problem, it can be used as a building block for a routing system that accounts for complex and dynamically changing traffic conditions.
The main challenge of neural algorithmic reasoning is the requirement of out-of-distribution (OOD) generalization. Since algorithms have theoretical guarantees to work correctly on any input from the domain of applicability, the difference between train and test data distributions for neural algorithmic reasoners is potentially unlimited. Usually, the model is required to perform well on much larger input sizes compared to the instances used for training. This feature distinguishes neural algorithmic reasoning from standard deep learning tasks: neural networks are known to perform well on unseen data when it has the same distribution as the train data, but changes in the distribution may significantly drop the performance. This challenge is far from being solved yet; for example, as shown by Mahdavi et al. (2023), even achieving the almost perfect test score, neural reasoners can rely on implicit dependencies between train and test data instead of executing the desired algorithm, which was demonstrated by changing the test data generation procedure.
The area of neural algorithmic reasoning has been developing rapidly in recent years, and graph neural networks (GNNs) have proven their applicability in this area in various contexts. In particular, Velicković et al. (2020) and Xu et al. (2020) analyzed the ability of GNNs to imitate classic graph algorithms and explored useful inductive biases. Further, Dudzik and Velicković (2022) showed that GNNs align well with dynamic programming. Also, several different setups demonstrated that neural reasoners benefit from learning related algorithms simultaneously, explicitly (Georgiev and Lió, 2020, Numerouso et al., 2023) or implicitly (Xhonneux et al., 2021, Ibarz et al., 2022) reusing common subroutines. A significant step forward in developing neural algorithmic reasoners was made with the proposal of the CLRS benchmark (Velicković et al., 2022) that covers 30 diverse algorithmic problems. In this benchmark, graphs are used as a general tool to encode data since they straightforwardly adapt to varying input sizes.
An important ingredient of the CLRS benchmark is the presence of hints that represent intermediate steps of the algorithm's execution. With hint-based supervision, the model is forced to imitate the exact trajectory of a given algorithm. Such supervision is considered a fundamental ingredient for progress in neural algorithmic reasoning (Veličković et al., 2022).
In this paper, we argue that there are several disadvantages of using hints. For example, we demonstrate that the model trained without hint supervision has a greater tendency to do parallel processing, thus following the original sequential trajectory can be suboptimal. We also show that while the model trained with supervision on hints demonstrates the alignment to the original algorithm trajectory, it could not capture the exact dependency between the hint sequence and the final result of the execution.
Motivated by that, we investigate whether it is possible to train a competitive model using only input-output pairs. To improve OOD generalization, we propose several modifications to the standard approach. As a first step, we propose a simple architectural improvement of the no-hint regime that makes it better aligned with the hint-based version and, at the same time, significantly boosts the performance compared to the typically used architecture. More importantly, for several tasks, we add a contrastive regularization term that forces similar representations for the inputs having the same execution trajectories. This turns out to be a useful inductive bias for neural algorithmic reasoners.
The experiments show that our approach is competitive with the current state-of-the-art results relying on intermediate supervision. Moreover, for some of the problems, we achieve the best known performance: for instance, we get the F1 score  $98.7\%$  for the sorting, which significantly improves over the previously known winner with  $95.2\%$ .
We hope our work will encourage further investigation of neural algorithmic reasoners without intermediate supervision.
Our work follows the setup of the recently proposed CLRS Algorithmic Reasoning Benchmark (CLRS) (Velicković et al., 2022). This benchmark represents a set of 30 diverse algorithmic problems. CLRS proposes graphs as a general way to encode data since any particular algorithm can be expressed as manipulations over a set of objects and relations between them. For example, for the Breadth-first search (BFS), the input represents the original graph with a masked starting node, and the output
Figure 1: Hints usage diagram for the Bubble sort. Green lines represent the encoders, and red — the decoders. The dotted lines illustrate the supervision. Hints at each step represent the current order of the elements and two indicated elements for the last comparison.
of the algorithm (traversal order of the graph) can be encoded via predicting the predecessor node for any node. For non-graph tasks, CLRS proposes an approach to derive the graph structure. For example, for sorting algorithms, each element of the array can be treated as a separate node, endowed with position scalar input for indexing, and the output is encoded as a pointer from each node to its predecessor in the sorted order.
An important component of the CLRS benchmark is hints, which represent the decomposition of the algorithm's execution into intermediate steps. Each hint contains all the necessary information to define the next execution step. For example, for Bubble sort, a hint encodes the current order of the elements and two indexes of the elements for the next comparison. We note that each algorithm defines a unique way of solving the problem, which can be learned separately.
Each transition between the consecutive trajectory states defines a new task, which can be learned separately or simultaneously. Recent work has demonstrated the utility of the encode-process-decode paradigm for learning the whole trajectory of the algorithm at once (Veličković et al., 2022, Ibarz et al., 2022) (Figure 1). At each time step  $t$ , all input data are encoded with simple encoders to a high-dimensional (we use  $d = 128$ ) latent space, producing vectors  $(x_{i}^{t}) \in \mathbb{R}^{n \times d}$  representing nodes,  $(e_{ij}^{t}) \in \mathbb{R}^{n \times n \times d}$  edges, and  $(g^{t}) \in \mathbb{R}^{d}$  graph features, where  $n$  is the task size (nodes count).
Then, a processor, usually a graph neural network (Veličković et al., 2022, Ibarz et al., 2022) or a transformer (Diao and Loynd, 2023, Mahdavi et al., 2023), operates in this latent space, executing a transition to the next step. The processor computes the hidden states  $h_i^t$  for the nodes (and optionally edges) from the combination of the encoded inputs and hidden states  $h_i^{t-1}$  from the previous step, where  $h_i^0 = 0$ . After this transition, the latent vectors are decoded to predict the hints for the current step. Then, the current step hints prediction can be discarded or fed back into the processor model for the next step, similarly to the encoded input, allowing the model to follow the predefined dynamic of the hints sequence.
At the last step of the trajectory, decoders project the output prediction to the problem. Similarly to the encoders, the decoders rely mainly on a linear layer for each hint and output. The decoders are also endowed with a mechanism for computing the pairwise node similarities when appropriate, which allows one to unify a method for predicting node-to-node pointers for graphs with different sizes. We refer to Velicković et al. (2022) for a more detailed explanation. It was also recently shown that the processor can be shared across tasks and thus reuse algorithmic subroutines (Ibarz et al., 2022).
Another recent work (Bevilacqua et al., 2023) shows a significant improvement in size generalization via a novel way of using hints called Hint-ReLIC. The method comes from the observation that for some algorithms, there exist multiple different inputs on which certain steps of the algorithm will perform identically. Designing causal structure on the algorithm's trajectory, the authors have demonstrated that the execution of each step depends only on a small part of the current snapshot, meaning that the invariance under changes of irrelevant parts can serve as a regularization for the model. For example, the Bubble sort at each step compares only two elements of the array, so any
changes of other elements or adding new elements to the array should not affect the execution of the given step. Such regularization forces better alignment of the model to the relevant computations, improving size generalization. Invariance is forced via a self-supervised objective which is a modified version of the ReLIC algorithm (Mitrovic et al., 2021).
3 Neural algorithmic reasoning without hints
Motivation There are several disadvantages of using hints in neural algorithmic reasoning. First, hints need to be carefully designed separately for each task and each particular algorithm. This may require a significant amount of human effort: e.g., the hint at the next step should be predictable by the hint from the previous step using a single round of the message passing (Veličković et al., 2022). Moreover, as this example shows, different model architectures may benefit from differently designed hints: for a processor based on a graph transformer, the optimal hints can be different. Mahdavi et al. (2023) demonstrated the comparison of the performances for different GNN architectures depending on hint usage, showing that for some algorithms the question of using hints is faced by each architecture separately.
As discussed above, hints force the model to follow a particular trajectory. While this may help the interpretability and generalization of the trained model, following the algorithm's steps can be suboptimal for a neural network architecture. For instance, it may force the model to learn sequential steps, while parallel processing can be preferable. Moreover, a model trained with direct supervision on hints could not necessarily capture the dependencies between hints and the output of the algorithm.
To demonstrate this disadvantage with a practical example, we compared the dynamics of the predictions for the model trained on the Insertion sort problem depending on hint usage, see Figure 2. For each model, we evaluated the intermediate predictions after each processor step and compared them with the final predictions of the model. We note that for hint-based models we have two decoders for predictions: the hint decoder and the output decoder, both of which can be used for this purpose.
We see that the model trained with supervision on the trajectory for Insertion sort struggles to learn the relation between the intermediate predictions of the pointers and the output predictions (blue line). In other words, the model potentially can consider the hint sequence (and even each transition between hints) as a separate task, not related to the output of the algorithm directly.
Figure 2: The ratio of pointers predicted after a given number of steps which are equal to the final model output. Insertion sort task, test data (64 nodes).
We also see that the execution trajectory for the models with hints inherits some sequential order of intermediate updates (orange line). The model without hints has a much stronger tendency to do parallel processing — almost all output predictions are obtained on the first steps of the execution.
Finally, we note that if a model is not restricted to a particular trajectory, it may potentially learn new ways to solve the problem, and thus by interpreting its computations, we may get new insights into the algorithms' design.
Simple improvements of the no-hint mode Let us describe several changes to the no-hint mode that make the comparison with hint-based methods more straightforward and, at the same time, significantly boost the performance compared to the typically used no-hint regime.
First, we propose an architectural modification of the no-hint mode that makes its computational graph more similar to the standard regime with hints. As described in Section 2, in the default implementation of the encode-process-decode paradigm of the CLRS benchmark, all hint predictions from the current step are fed back to the processor model (as illustrated in Figure 1). Thus, in addition
(a) Standard no-hint mode
Figure 3: One step of the algorithm's execution depending on the different hint usage modes for the Bubble sort algorithm.
(b) Hints supervision
(c) Proposed no-hint mode
to the intermediate supervision, hint predictions are also used in the next computational step. So, the transition of the hidden states from step  $t - 1$  to step  $t$  uses not only the encoded inputs and hidden states  $h_i^{t-1}$  but also the current hint predictions, which are derived from the hidden state via additional computations.
However, for the no-hint version, as there are no intermediate steps to supervise, hints are not predicted and not fed back to the model. While this approach seems natural, it leads to significant computation graph differences compared to the model with hints. We note that the computation of hint predictions and encoding them to the processor in the encoded-decoded mode can be expressed as an additional message-passing step over the nodes, edges, and graph features (which can differ from the message-passing procedure of the processor, if any). Thus, it can be reasonable to preserve this part even in the no-hint regime. So, we propose to run the no-hint model in the encoded-decoded $^1$  mode but do not supervise hint predictions, allowing the model to construct useful representations and use them in the next steps. This modification makes the computation graphs similar and also noticeably improves the performance of the no-hint regime. We refer to Figure 3 for the illustration of the proposed changes.
Second, we note that all algorithms (and, therefore, hint sequences) assume sequential data representation. However, for some algorithms, only the intermediate steps depend on the position encoding of nodes, while the outputs are independent of the position scalars. For example, the output for the sorting algorithms is represented as pointers to the predecessor in the sorted order (and the node with the minimal value points to itself). Positions are only needed to tie-break identical elements, but due to the data generation in the CLRS benchmark, such collisions do not occur. To avoid unnecessary dependency of models on the position encodings, we do not use them if the output is independent of positions. For some algorithms, however, the positions are needed: for instance, the Depth-first search traversal depends on vertex indexing. Thus, we keep the position scalars only for the DFS, SCC, and Topological sort problems.
4 Self-supervised regularization
4.1 Reducing the search space
In this section, we propose a self-supervised objective that can regularise intermediate computations of the model by forcing similar representations for the inputs having the same execution trajectories. The proposed technique can be directly applied to the following algorithms of the CLRS benchmark: Minimum spanning tree algorithms, Sorting, Binary search, and Minimum.
As we mentioned in the previous section, any predefined trajectory can be suboptimal for a given architecture. To design our alternative approach, we make an observation that helps us to combine the advantages of learning algorithms only from input-output pairs with inductive bias improving the size generalization capabilities of the models.
<sup>1</sup>We note that for models without hints there is no abstract meaning for encoding/decoding of the intermediate states, we use such wording only for consistency with hint-based models.
Figure 4: Self-supervised regularization for neural algorithmic reasoning.
In the absence of hints, the model does not have any knowledge about the algorithm's trajectory and its intermediate steps and potentially can align to a new way of solving the problem. Hence, obtaining any information about underlying computations of the model, for example, detecting the part of the data irrelevant to the current step, is non-trivial. Moreover, given the parallel nature of modern model architectures, all nodes can potentially be used at all steps. Thus, we propose constructing such a pair of different inputs that every step of a given algorithm is the same for both inputs.
Although constructing such pairs for an arbitrary algorithm is not always trivial, the underlying combinatorial structure of some tasks allows us to describe the properties of the desired algorithm. Such properties can reduce the search space of the model while remaining in the wide range of possible algorithms, which can solve the given problem.
For example, any comparison sort has the same trajectory for any pair of arrays with the same relative order of its elements. We clarify that we consider the sequence of the intermediate relative orders of the elements as a trajectory, not the array values themselves. A less trivial example is the minimum spanning tree problem. A wide range of algorithms (Prim (1957), Kruskal (1956), etc.) have the same execution trajectory on graphs with fixed topology (graph structure) and different weights on edges if the relative order of the corresponding edge weight is the same. As the set of all subsets of edges that do not contain a cycle forms a matroid (Harary, 1969), any greedy construction of the tree would lead to the optimal result. This observation allows us to use only the relative order of the edge weights, not the weights themselves. We note that the CLRS benchmark considers the set of edges in the tree (a mask on edges) as the result of the algorithm, not the total weight of the constructed tree.
One important remark is that for two different inputs, the exact matching of all the intermediate steps is generally a more strict condition than the matching of the result. For example, consider Kruskal's algorithm for the minimum spanning tree problem. The first stage of this algorithm is sorting the edges by their weight. The edges are then greedily added to the disjoint set of independent edges until the spanning tree is constructed. Thus, an edge with a weight that is larger than any weight from the minimum spanning tree does not affect the result but can cause the different states of the algorithm after the sorting stage.
Now we describe our idea more formally. Given a problem  $P: \mathcal{X} \to \mathcal{Y}$ , we define a set  $A$  of the algorithms solving this problem, and a decomposition of the set of inputs  $\mathcal{X} = \bigcup \mathcal{X}_k$  to the disjoint equivalent classes, such that for every two elements  $X_i, X_j \in \mathcal{X}_k$  from the same class  $\mathcal{X}_k$ , all algorithms from  $A$  produce the same execution trajectory. We note that the properties of the input data that we mentioned above produce such a decomposition of the input spaces for the specified tasks. Now we can use this to reduce the search space for our model to the space of functions acting the same way on any pair of inputs from the same equivalence class, being invariant under input changes that do not affect the algorithm's trajectory. Our idea is illustrated on Figure 4.
In contrast to Hint-ReLIC, we do not decompose each snapshot of the execution to the relevant and non-relevant nodes, as the dynamic of such decomposition is defined by the hints sequence and thus may inherit its disadvantages. Instead, we highlight the property of the whole input of the algorithm.
This idea allows us to give the model additional useful signals about the data without any explicit or implicit enforcing to any predefined dynamic, which can be suboptimal for a particular model architecture.
4.2 Adding contrastive term
Now, let us describe how classic contrastive learning techniques can be applied for this purpose. As mentioned above, for a given task, our goal is to learn representations of intermediate steps  $t \in [1 \dots T]$ , which are similar for any pair of inputs from the same equivalence class. One natural way to do this is to build a contrastive objective aimed at solving the instance discrimination task (Wu et al., 2018, Chen et al., 2020a, Mitrovic et al., 2021).
We note that there are different ways to define an instance. For example, one may consider the whole execution of the algorithm as a single classification task, where each class represents a trajectory choice. In this case, all the intermediate states of all the nodes can be used as a single instance. However, as the algorithm's execution can be described as a step-by-step decision-making process, a more fine-grained approach is suitable, defining each node representation at each step as a separate instance. We focus on the latter option.