On the Markov Property of Neural Algorithmic Reasoning: Analyses and Methods
Authors: Montgomery Bohde
Meng Liu
Alexandra Saxton
Shuiwang Ji
Venue: ICLR 2024 spotlight (2024)
Keywords: Neural Algorithmic Reasoning
Primary area: learning on graphs and other geometries & topologies
TL;DR: This work reveals the importance of aligning model design with the Markov nature in neural algorithmic reasoning tasks.

Abstract: Neural algorithmic reasoning is an emerging research direction that endows neural networks with the ability to mimic algorithmic executions step-by-step. A common paradigm in existing designs involves the use of historical embeddings in predicting the results of future execution steps. Our observation in this work is that such historical dependence intrinsically contradicts the Markov nature of algorithmic reasoning tasks. Based on this motivation, we present our ForgetNet, which does not use historical embeddings and thus is consistent with the Markov nature of the tasks. To address challenges in training ForgetNet at early stages, we further introduce G-ForgetNet, which uses a gating mechanism to allow for the selective integration of historical embeddings. Such an enhanced capability provides valuable computational pathways during the model's early training phase. Our extensive experiments, based on the CLRS-30 algorithmic reasoning benchmark, demonstrate that both ForgetNet and G-ForgetNet achieve better generalization capability than existing methods. Furthermore, we investigate the behavior of the gating mechanism, highlighting its degree of alignment with our intuitions and its effectiveness for robust performance. Our code is publicly available at https://github.com/divelab/ForgetNet.

Published as a conference paper at ICLR 2024
ON THE MARKOV PROPERTY OF NEURAL ALGORITH- MIC REASONING: ANALYSES AND METHODS
Montgomery Bohde, Meng Liu," Alexandra Saxton, Shuiwang Ji
Department of Computer Science & Engineering
Texas A&M University
College Station, TX 77843, USA {mbohde, mengliu, allie. saxton, sjij @t ammuu. edu
ABSTRACT
Neural algorithmic reasoning is an emerging research direction that endows neural networks with the ability mimic algorithmic executions step-by-step. A common paradigm in existing designs involves the use historical embeddings in predicting the results of future execution steps. Our observation in this work is that such historical dependence intrinsically contradicts the Markov nature of algorithmic reasoning tasks. Based on this motivation, we present our ForgetNet, which does not use historical embeddings and thus is consistent with the Markov nature of the tasks. To address challenges in training ForgetNet at early stages, we further introduce GG-ForgetNet, which uses gating mechanism to allow for the selec- tive integration of historical embeddings. Such an enhanced capability provides valuable computational pathways during the model's early training phase. Our extensive experiments, based on the CLRS-30 algorithmic reasoning benchmark, demonstrate that both ForgetNet and G-ForgetNet achieve better generalization capability than existing methods. Furthermore, we investigate the behavior the gating mechanism, highlighting its degree of alignment with our intuitions and its effectiveness for robust performance. Our code is publicly available at httpss//gghuub.com//vvcllbb///
INTRODUCTION
Neural algorithmic reasoning stands at the intersection of neural networks and classical algorithm research. It involves training neural networks reason like classical algorithms, typically through learning to execute step-by-step algorithmic operations (Veličković & Blundell [021 Veličkovid al. 2022a). Since classical algorithms inherently possess the power to generalize across inputs of varying sizes and act as "building blocks" for complicated reasoning pathways, learning to mimic algorithmic execution can confirm and amplify the generalization and reasoning abilities of neural
network models (Xu et al. 2020 Deac et al. 2021 Numeroso et al., 2023 Veličković et al.) (022b).
Existing works based on the CLRS-30 benchmark (Veličković et al. 2022a) have demonstrated the
effectiveness of mimicking algorithmic operations in high-dimensional latent space (Veličković et al.
20205 Georgiev Lio 2020; Veličković et al. 2020aa] 2022a; Ibarz et al. 2022 Diao & Loynd framework to learn the step-by-step execution of algorithms. At each step, the current algorithm state is first embedded in high-dimensional latent space via the encoder. The embedding is then given to the processor to perform one step of computation in the latent space. The processed embeddings are then decoded predict the updated algorithm state, namely hints. Within this paradigm, common practice is to use historical embeddings in the current execution step. Our insight in this work is that such historical dependence contradicts the intrinsic nature of classical algorithms.
2023). As detailed in Section 3.1 they typically employ
Our work is motivated by the Markov property of algorithmic reasoning tasks; that is, the present state is sufficient to fully determine the execution output of the current step. This observation led us to investigate if the use of historical embeddings in the existing paradigm is indeed useful as it does
Equual contribution

---

Published as a conference paper at ICLR 2024
not align with the underlying Markov property. Such misalignment introduces noise, thus hindering the model's generalization ability, especially in distti-iittiitttoo scenarios. To be consistent with the Markov property, we present ForgetNet, which removes the dependency on historical embeddings and explicitly embraces the Markov nature of algorithmic reasoning tasks. Such a modification, while simple, fundamentally realigns the computational graph of the neural model with the inherent structure of algorithmic processes, We observe that, although ForgetNet shows improvements across wide range of tasks, training such models may be challenging due inaccurate intermediate state predictions, especially at the early stages of training. To improve training, we further enhance our design with G-ForgetNet, in which a regularized gating mechanism introduced in order to align with the Markov property during testing while still allowing for beneficial computational pathways during training.
Our extensive experimental evaluations on the widely used CLRS-30 algorithmic reasoning bench- mark demonstrate that both ForgetNet and G-ForgetNet outperform established baselines. In partieu- lar, G-ForgetNet achieves robust and promising performance in many different tasks, showing the benefit the proposed gating mechanism. Further in-depth analyses of the training dynamics and gate behavior shed light on understanding of the advantages of the proposed approaches. Overall, the findings in this work demonstrate the importance of aligning model design with the underlying Markov nature achieve better generalization performance in neural algorithmic reasoning tasks.
2 RELATED WORK
Equipping neural networks with algorithmic reasoning abilities has gained increasing attention in recent research. Early attempts (Zaremba Sutskever] 2014: Graves et al. 2014 Kaiser & Sutskever 2015; Graves et al., 2016:ooulin Doulin & | networks with augmented memory mechanisms to mimic algorithms, showing that neural models could learn algorithmic patterns from data. With the use of graph-based representations, graph neural networks (GNNs) (Gilmer et al., 2017 Battagli et al. 2201)) can be applied naturally to algorithmic reasoning tasks (Veličković al. [200bb] Georgiev & Lio][202) Veličković al. 2020a Xu et al. 2020; Yan et al.) 2020 Dudzik & Veličkovié, 2022] Dwivedi et al. 2023). Intuitively, the message passing schema in various GNNs can naturally model the propagation and iterative nature of many classical algorithms. Recently, Veličković & Blundell (2021) outlines a blueprint for neural algorithmic reasoning, proposing a general propeddrrproesss--eeoodde framework. The framework, trained in algorithmic tasks, produces processors with potential applicability in real-world applications. Such generalization and transferable reasoning capabilities have been showcased in a few prior studies (Deac et al. | Numeroso et al.) 2023] Veličković et al. 2022b Beurer- Kellner jet al. 2022). In addition, Xhonneux et al. (2021) [bbrz et al. (2022) have explored the generalization ability of the processor across multiple algorithms.
in this direction typically use recurrent neural
[022]/Diivedi
To provide a comprehensive testbed for algorithm reasoning tasks, Veličković et al. (2022a) presents the CLRSS-3 benchmark, which covers 30 classical algorithms that span sorting, searching, dynamic programming, geometry, graphs, and strings (Cormen et al.) 2022). The CLRS-30 benchmark is known for its cottof-distribution (OOD) testing setup (Mahdavi al. 2022). where the input size of testing samples is much larger than those during training. Such a setup provides a rigorous test for the generalization capability of models, serving as standard testbed for algorithmic reasoning tasks. In the benchmark study, multiple representative neural models, including Deep Sets (Zaheer et al.) 2017). GAT (Veličković al. 2018). MPNN (Gilmer et al.) 2017), PGN (Veličković et al.) (002a)), and Memnet (Sukhbaatar et al. 2015). have been evaluated as the processor network within the mecoder-processor-dee
framework. Based on this benchmark, [barz et al. (2022) further proposes Triplet-GMPNN, which employs a triplet message passing schema (Dudzik & Veličković (022) and multiple improvements for the training stability of the encoder- encoder-processor- pro--psor- decoder framework. Recently, Beeiilaquua et al. (2023) proposes an additional self-supervised objective to learn similar representations for inputs that result in identical intermediate computation. a common design in the above methods is the incorporation of historical embeddings into current execution steps. In this work, we highlight that such historical dependence poses a misalignment with the Markoy nature of the algorithmic execution. This insight motivates our proposed ForgetNet and its enhanced version G- ForgetNet, which more faithfully align with the Markov nature by reconsidering the use of historical embeddings.
2

---

Published as a conference paper at ICLR 2024
Inpuc state
Lowal!
Lowefi
Inpet state
Low
Low
Inpuc state
Lozst0
Loss
3
D
3
D
D
D
(a) Baseline
(b) ForgetNet
(c) G-ForgetNet
Figure I: An illustration of (a) the baseline, (b) ForgetNet, and (c) G-ForgetNet methods. E, and D represent the encoder, processor, and decoder module, respectively.
ANALYSES ON THE MARKOV PROPERTY
In this section, first recap the existing prwooder-processo-deeooeee paradigm on the algorithmic reasoning tasks given in the CLRS-30 benchmark. Then, we emphasize the Markov characteristic of algorithmic executions. In addition, we highlight the existing misalignment between the use of historical embeddings and this Markov property. Motivated by this observation, we present ForgetNet, which removes such historical embeddings to achieve a closer alignment with the task's ss nature. An empirical study validates that ForgetNet achieves better generalization capability.
3.1 ENCODER- POODEERR-RROOSSSRRDDEOOOOS
PARADIGM
Following prior research, we consider the algorithmic reasoning tasks as formulated in the CLRS-30 benchmark (Veličković et al. (222a)). For certain algorithm, single execution trajectory serves as data sample, which is composed of the input, output, and h sints. Here, h are time series of intermediate states of the algorithm execution. Typically, a data sample is represented as graph with n nodes, where n reflects the size of particular sample. For example, in sorting algorithms, elements in the input list of length are denoted nodes. With such graph representation, the input, output, and hints s at a particular time step are either located node-level, edge-level, or graph-level features. As detailed in Veličković et al. (2022a), there are five possible types of features, including scalar, categoricaa.occccc eegeeiarra,,,,, mask, ma sk one, and pointer, each accompanied by its encoding/decoding strategies and associated loss functions.
Let us denote the node-level, edge-level, and graph-level features time step as {a; ). {e J. and
(()
(t),, respectively. Here, indicates the node index and ij specifies the index of the edge between node and j. Note that addition to the input, hints are also included these features when they are available. Most existing neural algorithmic learners (Veličković et al. 2020b Georgiev Liō 2020; Veličković et al., 2020a| 2022a Ibarz et al.) 2022 Diao & Loynd 2023) adopt the ponuudr-processo-ddeee
paradigm (Hamrick al. 2018). Specifically, each time step the encoder first embeds the current features into high-dimensional representations as
(t)
g ())))
(I)
Here, f.((), fe)., and fg(+) are the encoder layers, typically parameterized linear layers. The embeddings then fed into processor. which is parameterized graph neural network foNN(-). to perform one step of computation. The processor can formulated as
h "- (h())
(1=(0) J fANN
te g))
worth noting that the processed node embeddings from the previous step, (h (t-1)) ), are used at each time step t. Initially, h. = 0 for all nodes. Subsequently, the decoder, a linear model, uses the processed node embeddings {h (t) to either predict the hints for the next time step, the output if it is at the final time step. Note that the encoder and decoder should task-tailored based on the feature types in the particular task. Additionally, the learnable parameters of all neural modules shared over time steps. To train the described encoder proccsr-processor-- decoder model, the loss calculated based on the decoded hints at every step and the out output at the end.
where [:] denotes concatenation.
3

---

Published as a conference paper at ICLR 2024
Intermediat: Stries (Ilims)
1
Sartedlly-Srreed
2
3
5
3
3
Liets
Illustration of the execution steps in insertion sort. The top row represents the intermediate states, while the bottom row shows the corresponding partially sorted lists. At a specific step, the present state, denoted as the hintss includes the current order (the black pointers), the recently inserted element (the green pointer), and the current iterator (the blue pointer). The present state can fully determine the next intermediate state. The figure is adapted from Veličković el al. (2022a).
Figure 2:
During training, either the ground truth hi hints or the hints s predicted from the previous step can be fed into the encoder, depending on whether teacher forcing is used. During inference, the step-by-step hints are not available, and the encoder always receives the predicted hints from the previous step. In the benchmark study by Veličković et al. (2022a), the ground truth hints s are used with 50% probability during training, given that the training process would become unstable without teacher forcing. While using the actual hints can stabilize training, it introduces discrepancies between training and inference modes. Recently, Ibarz et al. (2022) proposes several techniques improve training stability, such using soft hint prediction, specific initialization, and gradient clipping tricks. More importantly, it demonstrates that, with such training stability, it is possible to completely remove teacher forcing and enforce the model rely on the hints predicted from the previous step, thus aligning the training with inference and achieving better performance. Therefore, as illustrated Figure | (a), our study in this work specifically adopts and builds on this pipeline that operates without relying on teacher forcing.
3.2 ALGORITHMIC NECESSITY OF HISTORICAL EMBEDDINGS
Markov nature of algorithmic executions. The Markov property refers to the principle that future states depend only on the current state and not on the sequence of states that preceded it. It is important note that such fundamental property holds in the context of algorithmic reasoning tasks formulated in the CLRS-30 30 benchmark because the entire algorithm state is given in each hints. s. To be specific, within an algorithm's 's sequential execution, the state information to unambiguously determine the state at the subsequent time step f. + 1, preventing the need to refer to any states preceding time step f. Let us take the insertion sort in Figure 2 as an example. At any specific step, the intermediate state, represented as the hints, completely determines the algorithmic execution output of that particular step, i.e., the next intermediate state.
time step encompasses all necessary
Misalignment with the use of historical embeddings. Given the Markov nature of the task, we revisit the necessity of using historical embeddings the existing paradigm for algorithm reasoning. As described in Section 3.1 prevalent practice in the existing encoder- processor-decoder framework is the incorporation of historical embeddings from previous steps into the current processor input. This practice, which might seem naturally borrow from design principles in graph neural networks (GNNs) and recurrent neural networks (RNNss, intends to capture and propagate potentially relevant information across time steps. However, it intrinsically contradicts the Markov nature of the task as highlighted above. Given the Markov property of tasks within the CLRS-30 benchmark, the progression of the algorithm should depend solely on the current state, given by the current hint The incorporation of historical embeddings from previous steps, while seemingly advantageous, might inadvertently add unnecessary complexity to the model. Such an addition not only complicates the model architecture but also introduces potential discrepancies and noise that might misguide our neural learners away from the desired algorithmic trajectory, consequently compromising the generalization ability.
removing the use of historical embeddings. As studied by Xu et al. (2020), it is easier for neural networks to learn reasoning tasks where the computational graph of the neural network aligns with the algorithmic structure the task since the network only needs to learn simple algorithm steps. Motivated by this intuition and the identified misalignment between the use of historical embeddings and the Markov nature of neural algorithmic reasoning tasks, we suggest removing the use of historical embeddings align the computational graph of the neural model
ForgetNet:
4

---

Published as a conference paper at ICLR 2024
100
Fargeth et
Beseline
80
Avergge Score (S)
Segments interseet
Br5
Masimum
Belmnon Ford
DAS Shortast Patts
Offkcrd
Aztulyy Setertor
Graham Scat
Deg
Invertion Sort
Mantt Chain Order
Breiees
Repllogical Sort
PST
Artitalttinn Faints
Tasi actordating
Jarvis March
4S HSTurtake
LCS Lengeh
Openaal (t1
Butble Sort
Binurr Search rei
Food Max Subarray
uaiv String Matcher
quicssor
SCC
Hepport
Reyd
kruth Moms- Pratt Quscgeeie Oveer all Avoraye
Figure 3: Comparison between ForgetNet and the baseline. Reported results are the average of 10 rums with random seeds. Numerical results can be found in Table
with the task's Markov nature. Specifically, following the notation in Eq. (2) we remove the use of (h!-") and only use the encoded node embeddings {x(") as the input node embeddings for the processor. Formally, the processor in Eq. (2) is replaced with
(1 (0) SANN ( J. (ele))))")))
(3)
While the modification of the model architecture seems simple, it non-trivially enables the updated model have a direct and coherent alignment with the underlying Markov nature of the neural algo- rithmic reasoning task. The parameterized processor can thus focus on learning the one-step execution of the algorithm, without the potential discrepancies introduced by using historical embeddings. This new streamlined framework, as illustrated in Figure (b), is termed ForgetNet.
Empirical validation. To verify our insight, using the full set of algorithms from the CLRS-30
benchmark, we train our ForgetNet alongside the existing architecture vs. Figure [ (a)). The only difference between these two models is that the historical embeddings are removed in ForgetNet. Using the standard OOD splits in the CLRS-30 30 benchmark, we perform 10 runs for each model on each algorithm task with a single set of hyperparameters. As demonstrated in Figure 3 ForgetNet improves the performance over the bascline across 23/30 algorithmic reasoning tasks. The improvements brought by removing historical embeddings are quite significant on several tasks. For example, the absolute margins of improvement on DFS, insertion sort, and bubble sort are 66.79% 24.57%, and 13.19% respectively. By focusing purely on the relevant signals at the current step, ForgetNet can generalize better OOD testing samples, fitting more useful signals for improved performance. In Appendix B. we further evaluate the performance of ForgetNet on the multi-task setup following [barz et al. (2022) These empirical studies directly verify our insight that it is effective to explicitly enforce the Markov property in neural algorithmic learners.
baseline (i.e. Figure D (b)
4 IMPROVED TRAINING VIA ADAPTIVE ALIGNMENT
In this section, we first identify the limitation of completely removing historical embeddings as suggested in ForgetNet. In particular, inaccurate intermediate state predictions at the early stage of the training will potentially lead to sub-optimal convergence. To alleviate this, we propose the G-ForgetNet model, which uses a learnable gating mechanism and an additional regularization term in order to capture the Markov property of ForgetNet without the subsequent training limitations.
4.1 LIMITATIONS OF ENTIRELY REMOVING HISTORICAL EMBEDDINGS
While our ForgetNet model demonstrates effectiveness on a diverse set of tasks, it underperforms the baseline on several tasks, such as the Warshalll algorithm. A closer examination suggests that during the early stage of training, the model struggles with producing accurate intermediate predictions for certain algorithm tasks, which could lead the model towards suboptimal convergence. To make this clear, with slight abuse of notations, we let x and (!)) denote the input state and the 1-th
5

---

Published as a conference paper at ICLR 2024
intermediate state, i.e., the hints, of an algorithmic trajectory sample, respectively. Accordingly, gio) represents the intermediate state predicted at timestep t. In addition, &, P, and 'D represent the computation included the encoder, processor, and decoder, respectively. our ForgetNet model, the computational pathway x - 3 - P + + y() naturally emerges as a desirable pathway for training. This is because both and y are accurate, facilitating high-quality back-propagation signal. However, as we progress into extended paths for subsequent execution steps, we expose the model to the pitfalls of inaccurate intermediate state predictions. For example, the computational pathway associated with the loss function for the second intermediate state, + P - D + g( + - P - D + y(2), is impacted by the inaccuracies of the prediction g()) Intuitively, this introduces noise for the processor P to learn the one-step execution of the algorithm, since the processor receives inaccurate input at the second time step. Such inaccuracies accumulate over time steps. This indicates that, during the early stages of training, the model is primarily navigating these sub-optimal pathways, hindering its optimization. Additionally, by removing the hidden states in ForgetNet, we have essentially removed a residual connection between processor layers, making more difficult for the model to backpropagate signals through many consecutive processor layers. As an empirical illustration, Figure 4 shows the training loss of ForgetNet and that of the baseline model for the Floyd- Warshall algorithm task. It indeed shows that the training losses for ForgetNet are elevated during the early stage of training. leading to sub-optimal convergence. We provide more results and deeper analysis of the training difficulties in ForgetNet in Appendix B.2 where we observe that elevated losses in NorgetNet are primarily incurred during the later steps of the hints time series, indicating the difficulties the model has with accumulation of inaccurate intermediate predictions.
4.2 G-FORGETNET: ADAPTIVE USE OF HISTORICAL EMBEDDINGS
In light of the aforementioned limitation, we further introduce G-ForgetNet with a regularized gating mechanism that restores important computational pathways during training and learns align with the Markov property. The core motivation behind this proposal is that while inclusion of information from previous layers does not align with the inherent Markov nature of the task, it can provide helpful support, especially during the early stage of training, where it can mitigate the effects of inaccurate intermediate predictions and facilitate higher quality backpropagation signals. Further, the added loss penalty encourages the model to obey the Markov property that was shown to be beneficial in Section 3.2. Specifically, by including (((11) in Eq. (2) as a component of the input for the processor at time step f, it can enrich the model with useful computational pathways, such as I -- P - D + g( associated with the loss function for the second intermediate state. In general, the introduced computational pathways x E P + there are / sequentially applied processors, are valuable for training the processor P to capture one- step algorithmic execution. This is because ( is the accurate output after executing the algorithm for steps from the input x. In essence, these pathways create an alternative route, cireumventing the challenges posed by inaccurate intermediate state predictions, especially at the early stage of training.
D - ((),, where
Based the above intuition, in G-ForgetNet, further introduce a learnable gating mechanism that modulates the use of historical embeddings. Formally, Eq. (2) is replaced with
"]. (h. = fANN (1z) (e) )."").
h
=
where the gate g. has the same dimensions as h (t-1) and O denotes element- product. Here, employ a simple multi-layer perceptron (MLP) obtain the gate
gf = MLP
(5)
where (-) is the sigmoid function. An illustration of G-ForgetNet is in Figure (c). Finally, introduce a modified hi nts loss function that includes a regularization term on the magnitude of (t) g
(os (0)
(6 (1)
+ Σ
9
(6)
Where C (g(().)) ()) is the standard hints s loss functions used in the CLRS-30 benchmark, which depends on the type and location of features contained in ((!), At the early stage of training,

---

Published as a conference paper at ICLR 2024
Table 1: Test OOD micro-F1 score of the baseline, ForgetNet, and G-ForgetNet methods. The cells are highlighted if their corresponding results are better than the bascline.
Algorithm
Baseline
ForgetNet
G-ForgetNet
Algorithu
Baseline
ForgetNet
G-ForgetNet
Activity Selector
93 1
97.17% 20
-9.03%
Jarvis' March
85. 14.11% 20
85, 26,211%42...
88, 88.53%
Articulation Points
a5.01
90. 30.16%
97.97%
Knuth-Morris Pratt
3.96%
28.9666
45% 3.32
BeilmannFoord
97. 2.
96. 4356m0.38
99.18%
LCS Length
76. =1
8 =0.09
85.43%
47
BTS
99.45%
99.32%
99.96%
Matrix Chuin Order
86.91%
94.83%
91 00%%
Binary Search
62.799
74.1156=2..1
-5.96%
Minimum
96. 12%8-=..0
99.01% =0.10
99.26%
Bradges
80.58% 4.79
94.74% 12.00
99.4 99.41%%.1111 10:15
87. 42.4%% 1.112
85.08% 11.12
91.25% 1.0
Bubble Sort
-1.99%---5.24
75.1256=229
83.19%++1.0
MST--Prim
94.14%
95.19%
DAG Shomest Paths
97.92% Louns
98%18% 10.34
99.37%% 1o.0
Naive String Mutcher
80. 80.32% 16.06
62. 22% 1a
97.02% 27
DES
301.33% 4-477
97.12%
74.31%
Optimal BST
78, 18.14%%
10%
83.58%
Dijkstra
961.88%-+0.r
98. =0.58
99.14%
Quicksclect
1.45%%=0084
1.6196=0..0 S8
6.30%
Find Max. Suburray
60.677%
64.57% 1.442
78.97% 19.70
Quicksort
47. 26.34
61.92% 1.6 22
73.28%
Floyd- Warshall
-3.0119-=2.177
14.0
56.32%%-+..6
Segments Intersect
9T. -7..11
98.55%
99.000 99.06%.
Grahum Scan
91. 01.82%% 11:20
95.40%% 10.27
97.67%% 10.14
SCC
44.83%% 1.965.777..4 12:74
50.80%
12.657
53.53% 1.2 2-18
Heupsort
48. -344
46. 06.90%
57.47%
+5.008
Task Scheclaling
80. 0:22
86,31 26-5%-...000
84. 11.555%
Insertion
70.62%%.1x.95
9551996 10.77
98.40% LOVR1
Topological Sort
84.87%61399
92.63%
99.92%
intuitively anticipate the gate to more "open", i..., the magnitude of g." to be large, thus enriching the model with the aforementioned beneficial pathways. As training progresses and the model starts predicting more reliable intermediate predictions, the dependence on historical embeddings should diminish, i.e., the gate becomes more "closed", to honor the Markov nature. Since the scale of the hints losses varies drastically for each algorithm, we use a heuristic to select the value of À for each algorithm based on the loss values; further details can be found in Appendix A.I
EXPERIMENTS
In this section, we perform comprehensive experiments to evaluate the proposed G-ForgetNet model, by addressing the following questions. (1) Can our G-ForgetNet model, equipped with the regularized gating mechanism, consistently perform better than the baseline model? How does it perform in the several tasks where ForgetNet underperforms the baseline? Does it help the early stage of training? (2) How is the G-ForgetNet model compared to a boarder range of prior methods? (3) What are the dynamics of the gating mechanism within G-ForgetNet? To what extent does it align with our expectations?
Datasets and setup. We perform experiments on the standard -iistiiiiitiittionnn (OOD) splits present in the CLRS-30 algorithmic reasoning benchmark (Veličković et al. 2022a). To be specific, we train on inputs with 16 or fewer nodes, and use inputss with 16 nodes for validation. During testing, for most algorithms, there are 32 trajectories with inputs s of 64 nodes. For algorithms where the ouputs s are associated with graph-level features, rather than node-level or edge-level, there are 64x more trajectories, ensuring a consistent number of targets across all tasks.
For the baseline, ForgetNet, and G-ForgetNet introduced Section 3.2 and 4.2 respectively, we conduct 10 runs for each model in each task, with single set of hyperparameters. Specifically, employ the Adam optimizer (Kingma & Ba, 2015) with a cosine learning rate scheduler and an initial learning rate of 0.0015. The models are trained for 10,000 steps with batch size of 32.
we further include more existing state-of-the of-the-art methods for comparison. Specifically, first involve the notable methods studied in the CLRS-30 benchmark, including Memnet (Sukhbaatar et al. 2015), MPNN (Gilmer) al. 2017). and PGN (Veličković et al.) (0200a). These models serve as the processor in the encoder- processor-decoder framework, which is trained with noisy teacher forcing in the benchmark setup. Furthermore, we include the recently proposed Triplet-GMPNN
More baselines. Beyond the baseline paradigm introduced in Section 3.1
method (Ibarz al., 2022). which develops set of techniques to stabilize training, thus removing teacher forcing completely to align the training and inference. Moreover, the processor network in Triplet-GMPNN a message passing network which incorporates messages from triplets of nodes. Our introduced baseline, ForgetNet, and G-ForgetNet, i.e., the three methods in Figure [ are built on the framework as developed by [barz al. 2022). Differing from the baseline described in Section 3.1 Triplet-GMPNN has additional update gate. It is worth noting that such a gate is different from our introduced gating mechanism in G- G-ForgetNet in terms of both motivation and architectural design. In particular, the update gate in Triplet-GMPNN is placed ahead of the decoder and aims to update the processed embeddings for a subset of nodes at each time step and keep the remaining unchanged, whereas our gate mechanism is
7

---

Published as a conference paper at ICLR 2024
Table 2: Test OOD micro-F1 score of G-ForgetNet and existing methods, The highest scores highlighted in bold, while the second-highed scores are underlined. Results for individual algorithms can be found in T
Algorithm
Memnet
MPNN
PGN
Triplet-GMPNN
G-ForgetNet
Div. & C.
13.05%+4.08
20.30%
65.23%=2.om
76.36%+o.43
78.97%:#0.70
DP
67.95%++2.19
65.10%=20.3
70.58%==0.81
81.99%+1..3
86.70%:t0
Geometry
15.14%.12.36
73.11%
61.19%::1.14
94.09%
95.09%:11... 1.16
Graphs
24.12%=#1.6
62.80%+2.55
60.25%%+1.5
81.41%+11.5
Greedy
53.42%+#1.13
82.39%#11.7
75.85%=1.2
91.22%:++0.40
91.79 0.23.
Search
34.35%
11.20%
36.11%6:....
58.61% L:005
63.84% %6:10.844
Sorting
71.53%++.99
11.83%##.b 91
15.46%=1.18
60. 60.38%++..2 2.27
78.09 %
Strings
1.52%
3.21%
2.04%:10.16
49.09%.444
54.74%
Overall Average
38.03%
51.02%
52.31%
75.98%
82.89%
99%
0/30
1/30
1/30
1/30
9/30
97%
0/30
1/30
1/30
5/30
13/30
95%
0/30
2/30
2/30
7/30
14/30
designed to enforce the Markov property of algorithmic reasoning and is supported by regularization term in the loss function. Another recent model, Hint-ReLIC (Bevilacqua al., (203)) uses an additional self-super learning objective based on data augmentations. Given such augmentations and different setups, our model and Hint-ReLIC are not directly comparable. We expect that a fusion of our model and Hint-ReLIC could further boost the performance, and we leave such an evaluation to future work as the code of Hint-ReLIC not yet publicly available.
G-ForgetNet vs. NorgetNet vs. the baseline. In Sec- tion 3.2 we have demonstrated the effectiveness of our ForgetNet model which removes historical embeddings to honor the Markov nature of algorithmic reasoning. Here, we further compare the three methods included in Figure [
3.0
Baseline
ForgetNet
G-Fo G-ForgetNet
Training Loss
2.0
to evaluate the effectiveness of our proposed gating mech- anism in the G-ForgetNet model. In T we report the average test results over 10 runs for each model in each algorithm. While ForgetNet surpasses the baseline across 23/30 tasks, G-ForgetNet consistently achieves improved performance over the baseline on all 30 tasks. In the sev- eral tasks where ForgetNet underperforms the baseline, such as the Floyd- -arrsall and naive string matcher tasks, G-ForgetNet demonstrates consistent improvements over the baseline. For example, in the naive string matcher task, while ForgetNet performs worse than the baseine, G-ForgetNet outperforms the baseline by an absolute mar- gin of 16.70% This demonstrates the effectiveness of the proposed gating mechanism, which is able to capture the benefits honoring the Markov property without the training difficulties of ForgetNet.
1.0
W
0.0
2000
4000
6000
8000
10000
# Training Steps
Figure 4: Training curves for the base- line, ForgetNet, and G- G-ForgetNet meth- ods on the Floyd- Wardsall task. The shaded region indicates the standard de- viation. Figure is smoothed for clarity.
As clarified in Section 4.2 the proposed gating structure is expected to enhance the early stage of training, thus improving the final convergence in many tasks. To empirically verify such intuition, in Figure 4 we illustrate the training losses of the baseline, ForgetNet, and G-ForgetNet models in the Warssall task. We observe that ForgetNet indeed faces challenges during the early stages, leading to sub-optimal convergence compared the baseline in this task. The G-ForgetNet model, can effectively sidestep the early training pitfalls, thereby leading to better convergence at the end of training in this task. This verifies our intuition that the additional computational pathways in G-ForgetNet can help enhance the early stages of training. In Appendix B we dive deeper into the loss curves corresponding to different execution steps for several algorithms and demonstrate that the loss experienced by ForgetNet each execution step tends to escalate more sharply as the algorithmic execution progresses than G-ForgetNet. This observation validates our earlier intuition in Section 4.2 that the gating mechanism in G-ForgetNer introduces computational pathways that act corrective
8

---

Published as a conference paper at ICLR 2024
signals against accumulated errors. By offering these pathways. G-ForgetNet can circumvent the pitfalls posed by inaccurate intermediate predictions, thereby facilitating the optimization of the losses corresponding to later execution steps. Overall, G-ForgetNet outperforms ForgetNet in 26/30 tasks and improves the overall average score from 78.98% in ForgetNet to 82.89% in G-ForgetNet.
Compared to more existing methods. We further extend our comparison of G-Forget to more existing methods, including the aforementioned Memnet, MPNN, PGN, and Triplet-GMPNI methods.
The results of these methods are obtained from the respective literature (Veličković et al. 2022a
Ibarz el al. 2022). As summarized in Table G-ForgetNet emerges as the top performer in 25/30 algorithmic tasks. Compared to the previous sfate-of-the the method Triplet-GMPNN improves the mean test performance across all 30 tasks from 75.98% to 82.89%. Additionally, G-ForgetNet surpasses the 99% threshold on 9/30 algorithms, compared to the prior best of just 1/30. Further, G- G-ForgetNet achieves large performance increases on several algorithms. For example, G-ForgetNet achieves a test score of 97.02% in the naive string matcher task, while the previous state-of-the-ar performance is 78.67%, and G-ForgetNet achieves test score of 98, 440 on insertion sort, compared to the previous state-of-the-art of 18. This comparison further demonstrates the effectiveness of our proposed method.
G-ForgetNet
Dynamics of the gating mechanism. In order to under- stand the behavior of the gating mechanism and gauge its alignment with our anticipations, we empirically in- vestigate its dynamics during training. Specifically, we
Hidden State Norm
compute the L2 norm of the hidden states, h before be-
Baseline G- ForgetNett
ing passed to the processor and then normalize by dividing
by the square root of the hidden state dimension. In G-
0.4
ForgetNet, the L2 norm is taken after gating g." @ h (1-1
so we are effectively measuring how much of the hidden states are allowed to pass into the processor. For every sample in the validation set, we consistently monitor the average L.2 norm over both nodes and algorithmic exe- cution steps, along the training iterations. In Figure [5] we illustrate the average L2 norm over all samples in the validation set during the training process for the Floyd- Warshall task for the baseline and for G-ForgetNet. We observe that the baseline hidden state norm is fairly con- stant and has a relatively large magnitude, indicating that it is fitting historical noise during training, whereas G-ForgetNet declines to nearly zero. This empirically validates that the dynamics of the gating mechanism align with our intuition in this task. That is, the gating mechanism is open during the beginning of training, thus enhancing early training while progressively focusing on the Markov nature of algorithmic tasks. We generally observe similar trends across all of the CLRS-30 30 algorithms, with more tasks shown in Appendix A. We further
0.2
0.0
2000
4000
6000
8000
10000
# Training Steps
Figure 5: Average L2 norm value throughout the training process on the Floyd-Warshall Warshall task. The shaded region indicates the standard deviation.
validate the importance of the loss penalty included in G-ForgetNet in Appendix A.3 where we investigate the behavior of the G-ForgetiNet model without the loss penalty. We observe that without the loss penalty, the model still exhibits declining trends in the hidden state norm, however it will not converge to 0 as desired. The performance of G-ForgetNet without the penalty still better than the baselines, however the performance is significantly improved with the penalty. This aligns with our intuitions since the penalty ensures that G-ForgetNet is consistent with the Markov property.
CONCLUSION
In this work, we highlight a key misalignment between the prevalent practice of incorporating historical embeddings and the intrinsic Markov characteristics of algorithmic reasoning tasks. In response, we propose ForgetNet, which explicitly honors the Markov nature by removing the use of historical embeddings, and its adaptive variant, G-ForgetNet, equipped with a gating mechanism and subsequent loss penalty in order to capture the benefits of the Markov property without the training difficulties found in ForgetNet. Our comprehensive experiments on the CLRS-30 benchmark demonstrate the superior generalization capabilities of both models compared to established baselines. In summary, this work reveals the importance of aligning model design with the Markov nature in neural algorithmic reasoning tasks, paving the way for more advancements in future research.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
This work was supported in part by National Science Foundation grants IIS-2243850 and IIS- 10S-200666.
REFERENCES
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, al. Relational inductive biases, deep learning, and graph networks. arXiv preprint 1806:080000110000 2018.
Luca Beurer- Kellne,, Martin Vechev, Laurent Vanbever, and Petar Veličković. Learning to configure computer networks with neural algorithmic reasoning. Advances in Neural Information Processing Systems, 35:730-742, 2022.
Beatrice Bevilacqua, Kyriacos Nikiforou, Borja Ibarz, Ioana Bica, Michela Paganini, Charles Blundell, Jovana Mitrovic, and Petar Veličkovié. Neural algorithmic reasoning with causal regularisation. In International Conference on Machine Learning, pp. 2202-222222 PMLR, 2023.
Thomas H Cormen, Charles E Leiserson, Ronald L Rivest, and Clifford Stein. Introduction to algorithms. MIT press, 2022.
Andreea-loana Deac, Petar Veličković, Ognjen Milinkovic, Pierre-Luc Bacon, Jian Tang, and Mladen Nikolic. Neural algorithmic reasoners are implicit planners. Advances in Neural Information Processing Systems, 34:15529 -15:5599--1552,,,, 2021.
Cameron Diao and Ricky Loynd. Relational attention: Generalizing transformers for graph- struturrd tasks. In The Eleventh International Conference on Learning Representations, 2023.
Andrew J Dudzik and Petar Veličković. Graph neural networks are dynamic programmers. Advances in Neural Information Processing Systems,
2022.
Vijay Prakash Dwivedi, Chaitanya K Joshi, Anh Tuan Luu, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Benchmarking graph neural networks. Journal of Machine Learning Research, 24 (43):1 143) 2023.
Dobrik Georgiev and Pictro Lió. Neural bipartite matching. arXiv preprint arXiv:2005.11300 2020.
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pp. 1263-1272. PMLR, 2017.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv preprint arXiv:1410.5401, 2014.
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwińska, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et al. Hybrid computing using neural network with dynamic external memory. Nature, 538(7626): 471-476, 2016.
Jessica B Hamrick, Kelsey Allen, Victor Bapst, Tina Zhu, Kevin R McKee, Joshua B Tenenbaum, and Peter W Battaglia. Relational inductive bias for physical construction in humans and machines. arXiv preprint arXiv: 180..002000..... 2018
Borja Ibarz, Vitaly Kurin, George Papamakarios, Kyriacos Nikiforou, Mehdi Bennani, Róbert Csordás, Andrew Joseph Dudzik, Matko Bošnjak, Alex Vitvitskyi. Yulia Rubanova, al. A generalist neural algorithmic learner. In Learning on Graphs Conference, pp. 2-1, PMLR, 2022.
Armand Joulin and Tomas Mikolov. Inferring algorithmic patterns with stack-augmented recurrent nets. Advances in neural information processing systems, 28, 2015.
Eukasz Kaiser and Ilya Sutskever. Neural gpus learn algorithms. arXiv preprint 1r110:151,.08228, 2015.
10

---

Published as a conference paper at ICLR 2024
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In The Eleventh International Conference on Learning Representations, 2015.
Sadegh Mahdavi, Kevin Swersky, Thomas Kipf, Milad Hashemi, Christos Thrampoulidis, and Renjie Liao. Towards better out-of-distribution generalization of neural algorithmic reasoning tasks. Transactions on Machine Learning Research, 2022.
Danilo Numeroso, Davide Bacciu, and Petar Veličković. Dual algorithmic reasoning. In The Eleventh International Conference on Learning Representations, 2023.
Sainbayar Sukhbaatar, Jason Weston, Rob Fergus, et al. End-to-end memory networks. Advances in neural information processing systems, 28, 2015.
Petar Veličković and Charles Blundell. Neural algorithmic reasoning. Patlerns, 2(7), 2021.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks. In International Conference on Learning Representations, 2018.
Petar Veličković, Lars Buesing, Matthew Overlan, Razvan Pascanu, Oriol Vinyals, and Charles Blundell. Pointer graph networks. Advances in Neural Information Processing Systems, 33: 2232-2244, 2020a.
Petar Veličković, Rex Ying, Matilde Padovano, Raia Hadsell, and Charles Blundell. Neural execution of graph algorithms. In International Conference on Learning Representations, 2020b.
Petar Veličković, Adrià Puigdomènech Badia, David Budden, Razvan Pascanu, Andrea Banino, Misha Dashevskiy, Raia Hadsell, and Charles Blundell. The CLRS algorithmic reasoning benchmark. In International Conference on Machine Learning, pp. 22084-22102. PMLR, 2022a.
Petar Veličković, Matko Bošnjak, Thomas Kipf, Alexander Lerchner, Raia Hadsell, Razvan Pascanu, and Charles Blundell. cosooning-modllttee representations. In Learning Graphs Conference, pp. 50-1. PMLR, 2022b.
Louis- Pascal Xhonneux, Andreea-loana Deac, Petar Veličković, and Jian Tang. How to transfer algorithmic reasoning knowledge learn new algorithms? Advances in Neural Information Processing Systems, 14:10000995222 2021.
Keyulu Xu, Jingling Li, Mozhi Zhang. Simon S Du, Ken-ichi Kawarabayashi, and Stefanie Jegelka. What can neural networks reason about? In International Conference on Learning Representations, 2020.
Yujun Yan, Kevin Swersky, Danai Koutra, Parthasarathy Ranganathan, and Milad Hashemi. Neural execution engines: Learning to execute subroutines. Advances in Neural Information Processing 33: 33:17298-17308 2020.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poezos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. Advances in neural information processing systems, 30, 2017.
Wojciech Zaremba and Ilya Sutskever. Learning to execute. arXiv preprint arXiv:1410.4400 2014.

---

Published as a conference paper at ICLR 2024
100
Average Score (N)
60
G-ForgetNat
FargetNeet
Baseline
Topelogical Sort
BPS
DAG Shartest Patts
Bridges
Minimum
Belimar-Fard
Segmants Intereett
Dijkatra
Actisity Selecror
Insertion Sort
anticulttin Paints
Naive Strong Matcher
Graham Stan
Maeris Chain Order
Prem
MSTAenakaal
jareis Mard ch
Binory Search
LeS Length
Tack Schedlaing
Optinal RST
Aind Max Subaray ay
Bubble Sort
DFS
Quicksort
Heassort
Royd-Warhhall
Kmith MorthMorri-PPar Pratt
sor
Quiceselect
Overall Average
Figure 6: Comparison between the baseline, ForgetNet, and G-ForgetNet. Reported results are the average of 10 runs with random seeds. Numerical results can be found in TableT
A FURTHER G-FORGETNET ANALYSIS
In Figure we compare G-ForgetNet with ForgetNet and the baseline model, and in Table 3 provide the full numeric results for G-ForgetNet compared with Memnet, MPNN, PGN, and Triple-- GMPNN.
A.1 LAMBDA HEURISTIC
Since the scale of the loss function varies drastically between algorithms, it is not possible to use single value for across all algorithms. In general, is non-trivial to select optimal values for this parameter, and in this work, select such that the gate penalty makes approximately half of the total training loss after 6,000 training steps. Intuitively, with such schedule, the model will spend the first 6,000 steps simply learning to execute the algorithm, then during the remaining 4,000 steps. the model will focus on learning single-step executions accordance with the Markov property. As demonstrated Table 2 this simple heuristic has quite robust performance across the entire set of CLRS-30 algorithms. We acknowledge that such simple penalty schedule and heuristic is unlikely be optimal and will be approved upon by future works.
heuristic select reasonable values for A. Specifically,
A.2 GATE ANALYSIS
In Figure 7 we include more figures the hidden state's 's L.2 norm for G-ForgetNet and the baseline, in Section 5 These further support that G-ForgetNet model does enforce the Markov property during testing as we observe the L2 norm converge to 0.
12

---

Published as a conference paper at ICLR 2024
Table 3: Test OOD micro-FI score of G-ForgetNet and existing methods. The highest scores are highlighted in bold, while the second-highest scores are underlined.
Algorithm
Memnet
MPNN
PGN
Triplet-GMPNN
G-ForgetNet
Activity Selector
24.10%.42.2
80.66 40.66%.13.1
06.80%=1.62
95.18%
99.03%.0..
Articulation Points
1.50%
50.91 %
49.53%=22.09
88.32%
97.97%
Bellman-Ford
40.04%
92.01%
92.99%++0.3
97 40.119
99.18%.4..1
BFS
43.34%.10.04
99.89%
99. 6..663%20002
99. 99.73%
99.96%:10:0
Binary Search
14.37%
36.83% 26
76.95%220.1
77 77.58% 2..3
85.96%11.59
Bridges
30.26%c+t0.oo
72.69% 4-7778
51.42%%77..
93. 13.99% +
99.43%::..15
Bubble Sort
73.58%.1..77
5.27%%20.60
6.01% 1111
67 77.68%%15.. 1:5.0
83.19%:2.59
DAG Shortest Paths
66.15%c+1..
96.21%++0.64
96.941
18.19%
99.37%:#0.03
DFS
13.36%
6.54%
8.71%=0.22
47.79% 1-149119
74.31%
Dijkstra
22.48%+12.30
91.50%+10.60
13.45711.775
96.05 16.05%%...00
99. 10.14%
Find Max. Subarray
13.05%f14.08
20.30%:10..9
23..3%%+2.5
76.36% 1-0.344
78.97%
Floyd- Warshall
14.17%.tt..3
26.74%=1.77
28. 76%
48. 52. 1.1.
56.32%22..
Graham Scan
40.62%6422.31
91.04%#+..31
56.87%
93 22.66%%--o...
97.67% 0.
Hcapsort
68.00%:1..77
10.94%
5.27%
31. 11.04%%.5..
67.47%%.1 08
Insertion Sort
71.12%++0..5
19.81 2.02
+1.37%
78.14% 18.14%
98.40% 00.02.1
Jarvis" March
22.99%::13.67
34.86%
49. 19%
91.01% 1:1300
88 28.53%%-12
Knuth-Moris-Pratt
1.81%%
2.49%
2.00%=0.12
19.51 tt4457
12.45%++..11
LCS Length
49.84%
53.23%%10.36
56.82 56.82%10.22
80.51 1.164
85.43%
Matrix Chain Order
81.96%
79.84%
83.91% L0..9
91.68% L..6%%....0 1.0.59
91.08% &00.51
Minimum
86.93%:to.11
85.34%220.6 88
87. 87.719 T025522
97. 20.55
99.26%
MST-Kruskal
28.84%:40.61
70.97%+1.50
66.96%
89.80% 4-0.777
91.25% 4..400
MST-Prim
10.29%%t:3.77
10..8%%7..6
63.33 13.33%.L...
86 15.39%%11.3
95.19%
Naive String Matcher
1.22%
3.92%++.30
2.08 % 00220
78.67%+-44 4-4.99
97.02%
Optimal BST
72.03%.+1.21
62.23%%10.44
71.01% 11.1%%11.8
73.77%
83.58%
Quickselect
1.74%%:t0.0
1.43%+40.6
0.47%
6.30%=+0..5
Quicksort
73.10%
11.30%
6. 6.
64. 14.64%%-55. 1.5.12
23.28%
Segments Intersect
71.80%+19.50
93.44%220.10
77.51 %o700.75
97 7..64% F-00000
09.06%%20030
SCC
16.32%
24.37%
20,80%
43.43% L:3.15
53.53% %
Task Scheduling
82.74%c40.04
84.11 %
84.89%220.0
87.25%
84.55% 70.00.1...
Topological Sort
2.73%
52.60%
60. 60.45%
87.27%
99.92%:+4.02
Overall Average
38.03%
51.02%
52.31%
75.98%
82.89%
99%
0/30
1/30
1/30
1/30
9/30
97%
0/30
1/30
1/30
5/30
13/30
95%
0/30
2/30
2/30
7/30
14/30
13

---

Published as a conference paper at ICLR 2024
Activity Selector
Bridges
1.0
Hidden State Norm
0.8
Hidden State Norm
0.8
0.6
Baseline
0.6
Baseline
G-ForgetNet
G-ForgetNet
0.4
0.4
0.2
0.2
0.0
0.0
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
# Training Steps
# Training Steps
Binary Search
Topological Sort
1.0
Baseline
1.0
G-ForgetNet
Hidden State Norm
0.8
Hidden State Norm
0.8
0.6
0.6
Baseline
ForgettNe
0.4
0.4
0.2
0.2
0.0
0.0
2000
4000
6000
8000
10000
0
2000
4000
6000
8000
10000
Training Steps
# Training Steps
Figure 7: Average L2 norm value throughout the training process on the activity selector, bridges, binary search, and topological sort algorithms. Shaded regions indicate standard deviation.
14

---

Published as a conference paper at ICLR 2024
Table 4: Test OOD micro-F1 score of the baseline, G-ForgetNet without penalty, and N-ForgetNet methods. The cells are highlighted if their corresponding results are better than the bascline. The highest scores are highlighted in bold.
Algarithon
Baseline
G-ForgetNet wif penaltty
G-ForgeiNet
Algarithmn
Baseline
G-Forgyt Net w/o penalty
G-Porget Not
Activity Salezeer
95 05.00 om1
99.035
Earvis Manch
85 25.6455an.0
58 5355
Asticalation Points
95.00 Re-7
98.99%
07 0055- 00
Knuth-Muri-PPatt
3. 3.9695e1.1
3. 3.84%
12.45%
Bellmon-Fard
97.67%
576 2050
99.
LCS Length
76. a
80) H0.3556
85. 85.43%
DFS
90 4555
00 57
99. 96
Marii Chain Order
86, 86.51%6a5.0
56. 53.5335
91.88%.11 S.
Binary Search
62. /9/ 6:3
82. 495
85.965
Minimum
96. a2.-225
99.50 29.56%
50 20.25555
Bridges
89 20.558%
99.43%
MIT Kruskal
ST. 425 42.-225.....22
91.58% 1...0
91 25.55% 15.40
Bubble Soet
61 LL.50F
66 664.78%
83.19%
MST Paim
92 222.555
944 20.5599...0
95, 95.19%
DAC Shoricst Puths
MT 12.22%
018.70 71/5
99.37
Naive String Matcher
80, 20.-2256444. 64.99
58, 7096 1.3
97, 77.822%
DFS
30. (0.586mmm
47.97%
74.31
Optinal BST
78, 78.1456-44
79.14%
83 15.5%% Gouu
Dijkers
06 (0.77555m
08 4655
99.14%
Quicksolect
1,655
2. 2.00%
6.30%
Find Mas Subormy
65 034675
77 6375--
78.97%
Quicksret
47. 9058
70. 175
73, 23,28%5
Ployd Wanshlll
53.00FL
54 0056..0...4
56.32%
Segments Intersect
97. 8356.....
99.27%
Gratam Scan
B1 12.22
07 02.552
97.675
SCC
19,78%
50 $0.85%612
08.00P6-55
59 39089%
Schadaaling
80.28%
84.55%.4..3
Insertion
70062%
87.90%
Ropological Sori
84.6775.43.0
90.30%61225
99.92%%.15.2
A.3 PENALTY ANALYSIS
In Table we report the performance our G-ForgetNet model without the loss penalty, i.e., with just the gate mechanism. We observe that, even without the loss penalty, the model still outperforms the baseline on 28/30 algorithms, however the average score is only 79.31% % compared with the penalty. Additionally, G-ForgetNet without penalty outperforms G-ForgetNet with penalty on algorithms, however these cases are very small improvements. Overall, this study empirically shows the importance of the penalty in the G-ForgetNet model, which aligns with our intuition that the penalty is necessary in order to enforce the Markov property. Finally, we provide a comparison of the L2 norm of the gated hidden states in G- G-ForgetNet with and without the penalty in Figure g On the activity selector algorithm, the G-ForgetNet model without the penalty still consistently decreases during training. however it does not reach the same final convergence the model with the penalty, and on Floyd-Warshand Warshall, G-ForgetNet without the penalty is fairly constant throughout training, again demonstrating the necessity of the loss penalty enforce the Markov property in G-ForgetNet.
52.88%
Activity Selector
Floyd-
1.4
Baseline
1.4
Baseline
G-ForgetNet w/o Penalty
G-ForgetNet w/o Penalty
1.2
G-ForgetNet
1.2
G- G-ForgetNet
Hidden State Norm
Hidden State Norm
1.0
1.0
0.8
0.6
0.4
0.4
0.2
0.2
0.0
0.0
2000
4000
6000
8000
10000
0
2000
4000
6000
8000
10000
Training Steps
# Training Steps
Figure 8: Average hidden state I.2 norm value throughout the training process on the activity selector and Floyd-Warshall
Warshall algorithms. Shaded regions indicate standard deviation.
B MORE EXPERIMENTAL RESULTS
B.1 MULTI-TASK EXPERIMENTS
Prior works Xhooneex et al. 2021 Ibarz et al. 2022) have investigated jointly learning multiple algorithms using a single processor. We follow the multi-tat task setup in [lbrz et al. (2022) and train a single ForgetNet processor all 30 CLRS algorithms while keeping separate encoders and decoders
15

---

Published as a conference paper at ICLR 2024
100
ForgetNet
Average Score [N)
60
20
Marri Chain Order
grs
Brdges
Minium
jacvi' March
Gratiam Scan
Actisity Selector
Dijketra
Insortion Strt
Topelogical Sort
DAG Shartest Pattis
Segments Intersect Articulation Points
NST -rer
MST-Anavkal
Binary So Search
Quicksort
LGS Lengts
Task schedblinn
Optinal EST
Bellmar- Ford
Heaport
Bubbiee Sort
Find Mox. Subarray
soc
Naive string Matcher
DrS
Knuth Morris art Flayd Flayd Warshall Quiceselect
Overail Average
Figure 9: Comparison between the ForgetNet multi-task model and the baseline multi-task model.
Average of 10 runs with random seeds. Numerical results can be found in Table
Table 5: Test OOD micro-F1 scores of the baseline and ForgetNet on the multi-task setting.
Algorithm
Baseline
ForgetNet
Algorithm
Baseline
ForgetNet
Activity Selector
88.14%
94.16%
Jarvis' March
74.84%=5..
95.79%=0.o
Articulation Points
94.70%=2.55
92.69%+0
Prath-Morris-Pratt
51.93%=4433
25.16%=0.88
Bellman-Foo Ford
$7.21%11.56
76.41%:3.45
LCS Length
83.68%=0.55
36.85%=0..1
BFS
99.77%
99.32%++..29
Matrix Chain Order
89.60%=0.44
97.83%.=..1
Binary Search
69.63%
89.52%(11.2
Minimum
92.73% 1.24
95.87 0.55
Bridges
96.38%
98.00%:#0.00
MST-Kruskal
89.30%=0..8
89.56% 1.68
Bubble Sort
53.24%
74.40%
MST--rim
86.54%=00.00
90.38% 0.22
DAG Shortest Paths
91.49%
93.87% 1-036
Naive String Matcher
46.74%
31.90%
DFS
36.35%c+2.00
38.51 18.51%
Optimal BST
74.23%==2.1
82.55%
Dijkstra
94.49%
94.34%+4.00
Quickselect
1.71%=#0.49
9.15%
Find Max. Subarray
51.15%::1.23
35.19%:10.6
Quicksort
75.31%=44.
88.48%=24
Floyd- Woyd-Warshlll
24.99%:11.42
20.39%:10.10
Segments Intersect
95.26%.10.25
92. 71 0.33
Graham Scan
68.05%%+1..0
95.42%%+t0.25
SCC
18.73%%=4400
66.93%
Heapsort
55.83%.17.9
78.39% 1:4.92
Task Scheduling
86.61 16..21
84.02% 0.
Insertion Sort
87.48%
93.98%co.99
Topological Sort
81.85% 1.21
93.88% 0.448
ForgetNet performs better than the baseline on 20/30 algorithms and has a higher average score across all algorithms. This confirms that enforcing the Markov property is also beneficial in the multi-task setting.
for each algorithm. As shown in Figure and in Table 5
B.2 MORE TRAINING CURVES
More training curves of the baseline, ForgetNet, and G-ForgetNet methods. In Figure 10 we illustrate the training curves, including the total loss and losses at different execution steps, of the three methods in the Floyd-Waad Warshall and activity selector tasks. We observe that ForgetNet does not converge to the same total loss value, shown in plot (a) on each figure. Further, we observe that the gap between ForgetNet and the baseline widens at later points in the hints time series, showing how the removal of connections between consecutive layers ForgetNet introduces training difficulties. The inaccurate intermediate predictions cause the model to struggle to optimize losses corresponding to later execution steps.
16

---

Published as a conference paper at ICLR 2024
3.0
Flayd-Warshall
Baseline ForgetNet
2.5
Training Loss
G-ForgetNet
1.5
1.0
2500
2000
4000
6000
8000
10000
(1)
Training Steps
15
(a) Total Loss
(b) Intermediate losses at different exccution steps
1.0
Activity Selector
Baseline
Training Loss
G-ForgetNet
0.6
0.4
0.2
0.0
2500 5000 7500 10000
2300 7500 10000
2500 5000
750010000
2500 5000 7300 10000
2000
Training Steps
4000
6000
8000
10000
201
(a)
Total Loss
(b) Intermediate losses at different execution steps
Figure 10: Training curves for the baseline, ForgetNet, and G- G-ForgetNet methods in the Floyd- Warshall and activity selector tasks, tasks (from top to bottom): (a) total loss and (b) losses different execution steps, i.e. the losses incurred at different points in the hints time series.
17


Reviews:
[{"review_id": "rhw5JSG3p0", "reviewer": "Reviewer_VNPs", "summary": "The paper focuses on the Markov property of neural algorithmic reasoning. More specifically, neural networks that learn to imitate algorithmic execution have thus far used as additional inputs data from the algorithmic traces at previous execution steps, not just the current step. The paper emphasises that, as the studied algorithms are Markovian, a better alignment between the neural network and the task should mean not using data from previous execution steps, which are named historical embeddings in the paper. This work uses an established architecture as the baseline, and with the removal of the historical embeddings proposes ForgetNet. Further, G-ForgetNet is obtained by using a gating mechanism to learn how much to use the historical embeddings within the baseline. The proposed changes outperform the studied baselines, and a simple ablation study on the gating mechanism shows the importance of using historical embeddings early on in training, when accumulating errors can be more harmful for final performance.", "questions": "Please see above.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is well motivated and clear. The proposed changes are supported by the empirical evaluation on the CLRS-30, outperforming the baseline in most algorithms.", "weaknesses": "The initial hypothesis — that better alignment with the markovian property should result in better performance — is strong. However, there are trainability issues (e.g. accumulating errors early in training) that can be addressed in different ways. The paper proposes one solution, gating, but some questions remain — why is gating better than the NN more generally learning how to combine historical and current embeddings, as in the baseline? Moreover, what are other solutions, possibly not involving learning?\n\nLastly, the TripletGMPNN includes a gating layer in the single-task setup. A discussion on how this differs from the gating in G-ForgetNet should be included.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "F65jrk035V", "reviewer": "Reviewer_PbVk", "summary": "This paper proposes ForgetNet, a GNN that enforces Markov structure across algorithmic reasoning steps.\n- Motivation: algorithmic tasks (e.g. sorting) can be modeled by transitioning through states which is Markov in nature, whereas existing methods use the history.\n- Solution: ForgetNet enforces such Markov structure by updating the representation at each step usingfeatures at the current step only.\n    - i.e. the hidden states at time $t$ is updated as $\\{h_i^{(t)}\\} = f_{\\text{GNN}}(\\{x_i^{(t)}\\}, \\{e_{ij}^{t}\\}, g^{(t)})$, for node features $x$, edge features $e$, and global feature $g$.\n- Further challenge: empirical results suggest that such Markov structure may cause training instability issues at early stage of training, likely because the intermediate steps are too unconstrained (note that supervision is provided only on the final state).\n- Solution: G-ForgetNet, which selectively and adaptively keeps some history using a gating mechanism.", "questions": "- Fig 4 (b): what are execution steps, and how are the labels acquired? Since there are labels to compute the per-step loss, can we also compare with per-step teacher forcing?\n- Fig 5: for better comparison, could you also provide the norm of the residual branch in the baseline?\n- I wonder if one can anneal the gating, i.e. gradually reducing the amount of the history (e.g. by controlling the gating value).\n    - One motivation is that in Table 1, when G-ForgetNet is not better than the baseline, its performance is very close to the baseline. This suggests that G-ForgetNet falls back to the baseline in these cases, hence ensuring the history embeddings are actually not being used may help with performance.\n    - Another reason is that according Fig 5, the gate value is ~0.38 even at the end of training. This means the history embeddings are still helpful, which weakens the claim that the task should be Markov, unless the history embeddings are effectively removed.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Another way to state the results is that\n1. for better generalization, the authors proposed to restrict the function class by excluding non-Markov solutions, leading to the proposed ForgetNet.\n2. the restricted function class raises optimization challenges. To address these challenges, the author incorporates a gating mechanism that selectively incorporate the history.\n\nBoth these changes are natural/simple yet effective: the proposed G-ForgetNet outperforms baselines in most tasks in the CLRS-30 benchmark.", "weaknesses": "- The discussion about the motivation could be improved.\n- About the motivation that \"algorithmic reasoning tasks are Markov\": it should be clarified that this work is about algorithmic reasoning tasks that can be modeled by finite-state automata. In general, whether the process is Markov or not depends the definition of the state space.\n- Even though the Markov observation motivates to remove history information, history information proves to be important for stabilizing training (hence the proposal of G-ForgetNet). This is similar to the effect of residual links, which can be discussed more in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Iyzl3h1WjB", "reviewer": "Reviewer_fqct", "summary": "This work studies a gap in the existing GNN literature -- while encoder-processor-decoder architectures typically have hidden states passed between calls to the processor, this work investigates the best way to handle those connections in the setting of algorithmic reasoning.", "questions": "1. In most algorithms ForgetNet beats the baseline (Figure 3). Is there any intuition or hypothesis around which particular algorithms the baseline is better? For example, for Naive String Matcher the baseline looks considerably better. Why might that be?\n2. I really like the visualization with bars in Figure 1, I think adding G-ForgetNet to that graphic or otherwise including a plot with the G-ForgetNet results would help. Can the authors provide such a visualization?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The work is original to my knowledge. In fact, I went through related work looking for experiments, claims, and ablation studies relevant to this paper. I expected to see some justification or discussion of the hidden states and the way they are used in other GNN architectures. I could not find a sufficient example and so I think question the architecture and showing that omitting hidden states all together or otherwise limiting the information flow through those connections is a valuable addition to the literature.  \n2. The quality of the writing is good. I think the paper is well motivated and clearly written.\n3. The conclusions are significant. The fact is that on a 30-problem benchmark suite, the two architectural changes proposed in this work generally improve performance.", "weaknesses": "1. The discussion of the Markov nature of algorithms could use some nuance. If it is the case that these algorithms, when processing entire graphs, therefore having some history in the 'state' are Markov, then the historical information in the hidden features shouldn't help. In fact, this work shows that a gated layer is better than no connection for passing information with the hidden states. Thus, I think the story arc is a bit confusing. It actually took me several readings to really piece together the fact that the treatment of the hidden state is under discussed in the literature and the two options presented in this paper are better than the existing approaches as a result of the confusing Markov narrative. This is perhaps a small point, and I'm only one reader -- if the Markov framework helps convey the story to others this is my weakness, and not the paper's but I though it was worth mentioning here.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lesoG3MWWC", "reviewer": "Reviewer_gWrY", "summary": "This paper aims to reconsider the Markov property of neural algorithmic reasoning. While classical algorithms inherently possess the Markov property, existing neural algorithmic reasoning approaches learn to execute algorithms using historical embeddings, which contradicts this property. Building on this observation, this paper proposes ForgetNet, a method that eliminates the reliance on historical embeddings to reintroduce the Markov nature. Additionally, the paper introduces G-ForgetNet, which adaptively integrates historical embeddings.", "questions": "Please see the weaknesses above.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. This paper focuses on an emerging and interesting research topic: neural algorithmic reasoning.\n2. This paper is well-written. Figures and formulas clearly illustrate the background and proposed method.\n3. The motivation for aligning current algorithms with their Markov nature is well stated. Additionally, this paper complements a crucial concept in the context of neural algorithmic reasoning.\n4. Experiments show that the proposed method can outperform baselines in most of the tasks (algorithms).", "weaknesses": "1. Although ForgetNet perfectly possesses the Markov property of algorithm execution, the newly proposed G-ForgetNet still utilizes historical embeddings but outperforms ForgetNet. Is there an inherent drawback of neural algorithmic reasoning in possessing the Markov property?\n2. For certain tasks, G-ForgetNet performs worse than ForgetNet, or ForgetNet performs worse than the baseline. Some task-specific explanations regarding this should be discussed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JHQ2XfG9zE", "reviewer": "Reviewer_QaJ5", "summary": "The paper points out that most algorithmic tasks should have Markov property and thus the optimal model should not rely on historical embeddings. The paper then proposes ForgeNet and GForgerNet which doesn't use or only limitedly use historical embeddings.", "questions": "See weakness,", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. There is a novelty in observing the Markov property of many algorithmic tasks and proposing an architecture that aligns with this property.\n2. The overall performance seems to be better than all compared baselines.\n3. The paper is quite well-written and easy to follow.", "weaknesses": "The authors claimed that the gating value should decrease because  \"As training progresses and the model starts predicting more reliable intermediate predictions\". But if there are no regularisation terms on the gating variable, why should the gating value decrease if there is still useful information in the historical embeddings?  The authors do show a plot of the gating variable trajectory as training progresses, but the gating value only decreases slightly and never to 0. Are there any cases where the gating variable decreases to 0? If not, why is this the case since the model can just simply ignore the historical embedding after training stabilizes?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)