SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Let us denote by  $\mathcal{X}_X$  the class of equivalence for  $X\in \mathcal{X}$ , which is a set of valid augmentations for  $X$ , under which the similarity of representations for the corresponding nodes at each step can be directly enforced. For this purpose, we use the following contrastive regularization term:
$$
L _ {t} = - \sum_ {X \in \mathcal {X}} \sum_ {i = 1} ^ {n} \sum_ {X _ {a} \in \mathcal {X} _ {X}} \log \frac {\exp (\phi (f _ {t} (X , i) , f _ {t} (X _ {a} , i)))}{\sum_ {j = 1} ^ {n} \exp (\phi (f _ {t} (X , i) , f _ {t} (X _ {a} , j))))},
$$
where  $f_{t}(X,i)$  denotes the representation of  $i$ -th node  $(i\in [1\dots n])$ , obtained at step  $t$ , and  $\phi (x,y)$  denotes a comparison function for the obtained representation, which we implement as  $\phi (x,y) = \langle g(x),g(y)\rangle$ , where  $g$  is a fully-connected neural network. We compute this term for each step  $t\in [1\dots T]$  and add it to the output loss.
We note that this term is similar to the contrastive part of the Hint-ReLIC objective with some differences. First, in Hint-ReLIC, certain hints are contrasted (such as intermediate prediction of the predecessor in the sorted order for sort algorithms, namely  $pred_h$ ) to be invariant under changes that affect only those parts of the data, which are not included to the current step of execution. Such an approach allows for highlighting the importance of different nodes at different execution steps. Our method does not specify any additional details about intermediate computations, so the whole latent representations of each node state are forced to be invariant under valid augmentations. Another difference is in the augmentations choice. While Hint-ReLIC is compatible with a wide range of interventions, such as connecting the underlying graph to any randomly generated subgraph, our augmentations for  $X$  are defined by the equivalence class  $\mathcal{X}_X$  in which by design all the inputs have the same size. To summarize, we propose a different way of applying the standard contrastive learning technique to the algorithmic reasoning setup. While Hint-ReLIC demonstrates an advanced way to align the model to the particular trajectory of the execution, our method can be applied as additional regularization to the model, which does not favor any particular trajectory from the wide range of possible algorithms.
In this section, we evaluate the performance of the proposed modifications to the no-hint mode and test the impact of contrastive regularisation on OOD generalization. To deepen our analysis, we also conducted additional experiments, evaluating the impact of the step count and demonstrating the influence of the contrastive term on similarity of trajectories for algorithmically equivalent inputs.
Baselines We compare our approach with several forms of hint usage. First, we consider the standard supervision of the hints, denoted as Baseline (Ibarz et al., 2022). Then, No hint is the model that completely ignores the hint sequence. Finally, Hint-ReLIC utilizes an advanced way of using the
Table 1: Training details and model hyperparameters
Train sizes	4, 7, 11, 13, 16
Validation size	16
Test size	64
Optimiser	Adam
Learning rate	0.001
Train steps count	10000
Evaluate each (steps)	50
Early-stopping patience (steps)	500
Batch size	32
Processor	Triplet-GMPNN
Hidden state size	128
Number of message passing steps per processor step	1
hints (Bevilacqua et al., 2023). All models use the Triplet-GMPNN architecture proposed by Ibarz et al. (2022).
Model and training details We train our model using the Adam optimizer with an initial learning rate  $\eta = 0.001$ , with batch size 32, making 10,000 optimization steps with early stopping based on the validation performance. Our models are trained on a single A100 GPU, requiring less than 1 hour to train. We report the results averaged across five different seeds, together with the standard error. We report all hyperparameters in Table 1.
We note that our proposed modification of the no-hint mode allows for the use of any type of representation replacing the hint prediction. For example, each sorting algorithm in the CLRS benchmark has a unique description of the intermediate computations inspired by the algorithm itself. Each step of the Bubble sort algorithm is described by the current elements order and two indexes for the current comparison. At the same time, hints for the Heapsort have a more complex structure that defines the current state of the heap and other details (Veličković et al., 2022). Such differences in the specifications of the algorithms produce differences in computational graphs for models trained for Bubble sort and Heapsort. While the proposed modification of the no-hint mode leaves freedom for such architecture choices, we consider a search of the optimal specifications (description of the representation of the input, intermediate steps, and output) for each algorithm beyond the scope of our work. Moreover, we do not use the original specifications for the hints from the CLRS benchmark, as their design (combination of types of hints) can potentially align the model to the original algorithm even without the supervision of the hints. Thus, we only construct a single type of intermediate representations to demonstrate the advantage of the proposed modification in the simplest way. We use a latent representation of type mask located at edges and use them for all the algorithms.
We also need to set the trajectory length (the number of message-passing steps) for our model. For the regime with hints, the number of steps can be obtained from the original trajectory. However, for the no-hint version, there are several options: take the number of steps from some algorithm, use a predefined steps count, or add a separate model that predicts if the algorithm execution is ended. We use the steps count linearly depending on the input size, with factors 1 and 5 for array and graph problems, respectively.
Data augmentations Recall that our proposed regularisation technique is aimed at forcing the invariance of the execution under the choice of the elements from the same equivalence class. To achieve this, we augment the inputs by sampling from the same data distribution, preserving the relative order of the elements for the array-based problems and the relative order of the corresponding edges for graph problems. We generate one augmentation for each data point.
Dataset We test our modifications of the no-hint mode on the same subset of the algorithms from the CLRS benchmark as in Bevilacqua et al. (2023). We additionally evaluate our contrastive objective on the Sorting, Minimum, Binary search, and Minimum spanning tree problems.
Table 2: Comparing performance for different hint usage modes. The table shows the mean and standard error for the test micro-F1 score across different seeds.
Algorithm	No hint	Baseline	Hint-ReLIC	No hint (ours)	Add contr. (ours)
Artic. points	81.97 ± 5.08	88.93 ± 1.92	98.45 ± 0.60	94.55 ± 2.47	—
Bridges	95.62 ± 1.03	93.75 ± 2.73	99.32 ± 0.09	98.91 ± 0.47	—
DFS	33.94 ± 2.57	39.71 ± 1.34	—	34.52 ± 1.17	—
SCC	57.63 ± 0.68	38.53 ± 0.45	76.79 ± 3.04	66.38 ± 7.12	—
Topological sort	84.29 ± 1.16	87.27 ± 2.67	96.59 ± 0.20	94.58 ± 0.60	—
Bellman-Ford	93.26 ± 0.04	96.67 ± 0.81	95.54 ± 1.06	95.25 ± 0.58	—
BFS	99.89 ± 0.03	99.64 ± 0.05	99.00 ± 0.21	99.95 ± 0.04	—
DAG SP	97.62 ± 0.62	88.12 ± 5.70	98.17 ± 0.26	98.18 ± 0.99	—
Dijkstra	95.01 ± 1.14	93.41 ± 1.08	97.74 ± 0.50	96.97 ± 0.32	—
Floyd-Warshall	40.80 ± 2.90	46.51 ± 1.30	72.23 ± 4.84	42.16 ± 2.14	—
MST-Kruskal	92.28 ± 0.82	91.18 ± 1.05	96.01 ± 0.45	92.89 ± 0.92	93.15 ± 1.01
MST-Prim	85.33 ± 1.21	87.64 ± 1.79	87.97 ± 2.94	85.70 ± 1.44	85.20 ± 1.77
Insertion sort	77.29 ± 7.42	75.28 ± 5.62	92.70 ± 1.29	90.67 ± 8.22	98.74 ± 1.12
Bubble sort	81.32 ± 6.50	79.87 ± 6.85	92.94 ± 1.23	90.67 ± 8.22	98.74 ± 1.12
Quicksort	71.60 ± 2.22	70.53 ± 11.59	93.30 ± 1.96	90.67 ± 8.22	98.74 ± 1.12
Heapsort	68.50 ± 2.81	32.12 ± 5.20	95.16 ± 1.27	90.67 ± 8.22	98.74 ± 1.12
Binary Search	93.21 ± 1.10	74.60 ± 3.61	89.68 ± 2.13	85.29 ± 4.52	87.12 ± 2.23
Minimum	99.24 ± 0.21	97.78 ± 0.63	99.37 ± 0.20	99.98 ± 0.02	99.99 ± 0.01
Table 3: Comparing test performance of the no-hint (ours) model on the sorting problem for different processor step counts.
Steps count	Test score
Linear (same as for Insertion sort from CLRS)	90.67 ± 8.22
Quadratic (same as for Bubble sort from CLRS)	90.97 ± 8.03
Linearthmic (same as for Quick sort from CLRS)	91.05 ± 8.17
Linearthmic (same as for Heap sort from CLRS)	89.66 ± 7.40
Table 2 compares the performances of several hint usage modes, including those proposed in Sections 3 and 4. We first note that our modification of the no-hint mode significantly improves the performance compared to the standard No hint version, achieving the best-known results on the BFS, DAG shortest paths, and Minimum problems. For instance, for sorting, the performance increases from  $68.5\% - 81.3\%$  for No hint to  $90.7\%$  for our version. Here we note that such a spread in performance between different sorting algorithms in the standard no-hint mode is caused by different step counts: No hint does not rely on hints, but uses a particular algorithm when setting the number of message-passing steps. In contrast, our modification uses a pre-defined step count depending on the input size. Also, compared to baseline supervision on hints, our no-hint model yields better performance on 14 out of 18 tasks. To conclude, our simple modifications of the no-hint mode already demonstrate competitive performance on most of the problems.
The last column of Table 2 demonstrates further improvements obtained due to the proposed self-supervised regularization. In particular, with this regularization, we achieve the best results for sorting  $(98.7\%)$ , significantly outperforming the existing models.
We also evaluate the effect of the step count for the sorting problem, see Table 3. For our experiments, we used the same step count as the Baseline model for Insertion sort. We demonstrate that mimicking the step counts from other baseline models still shows the advantages of the proposed no-hint changes.
The Hint-ReLIC results for DFS are not reported since Bevilacqua et al. (2023) achieve the perfect score with a simpler technique.
Finally, we show that for the sorting problem, the proposed contrastive term forces the model to have similar representations for inputs from the same equivalence class, as intended. For this purpose, we took different pairs of equivalent inputs and for each step of the execution we calculated the following ratio: if we take one node from one input and take the closest node from the augmented input, what is the probability that the closest node will represent the element with the same relative order. Figure 5 shows that while the model without the contrastive term (green line) is able to mainly rely on the relative order, our contrastive term is forcing much more confidence (and speed) in achieving that property.
Figure 5: Influence of the contrastive term on similarity of the representations of equivalent inputs depending on weight factor  $w$  of the contrastive term. Insertion sort problem, test data (64 nodes).
In this paper, we propose improvements of the no-hint regime for neural algorithmic reasoning and a method to regularise intermediate computations without enforcing the model to any predefined algorithm trajectories. While this approach seems natural and promising, some limitations remain.
Our contrastive learning approach applies to the limited set of problems. While for many algorithms it is possible to construct the inputs that produce the same trajectory, sometimes it can be non-trivial. For example, for the shortest path problems, linear transformation of each weight does not affect the resulting paths and all intermediate computations (such as the order of inserting nodes to the heap for the Dijkstra's algorithm (Dijkstra, 1959)). However, such augmentations are not informative due to the layer normalization, and constructing more diverse examples represents a separate task.
While the proposed self-supervised regularization aims to reduce the search space for the model, it does not give a model a strong inductive bias towards underlying computations. We believe that more advanced techniques can be useful for algorithms with long sequential reasoning, such as the minimum spanning tree. Developing stronger inductive biases is a promising direction for future research.
For our self-supervised regularization, we augment the inputs by sampling from the same data distribution. While this gives a signal to the model that it should account only for the relative order, this signal can be even more significant if we use different distributions. Investigating the effect of diverse data distributions on the performance of neural algorithmic reasoners is a good subject of a separate study.
In this work, we propose several improvements to learning algorithms without intermediate supervision. We also demonstrate the possibility of regularising all computational steps of the model without enforcing it to any predefined algorithm trajectory. The obtained models outperform previous no-hint models and show competitive performance with models that have access to the entire algorithm trajectory. We also achieve new state-of-the-art results on several tasks. In particular, a considerable improvement is obtained for the sorting algorithms, where we get the score of  $98.4\%$ .
We consider neural algorithmic reasoning without intermediate supervision a promising direction for further advances in neural algorithmic reasoners.
R. C. Prim. Shortest connection networks and some generalizations. Bell System Technical Journal, 36:1389-1401, 1957.
H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
A. Trask, F. Hill, S. E. Reed, J. Rae, C. Dyer, and P. Blunsom. Neural arithmetic logic units. Advances in Neural Information Processing Systems, 31, 2018.
P. Velicković and C. Blundell. Neural algorithmic reasoning. *Patterns*, 2(7):100273, 2021.
P. Velicković, R. Ying, M. Padovano, R. Hadsell, and C. Blundell. Neural execution of graph algorithms. In International Conference on Learning Representations, 2020.
P. Velicković, A. P. Badia, D. Budden, R. Pascanu, A. Banino, M. Dashevskiy, R. Hadsell, and C. Blundell. The CLRS algorithmic reasoning benchmark. In International Conference on Machine Learning, pages 22084-22102. PMLR, 2022.
J. Wei, Y. Tay, R. Bommasani, C. Raffel, B. Zoph, S. Borgeaud, D. Yogatama, M. Bosma, D. Zhou, D. Metzler, E. H. Chi, T. Hashimoto, O. Vinyals, P. Liang, J. Dean, and W. Fedus. Emergent abilities of large language models. Transactions on Machine Learning Research, 2022. ISSN 2835-8856.
Z. Wu, Y. Xiong, S. X. Yu, and D. Lin. Unsupervised feature learning via non-parametric instance discrimination. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 3733-3742, 2018.
L.-P. Xhonneux, A.-I. Deac, P. Velicković, and J. Tang. How to transfer algorithmic reasoning knowledge to learn new algorithms? Advances in Neural Information Processing Systems, 34: 19500-19512, 2021.
K. Xu, J. Li, M. Zhang, S. S. Du, K.-i. Kawarabayashi, and S. Jegelka. What can neural networks reason about? In International Conference on Learning Representations, 2020.
Y. Yan, K. Swersky, D. Koutra, P. Ranganathan, and M. Hashemi. Neural execution engines: Learning to execute subroutines. Advances in Neural Information Processing Systems, 33:17298-17308, 2020.
W. Zaremba and I. Sutskever. Learning to execute. arXiv preprint arXiv:1410.4615, 2014.