NeuralSolver: Learning Algorithms For Consistent and Efficient Extrapolation Across General Tasks
Authors: Bernardo Esteves
Miguel Vasco
Francisco S. Melo
Venue: NeurIPS 2024 poster (2024)
Keywords: Deep learning
algorithm synthesis
recurrent networks
algorithmic reasoning
sequential decision making
extrapolation
Primary area: deep_learning_architectures
TL;DR: A recurrent network that learns scalable algorithms that extrapolate on larger general tasks without fine-tuning.

Abstract: We contribute NeuralSolver, a novel recurrent solver that can efficiently and consistently extrapolate, i.e., learn algorithms from smaller problems (in terms of observation size) and execute those algorithms in large problems. Contrary to previous recurrent solvers, NeuralSolver can be naturally applied in both same-size problems, where the input and output sizes are the same, and in different-size problems, where the size of the input and output differ. To allow for this versatility, we design NeuralSolver with three main components: a recurrent module, that iteratively processes input information at different scales, a processing module, responsible for aggregating the previously processed information, and a curriculum-based training scheme, that improves the extrapolation performance of the method.
To evaluate our method we introduce a set of novel different-size tasks and we show that NeuralSolver consistently outperforms the prior state-of-the-art recurrent solvers in extrapolating to larger problems, considering smaller training problems and requiring less parameters than other approaches.

NeuralSolver: Learning Algorithms For Consistent and Efficient Extrapolation Across General Tasks
Bernardo Esteves"
Miguel Vasco
Francisco S. Melo INESC-ID, Instituto Superior Técnico, Universidade de Lisboa
INESC-ID. Instituto Superior Téenico, Universidade de Lisboa
Department of Intelligent Systems KTH Royal Institute of Technology Stockholm, Sweden
Abstract
We contribute NeuralSolver, a novel recurrent solver that can efficiently and consistently extrapolate, i.e., learn algorithms from smaller problems (in terms observation size) and execute those algorithms large problems. Contrary previous recurrent solvers, NeuralSolver can be naturally applied in both same-size problems, where the input and output sizes are the same, and in different- size problems, where the size of the input and output differ. To allow for this versatility, we design NeuralSolven with three main components: a recurrent module, that iteratively processes input information at different scales, a processing module, responsible for aggregating the previously processed information, and a curriculum-basee training scheme, that improves the extrapolation performance of the method. To evaluate our method we introduce a set of novel different-sin tasks and we show that NeuralSolver consistently outperforms the prior state-of- the-art recurrent solvers in extrapolating to larger problems, considering smaller training problems and requiring less parameters than other approaches. Code available at
Introduction
Humans can solve complex reasoning tasks by extrapolating: employing and combining elementary logical components build more elaborate strategies. Machine learning models excel at pattern recognition, often outperforming humans in classification 22 control 3 and prediction tasks 4 However, these models still struggle at reasoning which affects their ability maintain their performance for increasingly harder versions of the same task (5). A particular type of difficulty emerges from the increase in dimensionality of the input. Intuitively, learning algorithms to perform tasks, e.g., finding the goal in maze (Figure I becomes increasingly harder for larger problems.
We focus designing learning algorithms that able efficiently extrapolate, i.e., able learn algorithms to solve tasks over small problems (in terms the dimensionality of the input) and execute over arbitrarily large problems, without a significant loss in performance nor additional fine-tuning. Recently, a novel class of methods, which we denote by recurrent solvers, have been proposed that employ recurrent neural networks to extrapolate (5] o.. Compared
feed-forward networks, which have finite depth, recurrent networks can be iterated to an arbitrary depth at execution time. By doing so, these models can extrapolate to significantly larger problems than the ones seen during training by executing more recurrent steps. However, current recurrett solvers can only be applied to same-size problems, such as the case of image generation, where the dimensionality of the input (e.g., image size) is the same as the dimensionality of the output. This limitation inhibits the use of these methods wide range of tasks, which we call different-sizz tasks, where the input and output dimensions are different (e.g., classification and decision-making tasks).
"oorrespondence to: bernardo. www.eeetteeneii.ee...... ulisboa.pp
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

G2 50
Figure 1: Observations of the 1S-Maze environment of sizes 7x7, 7, 1 11x 11, x 33 and x29 129, where the agent (green square) must go to the goal (red square). The light green arrow represents the next target action that the model needs to predict, while the purple path represents the sequence of actions required to solve the maze.
We address this limitation and propose novel recurrent solver that is able consistently and efficiently extrapolate regardless of the size of the output and the type of problem. We name our approach NeuralSolver, novel architecture for general tasks (same-size and different-size). Our approach introduces several architectural changes to recurrent solvers: a recurrent convolutional block that, through multiple iterations, perceives information at different scales the input image, and processing block (with an aggregation function) that merges information for the output of the model. Furthermore, we introduce a curriculum-based training scheme to improve the extrapolation abilities of our model. We train standard supervised learning techniques and apply the learned algorithm at test time with arbitrarily large observations (e.g., 256 256 images, larger), with minimal loss in performance.
small observations (e.g.. 15 x 15 images) using
We evaluate NeuralSolver against prior recurrent solver in dierrture-standan same-size tasks and in novel set of different-size tasks. We show across all tasks that our model significantly outperforms previous approaches in extrapolation capability, being able execute learned algorithms in larger problems without loss in performance; in training efficiency, being able to learn algorithms from smaller problems; and in parameter efficiency, requiring 90% less parameters than similar approaches.
In summary, the contributions of our work are:
NeuralSolver: We propose novel recurrent solver that uses a recurrent convolutional
module, processing module (with an aggregation function) and curriculum-leane learminn to train algorithms small problems and extrapolate them to arbitrarily large problems. Contrary to previous recurrent solver, our model can be used in same-size and different-size tasks:
Different-size tasks:
We contribute a novel set of different-size classification tasks for recurrent solvers, with images of arbitrary size and label information;
* Consistent and efficient extrapolation: We show that NeuralSolver significantly outper- forms the prior recurrnt solvers regarding extrapolation capabilities, training efficiency and parameter efficiency.
2
Related Work
Schwarzschild et al. (S] introduce a class of recurrent neural networks that can execute at test time more iterations than the number of iterations used during training, solve problems than the ones seen during training, which we denote by recurremt-solvers. To do so, they propose a recurrent network architecture based on residual neural networks (ResNets) (2) that shows logical extrapolation abilities across same-size tasks by leveraging the spatial invariances in the inputs. However, these networks suffer from overthinking, where the network deteriorates in performance when the number of iterations is extended beyond the training distribution. To solve the overthinking problem Bansal et (6] introduces two new components to the previous architecture: (1) recall module, that integrates the input directly into specific layers of the recurrent network, safeguarding against potential loss or corruption of deep features; and (2) progressive loss (PL) training scheme that incentivizes recurrent networks to iteratively refine the feature representation, preventing the network from memorizing the number of recurrent module applications. The authors show that their improved network is able to extrapolate to problems 16 times larger than the training size with
complex

---

Q.1
0.9
0.0
0.0
Observation
Recurrent
Processing
Output
Observation
Recurrent
Processing Module
Output
Module
We design NeuralSolver with two fundamental components: (i) recurrent module (pur- ple), responsible for iteratively processing the input data regardless of its size; (ii) processing module (green), with an optional aggregation layer (A), responsible for generating the output and allowing architecture be used both in same-size and different-size tasks. Additionally, employ curriculum-based training scheme to improve the extrapolation performance of our architecture.
Figure 2:
more than 95% accuracy the same-size benchmark. We this method baseline against method. To address the problems of the computational efficiency and hyperparameter tuning of the progressive loss, 77] propose replacing the progressive loss with delta loss term [8] and achieve similar performance to the original one. We propose a novel architecture that is able to outperform Bansal et al. [6] without the need of any additional loss counter the overthinking problem (as shown in Appendix [.I). NeuralSol ver is the first recurrent solver able to consistently and efficiently extrapolate learned algorithms both same-size and different-size tasks.
3 The NeuralSolver Architecture
We address the challenge of designing model that is able to efficiently learn algorithms from small problems that consistently extrapolate to larger problems. Contrary to previous recurrent solvers [5} 7 we want to apply model to both same-size tasks, where the input and output have the same dimensionality, and different-size tasks, where the input and output have different dimensionality. To address these challenges, contribute NeuralSolver, depicted in Figure
3.1 Model Architecture
We design our model with three components: (i) we extract information from the input data using recurrent module, responsible for iteratively processing the input observation at different scales, allowing our model to cope with arbitrarily large input data; (ii) information then sent to processing module, responsible for aggregating (if necessary) the processed information and generating the output the network, allowing our model to used both same-size and different-size tasks; (iii) to train model employ curriculun-buct based scheme, that improves the extrapolation performance of NeuralSolver to larger problems.
This module consists of layer normalized convolutional LSTM [HII. LSTM that uses convolutional layers instead of linear layers in recurrent structure. The layer normalization is used to normalize the output of the convolutional layers and the cell state. We always provide the original input observations in the recurrent iterations of the LSTM, thus preventing the network from forgetting the input information over time, similar to the recall module in (6]. We maintain the dimensionality of the input and output of this module constant in order perform multiple iterations over the data, allowing our model process arbitrarily large observations. We show in Appendix B.I that our recurrent module allows NeuralSolver to not suffer from overthinking, despite not employing any progressive loss during training. Moreover, in Appendix C.5 we compare the effect different choices of recurrent architectures the performance our model.
Recurrent Module:
This module consists of block of three convolutional layers (following 6 whose input the last hidden state of the recurrent module. Their purpose is reduce the number of channels to desired output number of channels. Additionally, we have an optional aggregation layer (in our case a global max-pooling layer) that can employed in different-size tasks to reduce variable input size of the network to a fixed output size. For example, for two-dimensional input observation, we first reduce the input to a x tensor, where the desired output number of channels, and subsequently flatten the output tensor. This design allows our architecture be used
Processing Module:

---

iecration:
ilcraion::
Herratinn: 15
Meraaion: 20
RDLC RDLU
RDLE
RDLC
Figure 3: Propagation of information in NeuralSolver in maze-like environment: the goal is for the agent (green) to find the goal position (red). Top: the difference between the value of the internal state of the recurrent module each iteration step and the value the final iteration. Larger differences are shown in dark blue and smaller differences in white. Bottom: additionally, we show the action probabilities predicted by the processing module of the model at different iterations, where the agent can move right (R), down (D), left (L), or up (U).
both in same-size and different-size tasks, with minimum changes. In Section 5.3 compare the performance of different aggregation layers.
Curriculum Learning Training: We train our model on set of smaller-size size observations and test the performance of the model with larger dimensional observations (all belonging to the same task). Additionally, we employ curriculum learning approach in different-size tasks to counter the effect of the reduced training signal (due to the smaller output dimensionality) 12 [
we initially train the models on lower-dimensionalii observations, and then gradually increase the dimensionality of the observations every epochs. To reduce the risk of catastrophic forgetting, we sample a minibatch of observations with previously seen dimensionality (chosen uniformly at random) with a 20% chance, following previous work (12). In Section 5.3 we show the importance of the curriculum-base training scheme in the extrapolation performance of NeuralSolver.
We evaluate the role of each of these components in our overall performance in Section 5.3 For
additional details on the complete architecture of the method please refer to Appendix B
3.2 Propagation of Information in NeuralSolver
To understand how information is processed in the recurrent module of our approach, we focus on the value of the internal state of the convolutional LSTM as a function of the number of iteration steps performed. In Figure 3 we show maze-like environment different iteration steps, where the task is the find the next position in the optimal path between the current position (in green) and the goal position (in red). We highlight the difference between the value the internal state of the convolutional LSTM at each iteration step and the value of the internal state at the final iteration step (not shown in the figure). As the number of iterations increases, we observe that larger number of positions in the maze become white as, for those positions, the value of the recurrent state does not change anymore. The speed of this convergence depends on the receptive field of the convolation: for example, single convolution with a kernel size of 3 can capture the information from the adjacent pixels, propagating the information forward in one direction for a single pixel. As such, by using convolutional recurrent neural networks can propagate information across the input image by performing an appropriate number of iterations, regardless of the size of the problem.
This propagation of information also influences the output prediction. In Figure 3 (bottom), we highlight that the prediction of the output label is uncertain before iteration #10. However, in this iteration, the hidden state near the start goal (in green) has converged to the final value and the algorithm becomes certain the correct label (in this case, "Down"). We provide more examples of learned algorithms in Appendix
4 Evaluation
We evaluate deuralSolver to demonstrate how ourperforms previous state-of-the-art recurrent solvers accordingly three main criteria: (i) extrapolation capability, the ability to execute learned

---

(a) GoTo
(b) 1S-Maze
(c) Pong
(d) Doorkey
Figure 4: We introduce set of different-sizz classification tasks to evaluate the performance of recurrent solvers. In all tasks, the input is an image observation of the environment with arbitrary
size. The output is n-dimensional one-hot vector with: 4: b) = 4: c) 3; d) 4.
algorithms in larger problems; (ii) Iraining efficiency, the ability to learn algorithms from smaller problems and still maintain a high-level of extrapolation capability; and (iii) the higher parameter approach against previous recurrent solver baselines
efficiency of the model. We compare (Section 4.3) considering both same-size and different-size problems (Section 4.2). Appendix[B.. [B.3 present the architecture of our model and training hyperparameters used in our evaluation.
4.1 Methods
Model Training: We apply NeuralSol in supervised learning scenarios with datasets of input images/features and output images/features, in the case of same-size tasks, output labels, in the case of different-size tasks. As such, we train our models using cross-entropy loss on datasets with small image sizes and evaluate the classification performance of models in datasets with larger input images. In different-size tasks this approach is similar to typical imitation learning [14].
Model Checkpoint Selection: We consider a training and validation set split of 80% and 20%, respectively. To select the model checkpoint to be used for evaluation, we use the one that has the best performance on the validation set. In case of ties select the later model checkpoint.
Evaluation Metries: our evaluation we run the model for a large number of iterations (detailed in Appendix B.3 and consider the best accuracy obtained by the models at any iteration (6] We present the average best accuracy (in percentage) accompanied with the standard deviation. Additionally, we use the Almost Stochastic Order (ASO) test for statistical significance [15 [6], using significance level of a = 0.05 and employing the Bonferroni correction (17) for multiple comparisons. For more details on the ASO test and the statistical significance results, we refer the reader to Appendix D
4.2 Scenarios
For same-size tasks we use the benchmark recently proposed by Schwarzschild et al. [ consisting of three tasks: a logical toy task (Prefix-Sum) maze-solving task (Maze), and a chess puzzle task (Chess). Additionally, we modify the Maze task to allow smaller size inputs, which we denote by Thin-Maze. For detailed description of these tasks we refer the reader the original paper.
Due to the lack of literaturesstandard set of different-size tasks for algorithm extrapolation, we introduce four new classification scenarios. These new scenarios consider as input images of arbitrary sizes and as output vectors of fixed dimensions (more details Appendix (A]:
GoTo (Figure 4ap: Inspired by the Minigrid environment 18 the goal of the task to select the action that moves the agent (in green) closer the exit position (in red). The classification target is nou--dimensiona one-hot vector, corresponding to the available actions (up, down, left, right).
1S-Maze (Figure 4b): Inspired by the Maze environment from the same-size task benchmark, the objective of the task is to select the action that moves the agent (in green) closer to the goal position (in red). The classification target is four-dimensional one-hot vector, corresponding the available actions (up, down, left, right). The name an abbreviation of "1 Step Maze".
Inspired by the Atari Pong game, the goal of the task is to select the action that moves the paddle horizontally to underneath the ball. The classification target three- dimensional one-hot vector, corresponding the available actions (left, right, no action).
Pong (Figure 4
5

---

Table 1: Extrapolation accuracy the same-size tasks benchmark proposed in Schwarzschild et al. [ and the Thin-Maze environment, with corresponding training and evaluation sizes (8t, sr). Higher is better. All results averaged over 10 randomly-selected seeds. We highlight the best average results. We use (f) to indicate stochastic dominance (émin = 0) and (*) to indicate almost stochastic dominance (émin 0.5) of NeuralSolver over the baseline. We evaluate Bansal et al. 6
trained with progressive loss (PL) and without it.
Prefix-Sum
Maze
Thin-Maze (11,61)
Chess (8,8)
Model
(32, 512)
(24, 124)
NeuralSolver
100.00 +0.00
100.00 + 0.00
99.97 + 0.09
84.30 10.40
Bansal et al.
99.78 +0.78
86.63 +28.21 *
46.78 #37.49
79.99 #3.27 f
Bansal et al. 6 PL
100.00 00.01
91.13 1:7.37
42.76
137.56
82.96
10.19
FeedForward
0.00
0.00
0.00 +: 76.95 20.35
Very Small
Small
Medium
Large
100%
Accuracy
50%
0%
128
256
512
48
128
256
512
32
128
32
96
128
(a) NeuralSolver
(b) Bansal et al.
6
(c) NeuralSolver
(d) Bansal al. 6 Mazes
Prefix-Sum
Prefix-Sum
Mazes
Figure 5: Training efficiency of NeuralSolver and Bansal al. [6] same-size tasks: present the accuracy of the learned algorithms on extrapolating to problems with different dimensionality (columns), Each color represents a different training size, specific to each task, detailed in Ap- pendix A.3 In the dashed line we show the upper-bound on the performance.
Inspired by the Doorkey environment in Minigrid 18] the goal of this multi- step task is to reach the goal position (red pixel on the bottom right) by first picking-up key and using it to open a lock door. The classification target is four-dimensional one-hot vector, corresponding to the available actions (forward, rotate right, grab, toggle).
Doorkey (Figure 4d):
4.3 Baselines
Bansal et al. [6]: The current state-of-the-art recurrent solver architecture, that combines recurrent Resnet network with recall module and progressive loss. We employ the original source code from the authors and suggested training hyperparameters (when available).
FeedForward:
non-recurrent feed-forward network, based on the ResNet architecture, that has a fixed number of layers. This baseline is used to show the importance of the recurrent module for extrapolation.
Random: A lower- bound baseline, that randomly samples a classification label. This model used to help understand the worst-case performance of the models in different-size tasks.
As the baseline models cannot be applied directly in different-size tasks, we modify them by intro- ducing global pooling layer similar to that of our model.
5 Results
5.1 Same-size Tasks
Extrapolation Capability: Table T we present the extrapolation accuracy of NeuralSolver
against the baselines in same-size tasks [5. NeuralSolver achieves state-of-the-an extrapolation

---

Table 2: Total parameter count the used models on the same-size tasks benchmark proposed in Schwarzschild et al. [5] and the Thin-Maze environment, in the scale of millions of parameters.
Model
Prefix-Sum
Maze
Thin-Maze
Chess
NeuralSolver
0.168
0.053
0.053
0.699
Bansal et al. [6]
3.124
0.784
0.784
12.057
FeedForward
58.804
17.888
17.888
285.735
Table 3: Extrapolation accuracy on the different-size tasks with training and evaluation sizes (8r, 8T). We employ curriculum learning to train all methods, with training sizes indicated in the table. Higher is better. All results are averaged over 10 randomly-selected seeds. We highlight the best average results. We use (f) to indicate stochastic dominance (Emin = 0) of NeuralSolver over the baseline.
1S-Maze
GoTo
Pong
DoorKey
Model
([7, 15]], 129)
(66, 20). 128)
(66, 20), 128)
([6, 20), 128)
NeuralSolver
100.00 +0.00
100.00 10.00
100.00 0.00
100.00 0.00
Bansal al. 6
74.14 12.60 f
64.32 18.86 f
71.97 10.70 f
97.13 + 1.46
FeedForward
72.47 =0.85 f
56.82 #3.74 f
48.99 +11.5 f
79.22 11188
Random
29.74 +0.41 f
29.46 +0.40
37.87 0.47 f
29.52 + 0.58
performance on same-size tasks. In particular, our approach is able to achieve a 100% accuracy on the Prefix-Sut Summiix-mmmma and Maze tasks, when evaluated on scenario with observations 16 and times bigger, respectively, than during training. On the more complex Chess task, our model still outperforms the other baselines, achieving an average accuracy of 84.30%. In Appendix B.1 we show that NeuralSolver also does not suffer from overthinking.
Training Efficiency: In Figure we show the accuracy performance of the learned algorithm with model when trained with different problem sizes. The results show that NeuralSol vee outperforms Bansal et al. [6] in extrapolating when trained with smaller input observations, achieving consistent upper-bound performance when trained in problems with only 8-dimensional (Prefix-Sue Sum) and 16- dimensional (Mazes) input observatons. The results for the other environments are in Appendix C.I
Parameter Efficiency: In Table 2 we present the total parameter count of NeuralSolven ver against the baselines in same-size tasks [5]. The results show that model requires less than 10% of the total parameters of the baselines.
5.2 Different-size Tasks
Extrapolation capabilities: In Table we present the extrapolation performance of NeuralSoll Solr against the baselines in different-size tasks. Our model also outperforms prior state-of-the-ae recurrent solvers across all different-size tasks: approach is able to achieve upper-bound performance when extrapolating to observation sizes 9 (1S-Maze, and (GoTo, Pong, Doorkey) times larger than the provided during training. In Appendix C.7 explore the setting of extreme extrapolation, evaluating our model on very large image sizes (256 x 256 and 512 x 512). The results show that, in such challenging conditions, NeuralSolver is able to consistently extrapolate without losing performance, while the baselines fail to do so.
Training Efficiency: In Figure 6 we show the performance of model against the Bansal et al. 66 model with different training sizes. The results show that, while the previous state-of-the-an consistently under performs across most tasks on bigger test sizes, NeuralSolver is still able to learn algorithms with suitable extrapolation abilities (often close to the upper-bound performance) when trained with smaller observations. The observed gap of our method upper-boune performance also highlights that the novel set of different-size tasks can be employed to benchmark the extrapolation performance of future recurrent solvers when training using only extremely small problems.
Parameter Efficiency: NeuralSolver employs the same architecture across all different-size tasks, with 0.23 million total parameters. This corresponds to 92% reduction in the number of parameters
7

---

Table 4: Extrapolation accuracy of different ablated versions of NeuralSolver in the proposed different-size tasks. Higher is better. All results are averaged over 10 randomly-selected seeds. We highlight the best average results. We use (f) indicate stochastic dominance (Cmin 0) and (*) to indicate almost stochastic dominance (emin < 0.5) of our default model over the ablated versions.
Model
1S-Maze
GoTo
Pong
DoorKey
NeuralSolver
100.00
100.00 + 0.00
100.00
100.00
Use AvgPool
81.96 288.50
97.00 = 3.87
92.86 111.6 f
51.86
=38.98
Use 5L
78.00 + 5.78
84.16
116.39
100.00 + 0.00
92.99 = 7.62
No CL
95.26 7.29
64.12
123.50
100.00 + 0.00
97.23
4.73
No LSTM
81.87 110.09
67.28 +14.51
76.12 +17.66
93.68
7.88
Very Small
Small
Medium
Large
100%
NeuratSolver
Accuracy
50%
0%
33
97
129
32
64
96
128
64
128
64
96
128
Accuracy
Bansal
50%
0%
33
97
129
32
128
64
128
96
128
(a) JS-Maze
(b) Go
(c) Pong
(d) Doorkey
Figure 6: Training efficiency of NeuralSolver and Bansal et al. [6] different-size tasks: present the accuracy the learned algorithms on extrapolating to problems with different dimen- sionality (columns). Each color represents different training size, specifie to each task, detailed in Appendix A.3 In the dashed line we show the upper-bound on the performance.
against Bansal et al. 6 (3.15 million parameters) and reduction of 99.6% against the FeedForward baseline (71.57 million parameters).
5.3 Ablation Study
In Table we perform an ablation study on the components of NeuralSolver, particularly on the role of: (i) the aggregation function, (ii) the depth recurrent convolutional block, (iii) the training method and (iv) the type of recurrent layer in the convolutional block For (i), we replace the global max-pooling layer the processing module with an average-pooling layer (Use AvgPool). For (ii), we use 5 convolutional layers in the recurrent module (similar to Bansal et al. (6]) instead of the single convolutional layer (Use 5L). For (iii), we remove the curriculum- based training scheme (No CL). For (iv), replace the LSTM with ResNet block (No LSTM).
The results show that every component of NeuralSolver contributes the overall extrapolation performance of the method. The use of an LSTM layer instead of a ResNet block in the recurrent convolutional module results in significant improvement performance. This result is aligned with previous works that have shown that generalize better than recurrent ResNets [19]. The use max pooling as aggregation function results in better performance than average pooling. Removing curriculum learning also results in decrease in extrapolation performance of the method in some tasks (e-g., GoTo), highlighting the need to consider multiple (small) image sizes to extract relevant information during the training procedure.
recurrent neural networks empirically learn and

---

Table 5: Average reward returns the Minigrid Doorkey environment, with different sizes during execution. We show the average reward multiplied by 102. We employ curriculum learning with sizes (66, 20)) to train all models. Higher is better. We highlight the best average results. We use (f) to indicate stochastic dominance (émin 0) and (*) to indicate almost stochastic dominance (émin < 0.5) of NeuralSolver over the baseline. All results averaged over 10 randomly-selected seeds.
Model
x0
x2x32
64x64 64
128 x 128
Oracle
98.92 1:002
99.35
0.01
99.69 + 0.00
99.85 + 0.00
NeuralSolver
98.82 + 26
99.12 + 0.57
98.69 + 2.44
98.02 + 2.54
Bansal et al. [6]
98.47 #0.48
91.41 = 10.04 f
43.09 #31.18 f
24.71
#31.6 i
FeedForward
96.14 #1.63 63
63.51
7.63 f
23.53
7.34 f
7.19 = 3.18
(a) Oracle
(b) NeuralSolver
(c) Bansal al. 6
FeedForward
Figure 7: Example trajectories of the different methods when extrapolating to a Minigrid Doorkey environment with an image observation of size 64x64. The trajectory of the agents follows the gradient of the line (from darker to brighter). Additional examples in Appendix F
We also observe performance gain in using single convolutional layer over module composed of five layers (as used in (6), despite requiring five times more the number of iterations to compensate.
5.4 NeuralSolver Allows Extrapolation Sequential Decision-Making Tasks
We highlight the versatility of NeuralSolver by exploring visual imitation learning problems, in which, each time-step, the algorithm is provided with an image of arbitrary size, and needs output an action. These tasks can be solved through modern reinforcement learning methods [4 20 26
However, these methods lack the ability extrapolate, i.e., are unable learn algorithms (policies) on small problems, with lower-diments diver-dimeniional observations, and execute on larger problems, with gigeer-dimensionll observations. We consider the DoorKey environment and train an imitation- learning policy algorithm using NeuralSolver with dataser collected by pretrained oracle agent. Subsequently, we employ the learned algorithm in the Minigrid DoorKey environment [18), directly the agent's 's policy: at each time-step, mapping the observations of the agent (image of arbitrary size) into actions. We evaluate the performance of our method in extrapolating at execution time to larger observations against oracle policies specifically trained on each observation size.
The results in Table 3
highlight how NeuralSolver can achieve oracle-level performances in scenarios with larger observation sizes despite never being trained on such conditions. Moreover, the prior baselines struggle to maintain their performance for larger observations due the accumulation of errors during task-execution, resulting in sub-optimal trajectories as shown in Figure 7
Conclusion
We proposed NeuralSolver, a simple architecture that learns algorithms that can perform ex- trapolation across general tasks. We showed that our architecture consistently outperforms prior state-of-the-art recurrent solvers regard to extrapolation performance, training efficiency and parameter efficiency. In future work, we plan further improving the training efficiency of our method and explore as an architecture for reinforcement learning to learn perform sequential decision-man makinng tasks in online manner, while maintaining the ability to extrapolate.

---

Acknowledgements
This work was partially supported by national funds through Fundação para a Ciência e Tecnologia (FCT) with ref. UIDB/50021/2020 and the project RELEvaNT, ref. PTDC//CI--CO//5000/0000 The first author acknowledges the FCT PhD grant 2023.02298.B. BD. This work has also been supported by the Swedish Research Council, Knut and Alice Wallenberg Foundation and the European Research Council (ERC-BIRD).
References
[1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In 2015 IEEE International Conference on Computer Vision, ICCV 2015, Santiago, Chile, December 7-13, 2015, pages 1026-1034. IEEE Computer Society, 2015, doi: 20.1109/CCV2201..123..
[2] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 770-778, 2016.
[3] V. Mnih, K. Kavukcuoglu, D. Silver, A. Graves, Antonoglou, D. Wierstra, and M. Riedmiller. Playing Atari with deep reinforcement learning. CoRR, abs/ 2013.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figumov, Olaf Ronneherger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Amna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583 202.1
[5] Avi Schwarzschild, Eitan Borgnia, Arjun Gupta, Furong Huang, Uzi Vishkin. Micah Goldblum, and Tom Goldstein. Can you learn an algorithm? generalizing from easy to hard problems with recurrent networks. In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6- 14, 2021, virtual, pages 6695-6706, 2021.
[6] Arpit Bansal, Avi Schwarzschild, Eitan Borgnia, Zeyad Emam, Furong Huang, Micah Goldblum, and Tom Goldstein. mod-to--nndaa algorithm synthesis with recurrent networks: Extrapolation without overthinking. In NeurlPS, 2022.
[7] Alexandre Salle and Shervin Malmasi. simple loss function for convergent algorithm synthesis using mns. In Krystal Maughan, Rosanne Liu, and Thomas F. Burns, editors, The First Tiny Papers Track at ICLR 2023, Tiny Papers @ ICLR 2023, Kigali, Rwanda, May 5, 2023. OpenReview.net, 2023.
[8] Alexandre Salle and Marcelo O. R. Prates. Think again networks and the delta loss. CoRR, 0tss/9944.1186. 2019.
[9] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 173):1735-1780
1997. ISSN 0899-76677 doi: 10.1162/neco.1997 1997.9.8.
[10] Xingjian Shi, Zhourong Chen, Hao Wang, Dit- Yan Yeung, Wai-Kin Wong, and Wang-chun Woo. Convo- lutional LSTM network: A machine learning approach for precipitation noweasting. In Corinna Cortes. Neil D. Lawrence, Daniel Lee, Masashi Sugiyama, and Roman Garnett editors, Advances Neural Information P Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada, pages 802-8 810, 2015.
[11] Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. CoRR, abs/1607.064500 2016. doi: 10.485500arxiv.166 1601166677650
[12] Lukasz Kaiser and Ilya Sutskever. Neural gpus learn algorithms. In Yoshua Bengio and Yann LeCun, editors, 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016.
[13] Wojeiech Zaremha and Ilya Sutskever. Learning execute. CoRR, abs/14100466 2014. doi: 10.48550/ arxiv. 14...46...6
[14] R. Sutton and A. Barto. Reinforcement Learning: An Introduction. MIT press, 2018.
10

---

[15] Rotem Dror, Segev Shlomov, and Roi Reichart. Deep dominance how to properly compare deep neural models. In Anna Korhonen, David R. Traum, and Lluís Marquez, editors. Proceedingy the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Itll, July 28- August 2. 2019, Volume I: Long Papers, pages 2785, Association for Computational Linguistics, 2019. doi: 10. 18.38// 1/P19- 1266.
[16] Eustasio Del Barrio, Juan A Cuesta-Alberros, and Carlos Matrán. An optimal transportation approach for assessing almost stochastic order. The Mathematics of the Uncertain: A Tribute to Pedro Gil, pages 33-44, 2018.
[17] C.E. Bonferroni. Teoria statistica delle classi calcolo delle probabilità. Pubblicazioni del R. Istituto superiore di scienze economiche e commerciali di Firenze. Seeber, 1936.
[18] Maxime Chevalier-Boisvert, Bolun Dai, Mark Towers, Rodrigo de Lazcano, Lucas Willems, Salem Lahlou, Suman Pal, Pablo Samuel Castro, and Jordan Terry. Minigrid & miniworld: Modular & customizable reinforcement learning cuvironments for goal-oriented tasks. CoRR, abs/2306. 13831, 2023.
[19] Eric Price, Wojciech Zaremba, and Ilya Sutskever. Extensions and limitations of the neural GPU. CoRR, abs/1611 I. .00736,, 2016. doi: 10.48550/arxiv. 1611 1.00736.
[20] Peter R Wurman, Samuel Barrett, Kenta Kawamoto, James MacGlashan, Kaushik Subramanian, Thomas J Walsh, Roberto Capobianco, Alisa Devlic, Franziska Eckert, Florian Fuchs, et Outracing champion gran turismo drivers with deep reinforcement learning. Nature, 602(7896) 2022.
[21] Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Schoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and applications. arXiv preprint arXiv: 8/2.00882 2018.
[22] Volodymyr Mnih et al. reinforcement learning." " nature 518.7540 (2015): 529-533. Nature, 518:529-533, 2015.
[23] Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre,, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, (88(7839):604-699) 2020.
[24] Xu Wang, Sen Wang, Xingxing Liang, Dawei Zhao, Jincai Huang, Xin Xu, Bin Dai, and Qiguang Miao. Deep reinforcement learning: survey, IEEE Transactions on Neural Networks and Learning Systems, 2022.
[25] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling, Advances in nearal information prucessing systems. 15115084-50977 2021.
[26] Yuch-Hua Wu, Xiaolong Wang, and Masashi Hamaya. Elastic decision transformer. Advances Neural Information Processing Systems, 36, 2024.
[27] Yarin Gal and Zoubin Ghahramani. A thcoretically grounded application of dropout in recurrent neural networks. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett, editors, Advances in Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, pages 1019-1027, 2016.
[28] Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoder-decoder for statistical machine translation. arXiv preprint 4rxiv:1406.107: 2014.
[29] Vijay Veerabadran, Srinivas Ravishankar, Yuan Tang, Ritik Raina, and Virginia de Sa. Adaptive recurrent vision performs zero-shot computation scaling unseen difficulty levels. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 - 16, 2023, 2023. URI http:://pppperreeeee/ nips. cc/ cc/paper files/ |ttaa///gpgggg/////aaooaaaaaaaa0
Awsr::ettttaattttttt-tta--o Cow.eracctttoooeoeeeeeee html
[30] Dennis Ulmer, Christian Hardmeier, and Jes Frellsen. deep-significance casy and meaningful statistical significance testing in the age of neural networks. CoRR, abs/2204.06815 2022. doi: 10.48550/ARXIV. 2204.06815.
11

---

Table 6: Additional details on the datasets used in the evaluation of NeuralSolver. For the different- size tasks we show the dimensionality of the training examples used for curriculum learning.
Task
Train size
Test Size
Training examples
Test examples
GoTo
16.8.10.12.15.17..201
128
50,000
10,000
1S-Maze
[5,7,9,11,13]
121
50,000
10,000
Pong
16.8.10,,12.15.17.200
128
50,000
10,000
Doorkey
16.8,10,12,15.17:2201
128
50,000
10,000
Prefix-Sum
32
512
50,000
10,000
Maze
24
124
50,000
10,000
Thin-Maze
11
61
50,000
10,000
Chess
600,000
100,000
A
Additional Details on the Evaluation Scenarios
We evaluate Neural Sollllollve ver on the same-size tasks benchmark introduced in Schwarzschild et al. [S] and on novel set of different-size tasks. In Table we show the training and test sizes for each task, as well as the number of training and test examples. For the different-size tasks, we employ curriculum learning during training and show the dimensionality of the training examples in the table.
A.1 Different-size Tasks
Due the lack of benchmark for different-size tasks, we introduce set of four classification tasks with arbitrary input (image) observation dimensionality and fixed output size (classification labels). For each task generate dataset, consisting of 50,000 training examples, with 80% used for training and 20% for validation. Additionally, we include 10,000 test examples for evaluation purposes, containing observations with a larger dimensionality than the training examples.
GoTo: Inspired by the Minigrid Environment [8], the GoTo task is a simple grid setting, with a 1 1-pixel white border, and 1 1-pixel green player, and 1-pixel red goal, as shown in Figure [al The classification target is a one-hot vector of 4 dimensions, corresponding to the next-step actions available the agent: up. down, left, right. To generate the dataset algorithm from the player to the goal, where the agent starts by minimizing the vertical distance to the goal until it reaches the same vertical position, and then minimizes the horizontal distance until it reaches the goal.
a simple path-finding
This task is based on the Maze task from [5]. The input is maze with the player and goal positions, and the objective is to move the player (in green) to the goal (in red) position by solving the maze. We use the official dataset from the original paper, where made the green player have random start position so that it does not become trivial to solve. We also changed the thickness the walls and the border to 1 pixel, as shown in Figure 4b) The classification target one-hot vector of 4 dimensions, corresponding to the next-step actions available to the agent: up, down, left, right. The dataset is generated using depth-firs first search algorithm, where the label is the next action to take to reach the goal.
MS-Maze:
Pong: Inspired by the Atari game, the Pong task is simplified version of the game, where the objective is to move the paddle horizontally to the ball position, as seen in Figure The classification target is one-hot vector of 3 dimensions, corresponding to the next-step actions available to the agent: left, right, no action. The dataset is generated by a simple agent that follows the ball, keeping the paddle centered on the ball.
The goal is move the player (red pixel on top) to the goal (red pixel on bottom right) through sequence of steps, seen Figure [d]] moving to the key (yellow pixel), grabbing it, moving to the door (different color pixel on middle separator), opening it and finally moving to the goal. The classification target is a one-hot vector of dimensions, corresponding to the next-step actions available to the agent: forward, rotate right, grab, toggle. We generate the targets using simple oracle agent that solves the task. In the
Doorkey: This task is inspired by the Minigrid Doorkey environment
12

---

NeuralSolver
Bansal et al.
FeedForward
Random
100
Accuracy
200
400
0
500
1000
0
500
1000
0
100
200
(a) Prefix-Suu Sum
(b) Maze
(c) Thin-Maze
(d) Chess
Figure 8: Extrapolation accuracy of the different methods as function of the number of iterations in the same-size task benchmark. All results are averaged over 10 randomly-selected seeds.
original reinforcement learning environment, upon reaching the goal the agent receives reward of r - 0.9 (1), where t is the number of steps of the agent and T is the maximum episode length (if > the agent receives no reward and the episode restarts), In conjunction with the oracle agent, we use this environment to evaluate the performance of the learned algorithm in the original environment.
A.2 Same-size Tasks
We employ the same-size sine-size task benchmark introduced [5]. We present short description of the tasks used in the benchmark. For more detailed description of the tasks, please refer to the original paper.
Prefix-Sume This task consists of computing the prefix-sum module two of an input binary array, where the in bit of the output is the module two sum of the current and the previous n bits of the input.
Maze: This maze-solving task where the output is the path from the player to the goal. The training is done on mazes of size 24 x24 pixels, and tested on mazes of size 124x 124. While on the original paper, it is referred that the mazes have training size of and test size of 59, each maze has a border of pixels and the paths have width of 2 pixels, thus 9 x 2 + 3 x 2 24 and 59 3
124.
Thin-Maze: Similar to the 1S-Maze different-size task, this task adapts the previous scenario and modifies the thickness of the walls and the border to pixel. Thus the training and test sizes 11x11 and 61 x61, respectively.
The chess task comprises of chess puzzles, where the input is an 8x8x x 12 array indicating the position of the pieces on board, and output an 8x8 binary array indicating the optimal move origin and position to solve the puzzle. Using the original dataset, the training involves 600,000 puzzles below a rating 1,385, while testing employs 100,000 examples above this threshold.
Chess:
A.3 Training Efficiency Task Sizes
To simplify the interpretation, we omitted the training sizes of each task used in Figure and Figure In Table we detail the training sizes of each task.
Table 7: Training sizes of each task used for the training efficiency evaluation
Task
Very Small
Small
Medium
Large
Prefix-
4
16
32
Maze
16
20
24
28
1S-Maze
9
11
13
15
GoTo, Pong, Doorkey
8
12
15
20
13

---

NeuralSolver
Bansal et al.
FeedForward
Random
100
Accuracy
75
1000
2000
500
1000
200
400
0
500
1000
(a) 18-Maze
(b) GoTo
(c) Pong
(d) Doorkey
Figure 9: Extrapolation accuracy different methods as function the number of iterations in the different-size tasks. All results are averaged over 10 randomly-selected seeds.
B Additional Details on NeuralSolver
B.1 NeuralSolver Does Not Suffer From Overthinking
We evaluate the accuracy performance of our model and the baseline models across all same-size (Figure [ and different-size (Figure ? tasks function the number of iterations performed. The results show that NeuralSolver does not suffer from overthinking: the performance of the model remains constant as we increase the number of iterations.
B.2 Changes We Tried That Did Not Improve NeuralSolver
We present a list of architectural changes we explored for NeuralSolver that did not bring any substantial improvement to the extrapolation performance of our method:
Using batch normalization in the ResNet blocks of the recurrent module;
Increasing decreasing the warming up period of the training;
Reducing the kernel sizes of the final output head;
Having in the ResNet block one convolutional layer with filters and the remaining ones with 3 x 3 filters;
Reducing the amount ResNet blocks in the original Bansal et architecture from two to one;
Removing the ResNet block, and using just single recall convolution;
Using an Tanh regularization in the output of the recurrent module, avoid the drift or explosion of the recurrent memory;
Using the progressive loss to improve overthinking problem in different-size tasks (did not help in standard Bansal et al. 6 and initial experiments with LSTMs)
B.3 Model Implementation and Training Hyperparameters
In Table 8 present the model implementation used across all tasks. In Table we present the hyperparameters used for each task. In the same-size tasks we employ the same hyperparameters as in Bansal et al. [6]. with the exception of the Maze task, where we increase the training to 1,50 epochs and keeping the decay schedule at epoch number 100, as present in the source code. We also removed the warm-up for all our models, since they brought no performance improvement to the models.
B.4 Computational Complexity
We evaluate the computational complexity the models, by computing the amount of Multiply- Add Operations (MACs) for a single training example across all tasks in Table 10 in the scale of billions or gigaMACs. The results show that NeuralSolver more computationalyy efficient than Bansal et al. (6). requiring less than 30% of the operations per training example.
14

---

Table 8: Model implementation details of NeuralSolvee across all tasks. For different-siz tasks, the output size is variable. We employ the same architecture in the Maze and Thin- Maze environments.
Dataset
Width # Channels in Hidden Layers
Different-size Tasks
64
64, 64,
Prefix-Sum
100
400, 200,
Maze
32
32, 8, 2
Chess
128
32, 8, 2
Table 9: Training hyperparameter values for NeuralSolver across all tasks.
Hyperparamter
1S-Maze
Pong
GoTo
Doorkey
Prefix-Sum
Maze
Chess
Optim.
Adam
Adam
Adam
Adam
Adam
Adam
SGD
Learning Rate
le-3
2.5c-4
le-3
lc-3
le-3
le-3
1c-2
Decay Schedule
[100]
[60, 100]
[100]
[100, 110]
Decay Factor
0.1
0.01
0.1
0.01
Warm-Up
10
3
Epochs
150
50
50
50
150
150
120
Clip
2.0
2.0
2.0
2.0
1.0
Curriculum Epochs
4
Weight Decay
2c-4
2c-4
2c-4
2c-4
2c-4
2c-4
Std Dropout (LSTTM)
0.3
0.3
0.3
0.3
0.3
0.3
0
Gal Dropout 27 (LSTM)
0.4
0.4
0.4
0.4
0.4
0.4
0
Eval Iterations
2000
1000
500
1000
500
1000
200
Extreme Eval Iterations
20000
2000
2000
2000
B.5 Computational Resources
All experiments were executed on workstations with Nvidia RTX 3090 and RTX 4090 with 24GB of vram each. Training and evaluating (with a maximum size of 128) a single run takes less than day to run in all tasks, with the exception of the chess task, where each run takes around 58 hours for the Bansal
6 model and 26 hours for NeuralSolven
B.6 Comparison with Bansal et al.
In this section we highlight the major differences between NeuralSolver and Bansal et al. [6] shown Figure 10 The main architectural differences between our model and the prior Bansal et al. [6] architecture are: (i) the recurrent module used; (ii) the removal of the initial projection layer; and the design processing module(iii).
The recurrent module NeuralSolver is a layernorm convolutional LSTM (Figure 11 where is the input, h is the hidden state, is the cell state of the LSTM. All convolutional layers in all models use the same length-three filters and padding scheme, without bias terms except for the convolutional layers in the LSTTM block. Since the input is always the same across recurrent iterations, we pre-compute the convolution and layernorm passes, so that at each iteration we only sum the pre-computed value.
Recurrent Module:
We observed a small negative impact in extrapolation when using an initial projection layer for the recurrent module, and such we removed from the architecture. This results can be seen in Table [8 of Section C.6
Removal of Projection Layer:
The processing module consists of three convolutional layers with widths specified in Table with ReLU activations applied after the first two layers. The last convolutional layer produces two-channel outputs used for binary pixel classification in the case of the same input and output size tasks, or N channel outputs in the case of fixed output size tasks, where N is the number of possible outputs. Our contribution introduces aggregation layer at the end of the model,
Processing Module:
15

---

Table 10: Computational complexity (in gigaMACs) of different models across all tasks, measured in terms of the amount of Multiply-Ade Operations necessary to run single training example.
Model
13-Maze
GoTo
Pong
Doorkey
Prefix-Sut Sum
Maze
Thin Maze
Chess
NeuralSolver
4.12
9.76
9.76
9.76
0.62
3.39
0.71
5.74
Bansal al. 6
15.93
37.71
37.71
37.71
3.00
13.48
2.83
23.05
FeedFoorward
15.93
37.71
37.71
37.71
3.00
13.48
2.83
23.05
NeuralSolver
X
Projection Layer
Recument Black (Resheer)
Procossing Rinck
Recumment Binck (L5TT))
Processing Riock
Aggregation Layer
y
Figure 10: Simplified comparison of the architectural differences hetween NeuralSolver and Bansal al. 66
used for tasks with different target size from the training size. The aggregation layer consists of an adaptive max pooling that reduces the entire input size to 1, with exception for the input channels that are kept the same.
complete description of the final architecture can be described as follows: for same-size tasks, given an input image o, pass the input through a single layer ConvLSTM (q) to obtain the next hidden state (h₃ and cell state (c1), h1, 61 = g(o, lgg co). The initial hidden (ha) and cell (cg) states initialized at zeros. The recurrent process is given by feeding back the previous hidden and cell states to the LSTM, hi, C. = g(o, - Q-1). After a fixed size of recurrent iterations, we can obtain the model prediction at timestep t by passing the hidden state h, through the processing module. The processing module is composed of three convolutional layers with 3x3 kernels denoted by Wi. W2, W3 and ReLU activation y. As such, we perform p: W3 >WW2 >WW h4 ). and final softmax activation (a), Dt (pt). For different size tasks, before the softmax we perform a global maxpool operation (G) that reduces the processing output height and width Øc o (G(p:)).
C Additional Results
C.1 Extended Training Efficiency on Same-size Tasks
we show the performance of the learned algorithm with our model when trained with different training sizes on the Thin Maze and Chess tasks. The results show that only NeuralSolver outperforms Bansal et al. (6] in extrapolating when trained with smaller input observations on Thin Mazes. On the other hand, in the chess task, both models perform very similarly, with Bansal et al. 66) being slightly more performant on easier puzzle data, and NeuralSolver being slightly better with more data.
In Figure 12
C.2 Training and Validation Accuracy on Different-size Tasks
In Table II and Table 12 we show the training and validation accuracy for all models in the different- size tasks. The results show that all models are trained to optimal performance.
C.3 Fine-tuning of the Progressive Loss Parameter
The Bansal al. (6] model introduces a progressive loss term help the model extrapolate. However, this term has an alpha (a) parameter that needs to be tuned to each individual task. In Table [3] we show the results of the fine tuning of the alpha parameter for different-size tasks. For each different-size task we select and present the best-performing alpha value (in bold). For the Thin-Maze task, we selected both the a and the a 0.1 for the same-size benchmark comparison. For the remaining same-size tasks, employ the author's 's suggested values for
16

---

Layemorm
LSTTM
gala logic
Convolution
Dropout
Died
pre-computed
Processing Block
Convolution
Layerrom
Diagram of layernorm convolutional LSTM used in NeuralSolven After a desired amount of recurrent iterations, we the hidden state h as input to the processing module, obtain the output.
Figure 11:
7/100k
9/200k
11/40K
13/600k
100%
Accuracy
50%
0%
33
65
97
129
33
65
97
129
700000
700000
(a) NeuralSolver
(b) Bansal et al.
6
(c) NeuralSolver Chess
(d) Bansal et al. 66 Chess
Thin-Maze
Thin Maze
Figure 12: Training efficiency of NeuralSolver and Bansal et al. 66 on same-size tasks: we present the accuracy of the learned algorithms on extrapolating to problems with different dimensionality (columns). Each color represents a different training size, specific to each task (Thin Maze, Chess). In the dashed line we show the upper-bout bound on the performance.
C.4 Ablation on Model Size of NeuralSolver
We conducted an additional evaluation of the extrapolation performance of our method, done in Section [.22 for different sizes of the LSTM's output channel in the different-size tasks. As expected, we can observe reduction in the performance of our method across most tasks as we reduce the capacity of our LSTM. However, we point out that our smallest model still achieves competitive performances against Bansal et al. [6] as shown in Table 3
C.5 Ablation the Recurrent Module for NeuralSolver
We compare the performance of our model with two different recurrent modules: LSTM (our default method), GRU 23 and LocRNN 29] In Table [5] present the results of this comparison, following the same procedure of Section 5.2 The results show that our LSTM module consistently outperforms the other alternatives. We also show the number of parameters and computational complexity in Table [6] and Table 17 respectively.
C.6 Extended Ablation Results
In Table [8] we performed additional ablation on the extrapolation abilities when we removed the layer norm component of the recurrent neural network (No LN) and when disabled the standard Py Torch dropout and Gal Dropout (27). The results show that both components help improve extrapolation, especially in the 1S-Maze task. We explored the use of a projection layer of the input like used in the Bansal et al. 6 baseline 6 and noticed a slight negative impact on the extrapolation.
C.7 Extreme Extrapolation Results
We evaluate the performance of NeuralSolver and Bansal et al. on extrapolating very high- dimensional input observations. In particular, evaluate the performance of the methods
17

---

Table 11: Train classification accuracy of all models in the different- sizze tasks.
Model
1S-Maze
GoTo
Pong
Doorkey
NeuralSolver
99.99 =0.00
99.99 +0.01
100.00
=0.00
99.69 #0.03
Bansal et al. 6
100.00
=0.00
99.99
+0.01
100.00
=0.00
99.97 #0.02
FeedForward
100.00 =0.00
100.00 +0.00
100.00 =0.00
99.85 #0.05
Table 12: Validation classification accuracy of all models in the different-size tasks.
Model
1S-Maze
GoTo
Pong
Doorkey
NeuralSolver
100.00 10.00
100.00 10.00
100.00 10.00
100.00 +0.00
Bansal al. 6
100.00 =0.00
100.00
100.00 =0.00
100.00 +0.00
FeedForward
100.00 +0.00
100.00 #0.00
100.00
#0.00
99.94 #0.04
extrapolating to 256 256 (Table 19 and 512 x 512 (Table 20 observations. Due to the large dimensionality of the tasks and the large number of iterations required to solve the tasks, we only consider 100 testing examples The results show that only NeuralSolver able to maintain perfect accuracy across all tasks. Moreover, our model does not suffer from overthinking, shown in Figure
C.8 Hyperparameter Scan
In Fig 14 we perform small hyperparameter scan the choosen default values in Table 9
'Each test the 18 -aae task with observations of size 512 x 512 took on average 28 hours on a Nvidia
4090 GPU.
18

---

Table 13: Extrapolation accuracy for different values of on the different-size tasks and in the Maze task.
a
1S-Maze
GoTo
Pong
Doorkey
Thin-Maze
0.00
67.16
1: 5.80
62.98
112.21
69.24 +: 7.60
93.81 16.84
46.78
1:7.49
0.01
63.47
112.73
64.69
9.23
71.97 110.70
97.13 11.46
42.76
#37.56
0.10
73.36 + 7.56
62.09
13.38
62.53
100.1
93.44 14.71
22.26
227.65
0.50
74.14 + 2.60
58.45 à 5.81
68.13 113.27
93.19
18.70
0.22 0.17
1.00
66.55
13375
60.52 #11.75
71.48 14.88
94.71 #2.96
5.32 = 5.84
Table 14:
Extrapolation performance NeuralSolver in the proposed different-size tasks with different number of channels in the LSTM's output and hidden state (model width). Higher is better.
Model Width
MazM
GoTo
Pong
Doorkey
64 (Default)
100.00 10.00
100.00
0.00
100.00 10.00
100.00 + 0.00
48
98.79 +3.82
99.94 + 0.20
100.00 +0.00
99.76 = 0.63
32
78.75 #3.82
96.56 130.3
100.00 #0.00
98.65 = 2.89
24
80.05 #3.31
96.96 + 7.91
100.00 #0.00
91.21 11..99
Table 15: Extrapolation accuracy of NeuralSolver with different recurrent modules in the proposed different-size tasks. Higher is better.
Model
1S-Maze
GoTo
Pong
Doorkey
LSTM (Default)
100.00 +0.00
100.00
100.00 00.00
100.00
GRU 28
83.24 10.09
99.94 1 0.13
100.00
00.00
95.85
4.34
LocRNN 29
87.65 19.13
82.56 +17.73
94.08 8.
86.02 111.7
Table 16: Total parameter count of NeuralSolver with different recurrent modules in the proposed different-size tasks, in the scale of millions of parameters. Lower is better.
Model
1S-Maze
GoTo
Pong
Doorkey
LSTM (Default)
0.231
0.231
0.230
0.231
GRU 28
0.192
0.192
0.192
0.192
LocRNN 29
0.236
0.236
0.236
0.236
Table 17: Computational complexity (in gigaMACs) of NeuralSolver with different recurrent modules in the proposed different-size tasks. Lower is better.
Model
1S-Maze
GoTo
Pong
Doorkey
LSTM (Default)
4.12
9.76
9.76
9.76
GRU 28
3.19
7.55
7.54
7.55
LocRNN
29
4.33
10.25
10.25
10.25
19

---

Table 18: Extended extrapolation accuracy of different ablated versions NeuralSolver in the proposed different-size tasks. Higher is better.
Model
1S-Maze
GoTo
Pong
DoorKey
NeuralSolye
100.00 +0.00
100.00 +0.00
100.00 +0.00
100.00 +0.00
No LN
75.53 + 0.87
99.55 #1.
100.00
#0.00
97.11 #1.99
No Dropout
90.61 111.23
98.48 33005
100.00
+0.00
99.65 10.89
With Projection
97.24 8:25
99.93 #0.21
100.00 =0.00
99.45 +1.02
Table 19: Extrapolation accuracy on different-size tasks considering 256 x 256 observation size.
Model
1S-Maze
GoTo
Pong
Doorkey
NeuralSolver
100.00 10.00
100.00 + 0.00
100.00 10.00
100.00 10.00
Bansal et al. 6
77.20
60.33 12.38
65.00 49.01
100.00
Table 20: Extrapolation accuracy on different-size tasks considering 512 x 512 observation size.
Model
1S-Maze
GoTo
Pong
Doorkey
NeuralSolver
100.00 +0.00
100.00 + 0.00
100.00 +0.00
100.00 +0.00
Bansal et al. 6
79.20 =3.81
57.22 110.1
63.60 +2.42
99.71 +0.45
GoTo
IS-Maze
Pong
DoorKey
100
100
80
Accuracy
60
60
40
40
0
500
1000
1500
2000
500
1000
1500
2000
(a) NeuralSolver
(b) Bansal et al. 6
Figure 13: Accuracy results of extrapolation to large observation sizes (512x512). All results are averaged over 10 randomly-selected seeds. In the 1S-Maze environment, the real number of iterations 10x larger. Higher is better.
20

---

Optimizer
Learning
Up
95
90
80
Mosn
70
80
75
60
70
AdaGrad
Adom
AdomW RMSPtop
SGD
2.56-4
Le:3
3e-3
1e-2
0
10
Epochs
Clip
Curricnium Leamnngg Epochs
100.00
100
100
99.98
Mean Value
99.96
90
99.92
80
99.90
clip
JouSE
Woight Decay
Standaud Pytorch Dropout
Gel Dropout
100
100
100.0
95
99
99.0
Moan
98
98.5
85
97
97.5
Figure 14: Hyperparameter scan of NeuralSolver on the GoTo task. Each plot makes a single change from the hyperparameters in Table 9 The blue whiskers represent the confidence intervals at 95%. The green and red arrows represent the maximum and minimum values of the bootstrap distribution, respectively.
21

---

1.0
Forwadd
1.00
1.00
1.00
1.00
FeedForward
1.00
1.00
1.00
1.00
0.8
Bansal
1.00
0.51
Bansal ct PL.
0.00
1.00
1.00
0.6
-04
Ronsal et al.
0.00
1.00
1.00
1.00
Bansal et al.
0.00
1.00
1.00
1.00
0.2
NeuralSolver
0.00
0.41
0.28
1.00
NeuralSolver
0.00
0.03
0.03
1.00
0.0
Food Forward
et al. PL
Bansal et al.
NeuralSolveer
FeedForward
Bonsal et al. + PL.
NeuralSolver
Bansal es
Thin-Maze
Chess
1.0
Foodrdr
1.00
1.00
1.00
1.00
PeodForward
1.00
1.00
1.00
1.00
0.8
Ransal
1.00
PI.
0.6
0.4
Dansal et al.
0.00
0.78
1.00
1.00
Bansal et al.
0.35
1.00
1.00
1.00
0.2
NeuralSolver
NeuralSolver
1.00
0.0
FoodForward
Dansal et al
NeuriSollvrr
FeedForward
Bonsal et al. PL
NearnSoolvr
Bonsal
Bansal
Figure 15: Almost Stochastic Order scores the same-size task results presented in Section 5.1 ASO scores are expressed in émin. with significance level 0.05 that is adjusted accordingly by using the Bonferroni correction [17). Read from row to columm: e.g., NeuralSolver (row) is almost stochastically dominant over Bansal et al. [6] (column) in the Prefix-Sur task with emin of 0.28.
D Statistical Significance
The Almost Stochastic Order test (ASO; 15
recently introduced assess statistical significance in Deep Neural Networks across multiple runs. The ASO test measures the presence of stochastic order between two models or algorithms based on their respective sets evaluation scores. Using the individual scores of algorithm A and B across various random seeds, the method calculates test-specific value (emin), representing how far algorithm algorithm B. The lower the value of Cmin. the more likely it is that A is better than B. For Cmin 0.5, A said be almost stochastically dominant over considered superior to B, although with less confidence. We claim that A is stochastically dominant over B with a predefined significance level when 6min 0.0. We use the implementation by 30) to perform these statistical tests. We present the comparison of all pairs of models using ASO with confidence level of a 0.05 (before adjusting for pair-wise comparisons using the Bonferroni correction [7
16]
from being significantly superior to
most cases than vice versa, and thus could be
22

---

15-Maze
GoTo
1.0
Random
1.00
1.00
1.00
1.00
Random
1.00
1.00
1.00
1.00
0.8
PoodForward
0.00
1.00
1.00
1.00
FoewForward
0.00
1.00
1.00
1.00
0.6
0..
Ransal
0.00
0.38
1.00
1.00
0.00
0.19
1.00
1.00
0.2
NeuralSolver
NeuralSalver
0.0
Rendom
FeedForward
Bansal et al
NeuralSolver
Random
PoodForward
NeuralSalver
Bansal et
Pong
Doorkey
Random
1.00
1.00
1.00
1.00
Random
1.00
1.00
1.00
1.00
Forward
0.44
1.00
1.00
1.00
FeedForward
0.00
1.00
1.00
1.00
Bansal et
0.00
0.00
1.00
1.00
Bansal al.
0.00
0.00
1.00
1.00
NeuralSolver
0.00
0.00
0.00
1.00
NeuralSalver
0.00
0.00
0.00
1.00
Random
FeedForward
Bansal al.
NeuralSolver
Random
FoodForward
al.
NeuralSalve
Figure 16: Almost Stochastic Order scores of the different-size task results presented in Section 5.2
ASO scores are expressed in 6mins with significance level a 0.05 that is adjusted accordingly by using the Bonferroni correction [17]. Read from row to column: e.g., NeuralSolver (row) is stochastically dominant over Bansal et al. [6] (column) in the 1S-Maze task with mmnn of 0.00.
23

---

18-Maze
GoTo
1.0
No LSTM
1.00
1.00
0.50
0.75
1.00
No LSTM
1.00
1.00
1.00
1.00
1.00
0.8
No CL.
0.06
1.00
0.00
0.37
1.00
No CL
1.00
1.00
1.00
1.00
1.00
0.6
Use 5L
1.00
1.00
1.00
0.81
1.00
Use 5L
0.14
0.22
1.00
1.00
1.00
0..4
Use AvgPool
1.00
1.00
1.00
1.00
1.00
Use AvgPool
0.00
0.00
0.08
1.00
1.00
0.2
NeuralSolver
0.00
0.00
0.00
0.00
1.00
Nouver
0.00
0.00
0.00
0.00
1.00
0.0
No LSTM
No CL
Use 5L
Use AvgPool
NeuralSolvor
No LSTM
No CL
Use 5L.
Use Pool
Pong
Doorkey
No LSTM
1.00
1.00
1.00
1.00
1.00
No LSTM
1.00
1.00
0.87
0.02
1.00
No CL
0.00
1.00
0.50
0.07
0.50
No CL
0.59
1.00
0.36
0.02
1.00
Use 5t.
0.00
0.50
1.00
0.07
0.50
Use 5T.
1,00
1.00
1.00
0.02
1.00
Use AvgPool
0.25
0.97
0.97
1.00
1.00
Use AvgPool
1.00
1.00
1.00
1.00
1.00
NeuralSolver
0.00
0.50
0.50
0.00
1.00
NeuralSolver
0.00
0.03
0.00
0.00
1.00
No LSTM
No CL
Use 5L
Use AvgPool
NeuralSolves
No LSTM
No CL
Use AvgPool
NeuralSalver
Use
Figure 17: Almost Stochastic Order scores of the ablation study on the different-size tasks presented in Section 5.3 ASO scores are expressed mmnn with significance level a 0.05 that is adjusted accordingly by using the Bonferroni correction [17]. Read from row to column: e.g., NeuralSoll Solver (row) stochastically dominant over No LSTM (column) in the 1S-Maze task with emin of 0.00.
24

---

Perdirewert
1.00
100
1,00
1000
FandForward
1.00
1.00
1.00
1.00
FeedForeadd
1.40
1.00
000
1,00
1.00
0.00
1.00
0.35
1.00
1.00
-0.
NezratSoleve
1.00
100
NoursSScheer
1.00
1.00
NeummiSotoer
0.00
1.00
Deacle
0.06
0.02
0.06
1.00
Cracie
0.00
0.02
0.08
1.00
Crricee
0.00
0.00
0.38
1,00
Banno weaoteeaa
Onade
RodFoowc
Neund
Peeffooment
3.0
PuotPurward
1.00
1.20
1.00
1.00
PeodPooward
1,00
1.00
1.00
1.00
0.6
2.4
NeuraSSbberr
NeuralSolove
Oradle
0.00
0.20
0.22
1.00
Oracle
0,00
0.00
0.00
1.00
0.0
Oheck
PadPormard
Osuca
Rosst
Figure 18: Almost Stochastic Order scores of the Doorkey sequential decision crrisiionmmaiing task results
presented in Section 5.4 ASO scores are expressed <mnn with significance level = 0.05 that is adjusted accordingly by using the Bonferroni correction [17
Read from row to column: e.g.. (column) starting scenarios
NeuralSolvve (row) stochastically dominant over Bansal
6
with observations of size 32 with an Cmin of 0.00.
Pong
1,00
1.00
1.00
1.00
0.71
1.00
1.00
0.66
1.00
1.00
1.00
1.00
1.00
0.65
0.34
1.00
0.70
1.0
0.50
0.23
1.00
0.80
0.06
0.02
0.50
1.00
1.00
1.00
1.00
1.00
0.50
1.00
1.00
0.49
1.00
1.00
1.80
0.60
1.00
0.10
1.00
1.00
1.00
1.00
1.00
1.00
0.01
0.90
0.30
1.00
1.00
1.00
0.01
1.00
0.58
0.28
1.00
0.58
0.44
1.00
0.00
0.35
1.00
0.00
0.38
0.0
0.50
0.01
0.50
0.00
0.50
0.10
0.00
Doorkey
Thin-Maze
1.00
1.00
0.75
0.59
1.00
0:76
1,0
1.00
0.10
1.00
1.00
1.00
1.0
6.8
0.30
1.00
1.00
1.00
1.00
1.00
0.50
1.00
1.00
1.00
1.00
1.00
0.10
1.00
0.10
0.01
0.09
0.21
0.07
1.00
0.35
0.01
0.00
0.43
1.00
1.00
6.2
0.00
1.00
0.88
1.00
1.00
1.00
0.00
0.01
0.00
0.39
0.76
1.00
0.0
1.00
0.50
0.30
0.01
1.00
0.50
0.10
0.01
Figure 19: Almost Stochastic Order scores of the fine tuning progressive loss results presented in Appendix C.3] ASO scores are expressed in Gmin- with significance level a 0.05 that is adjusted accordingly by using the Bonferroni correction [17]. Read from row to column: e.g., a progressive loss alpha value 0.5 (row) almost stochastically dominant over the value of 1. (column) in the 1S M-Maz task with émin of 0.23.
25

---

Iarabion:
Iarabonn
Eandber:30
Tarationc 50
Ieretion 00
Earobarr:00
Bersbonn 90
Bersiion 300
Berol ART
Figure 20: Propagation of information in NeuralSolver for the -a-aa task: Top: we highlight the difference between the current iteration and last iteration (not represented) of the internal state of the recurrent module of our model positions of the recurrent state that converged to fixed final state. Larger differences are represented in deep blue color. Bottom: the predicted action probabilities by the model at different iterations, where the agent can move right (R), down (D), left (L), up (U).
x 35 environment. The white pixels represent the pixel
E Learned Algorithms with NeuralSolver
Similar Section 3.2 show visualizations of the recurrent space of NeuralSolver while solving the different-size tasks. More specifically, we visualize the difference between the current recurrent state of the LSTM and the last recurrent state, for each iteration. The following figures show two plots per iteration: the one on top shows the difference between the recurrent states in a certain color superimposed
the input image, while the one below shows the predicted action probabilities outputted by the model.
With these visualizations giving a hint of the algorithm that the model performing. An interesting observation across these visualizations is how the final prediction of the model only converges to the correct action after the recurrent iterations have converged to fixed state around the player position.
notice how the information propagates through the input problems,
E.1
1S-Maze
By looking at how the information propagates through the maze paths in Figure 20 and 21 the model appears learn a parallel algorithm that starts with the dead ends until finding the optimal path between the player (green) and the goal (red). At iteration #90 in Figure 20, the visible difference between the current and last recurrent states actually represents and the goal. The algorithm converges to the optimal task after the difference of the state converges near the player. Figure 2T shows that finding the optimal path is not necessary to predict the next action.
optimal path between the player
26

---

Itaration:
Laration::
Fvers 300
Tteration 350
Faers 200
Elerstion 2010
Poarol $000
Figure 21: Propagation of information in NeuralSolver for the 1s Maze task: Top: we highlight the difference between the current iteration and last iteration (not represented) of the internal state of the recurrent module of our model
x environment. The white pixels represent the pixel positions of the recurrent state that converged to fixed final state. Larger differences are represented in deep blue color. Bottom: the predicted action probabilities by the model at different iterations, where the agent can move right (R), down (D), left (L), up (U).
Itarotio:
Earotior: 30
Bar Merstom::22 20
handbir:: 40
Baretion: 60
Figure 22: Propagation of information in NeuralSolver for the GoTo task: Top: we highlight the difference between the current iteration and last iteration (not represented) of the internal state of the recurrent module of our model in 64 64 environment. The black pixels represent the pixel positions of the recurrent state that converged to fixed final state. Larger differences are represented in deep blue color. Bottom: the predicted action probabilities by the model at different iterations, where the agent can move right (R), down (D), left (L), up (U).
E.2 GoTo
We highlight three examples of possible player and goal positions that influence the final predicted action. Namely, if the player (green) is above the goal (red), below the goal, or at the same level in height.
In Figure 22 we can notice that the model chooses by default the up action, while across the iterations signal sent from the goal position to the positions above, propagating in cireular/oval shape. Thus the model appears to learn an algorithm that attempts to signal the player if it should go down. Similarly, in Figure the recurrent state does not change near the player, as such the predicted action remains the same.
In Figure 24 we see a case in which the player is the left of the goal. In this setting, the horizontal line with slightly different contrast ranging from the goal position appears to communicate across that line the player should go left or right to reach the goal.
27

---

Iaraaion:
Fandior::
tarabton:
Earutior: 40
Iaration:
Propagation of information in NeuralSolver for the GoTo task: Top: we highlight the difference between the current iteration and last iteration (not represented) of the internal state of the recurrent module of our model in 64 x 64 environment. The black pixels represent the pixel positions of the recurrent state that converged to fixed final state. Larger differences are represented in deep blue color. Bottom: the predicted action probabilities by the model at different iterations, where the agent can move right (R), down (D), left (L), up (U).
Figure 23:
baretion:
Ioretion:
Figure 24: Propagation of information in Neural for the GoTo task: Top: we highlight the difference between the current iteration and last iteration (not represented) of the internal state of the recurrent module of our model in 64 x 64 environment. The black pixels represent the pixel positions of the recurrent state that converged to a fixed final state. Larger differences are represented in deep blue color. Bottom: the predicted action probabilities by the model different iterations, where the agent can move right (R), down (D), left (L), or up (U).
E.3 Pong
We present two examples our learned algorithms solving the simplified game of Pong: one when the ball is on the left of the paddle, and the other when on the right. We can notice that the model learns similar algorithm to the model in the GoTo task. In both cases the model starts by predicting that the paddle should move right. In Figure 25 the model prediction changes once the information propagation reaches the paddle, changing it the left action. In Figure 26,, see that this propagation of information does not reach the paddle, and as such the action remains unchanged.
We can also see line with slightly different contrast leaving the center of the ball, thus signaling that the model should instead predict the stay action. This could indicate that different contrasts might represent different algorithms running simultancously.
28

---

Iaanaion:
Fandiar::
Earutirr: 30
Taration:
Ieretion: 30
harotior:
Ieraaion: P0
handtier::0
Ieration 90
Figure 25: Propagation of information in NeuralSolver for the Pong task: Top: we highlight the difference between the current iteration and last iteration (not represented) of the internal state of the recurrent module of our model positions of the recurrent state that converged to fixed final state. Larger differences are represented in deep blue color. Bottom: the predicted action probabilities by the model at different iterations, where the agent can move left (L), stay (S) and right (R).
64 environment. The black pixels represent the pixel
E.4 Doorkey
We provide two examples of learned algorithms in Doorkey: in Figure we show case when the player is almost reaching the goal, and in Figure 28 case when it already captured the key and is reaching for the door.
These visualizations are harder to interpret than the other ones. This phenomenon might occur since, to solve this task, the agent needs to follow complex sequence of tasks, thus requiring multiple algorithms to find where the positions of the objects are and which action to do next. We can notice this in Figure 28 where there is line that shoots vertically from the door position that is present in Figure 27 This shows that for more complex tasks, better visualization and interpretation techniques of the recurrent states required.
29

---

ltaratio::
Eerutier:
Heration
beratior:
Ieration:
Figure 26: Propagation of information in Neural for the Pong task: Top: we highlight the difference between the current iteration and last iteration (not represented) of the internal state of the recurrent module of our model in 64 x 64 environment. The black pixels represent the pixel positions of the recurrent state that converged to fixed final state. Larger differences are represented in deep blue color. Bottom: the predicted action probabilities by the model at different iterations, where the agent can move left (L), stay (S) and right (R).
terabion:
baratio::
Buration:
Figure 27: Propagation of information in Neural for the DoorKey task: Top: we highlight the difference between the current iteration and last iteration (not represented) of the internal state of the recurrent module of our model in 64 x 64 environment. The black pixels represent the pixel positions of the recurrent state that converged to fixed final state. Larger differences are represented in deep blue color. Bottom: the predicted action probabilities by the model at different iterations, namely Forward (L), Rotate Right (R), Pickup (P), and Toggle (T).
30

---

llaration:
Fanatior::
Eoretio:600
Roratior:80
Figure 28: Propagation of information in NeuralSolven for the DoorKey task: Top: we highlight the difference between the current iteration and last iteration (not represented) of the internal state of the recurrent module of our model in 64 x 64 environment. The black pixels represent the pixel positions of the recurrent state that converged to fixed final state. Larger differences are represented in deep blue color. Bottom: the predicted action probabilities by the model at different iterations, namely Forward (L.), Rotate Right (R), Pickup (P), and Toggle (T).
31

---

Additional Example Trajectories
In Figure 29 present additional trajectories of agents that execute policies learned using
NeuralSolver and other baselines, following the same procedure as in Section 5.4
Oracle
NeuralSolver
Bansal et FeedForward
Figure 29: Visualization of the trajectories of TeuralSolver Bansal et al. [6], and FeedForward models against the Oracle trajectory. The first row is task of size 32 32, second and third of 64x64, and fourth and fifth of 128x 128.
32

---

G Examples of tasks with observations of size 512 x 512
We present additional visualizations of examples of tasks with large observations (512 512).
Figure 30: Example of a -S-aa task with 512 x 512 observations.
33

---

Figure 31: Example of Doorkey task with 512 x 512 observations.
34

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our claims are properly justified through the paper, with experimental support in Section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Even though we did not provide a specific section to discuss the limitations, through the article our experiments show possible improvements the extrapolating from smaller training sizes, Section 5.1 and Section 5.2 as well slightly worse performance when training on fewer training examples, Section
Justification:
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
35

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We include the appendix all the details necessary to successfully implement the NeuralSolver architecture and different-size tasks, even without the provided code. In Section B.3 we provided the hyperparameters used in the training of all tasks and in Sec- tion B.6 we highlight the main differences between the Bansal al. [6] and NeuralSolwe implementations that allow also to faithfully implement our model based on the available Bansal al. [6] descriptions and available code. In Section A we give the details how to implement the different-size tasks.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
(b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
36

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide a GitHub repository with the code, the data, and a Readme file detailing how perform the main experimental results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ //nips. http//uublcggguueeeees http://///uddee//eeeuuubbiiii /at////ooeSSsiiseeeeeeeeeeeeeeee ississsinn ssionPolicy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide complete details of the training and test details in Appendix B
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: All our experiments are accompanied with standard deviation errors and with the ASO [6] significance test described in Section
Guidelines:
37

---

The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errоrs is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Section B.4 and Section
B.5
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics httttps://neurip..ooo httparrss....pppp.a./hhhcccceela
Answer: [Yes]
Justification: Our research does not involve human participants. As this is an empirical work that studies the extrapolation abilities deep learning models when learning algorithms we do not have direct negative society impacts. Our training tasks and use of previous datasets do not contain any personally identifiable information, do not use or refer to any sort real people information, well respect the previous task licenses.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
38

---

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: As this is an empirical work that studies the extrapolation abilities of deep learning models when learning algorithms we do not have direct negative society impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups). privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Our models and tasks are based on previous work that are all properly cited and referenced.
Guidelines:
39

---

The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperssi thcode. com/ com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: In Section give the details how implement the different-size tasks. We also provide the code and description to create these tasks.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/code/moodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
40

---

Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
41


Reviews:
[{"review_id": "VpPhPSjFBj", "reviewer": "Reviewer_tqxp", "summary": "The authors propose a new architecture that improves on the Deep Thinking (DT) architecture by Bansal et al, 2022. It replaces the recurrent ResNet block with a convolutional LSTM+layernorm. They also use a curriculum-based learning schema. This improves the model's extrapolation capability significantly. The authors test their model on identical tasks to those of Bansal et al, and show that they outperform the DT model on all of them, often significantly.", "questions": "It would be nice to see an ablation on the layernorms in Tab. 4.\n\nIn line 168, the authors write \"consider the best accuracy obtained by the models at any iteration\". Is this also done for NeuralThink? It should not suffer from overthinking, right?\n\nWhy is think maze more difficult than the thick one, especially for the baselines?", "limitations": "It is not clear how these architectures can transfer from simple grid world problems to some more real-world applications. The authors could discuss this more in the paper.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The authors propose a strong model that outperforms the baseline in all tasks they test for. Their model is more general than the baseline, allowing generalization also in tasks whose output grid is of a different size compared to the input. The model doesn't suffer from overthinking and does not require special regularization.", "weaknesses": "The writing can be improved: the authors do not clearly state what is novel in their architecture in the main text. They should emphasize the exact differences to the literature more clearly: e.g by stating that the difference in their architecture and DT is replacing the ResNet by a convolutional LSTM.\n\nFigures with the double meaning for the color are confusing (Fig 5 and Fig. 6). Also, for Fig 6, there are 3 options but just 2 numbers for each color. \n\nThe exact details of the architecture are not entirely clear. It would be better if the authors could describe their architecture with equations, at least in the appendix.\n\nThe authors are using LSTMs but do not cite [1].\n\n[1] Hochreiter et al, 1997:  Long Short-Term Memory", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments and for the time spent reviewing our paper. We now address the weaknesses (W), questions (Q) and limitations (L) pointed out by the reviewer:\n\n+ **(W1) Authors do not clearly state what is novel in their architecture in the main text.** Due to lack of space in the main body, we have an extended discussion on the differences between NeuralThink and DeepThink in Appendix B.6. Nonetheless, we have improved the description of the novelty of our model in line 100 of the updated version of the paper. \n+ **(W2) Color labels of  Fig. 5 and Fig. 6 are confusing:** To improve our figures we modified the labels of the colors to “very small”, “small”, “medium” and “large”, and defined these training sizes for each task in Appendix.  \n+ **(W3) The exact details of the architecture are not entirely clear:** We can describe our architecture as follows: for same size tasks, given an input image $o$, we pass the input through a single layer ConvLSTM ($q$) to obtain the next hidden state ($h\\_1$ and cell state ($c\\_1$), $h\\_1, c\\_1 \\= q(o,h\\_0,c\\_0) $. The initial hidden ($h\\_0$) and cell ($c\\_0$) states are initialized at zeros. The recurrent process is given by feeding back the previous hidden and cell states to the LSTM, $h\\_t, c\\_t \\= q(o,h\\_{t-1},c\\_{t-1})$. After a fixed size of recurrent iterations, we can obtain the model prediction at timestep $t$ by passing the hidden state $h\\_t$ through the processing module. The processing module is composed of three convolutional layers with 3x3 kernels denoted by $W\\_1,W\\_2,W\\_3$ and ReLU activation ($\\\\gamma$). As such, we perform $p\\_t \\= W\\_3 \\\\ast \\\\gamma(W\\_2 \\\\ast \\\\gamma(W\\_1 \\\\ast h\\_t)) $, and a final softmax activation ($\\\\sigma$), $\\\\hat{y}\\_t \\= \\\\sigma(p\\_t)$. For different size tasks, before the softmax we perform a global maxpool operation ($G$) that reduces the processing output height and width to 1, $\\\\hat{y}\\_t \\= \\\\sigma(G(p\\_t))$. In the updated version of the paper we added this description to Appendix B.6.  \n+ **(W4) The authors use LSTMs but do not cite the original article:** We have updated the paper to include the citation of the original LSTM paper.  \n+ **(Q1) It would be nice to see an ablation on the layernorms in Tab. 4\\.** In the original version of the paper we included the ablation on layernorm in Table 15 of Appendix C6, due to lack of space in the main body of the paper. In the updated version of the paper we have emphasized the existence of additional ablation studies available in Appendix, in the beginning of Section 5.3.  \n+ **(Q2a) Is the best accuracy obtained by the models at any iteration used in NeuralThink?** Yes, we use this evaluation methodology for all methods.  \n+ **(Q2b) Does NeuralThink suffer from overthinking?** In Appendix B.1. we empirically show that NeuralThink does not suffer from overthinking, across all evaluation tasks. We have updated the text in line 213 to clarify this point.  \n+ **(Q3) Why is the thin maze more difficult than the thick one, especially for the baselines?** We believe that the difficulty arises from the size of the receptor field in the convolutional layers required to perceive the environment. The original DeepThink architecture uses 5 convolutional layers with a kernel size of 3 for its recurrent module. This means that at each recurrent step, the network is capturing  information from 5 adjacent pixels in each direction (i.e., a square of 11x11 pixels). This fact makes it harder for the model to generalize from smaller problems as the model has to learn an iterative algorithm from a  square of 11 pixels at each iterative step. On the other hand, the ConvLSTM used by NeuralThink only captures information from the next adjacent input pixel (or squares 3x3 pixels). This design choice allows NeuralThink to learn general iterative algorithms from smaller sizes.  \n+ **(L1) It is not clear how NeuralThink can work on real-world applications:** We have evaluated NeuralThink on both existing benchmarks of DeepThinking models as well as on novel benchmark scenarios that we contribute for different-sized tasks. We have also explored the use of our model in sequential decision-making tasks, such as the DoorKey environment of the MiniGrid suite in Section 5.4, which is still a standard environment for the evaluation of modern reinforcement learning agents (see \\[1a-1b\\]).\n\nReferences:  \n\\[1a\\] Nikulin, Alexander, et al. \"XLand-minigrid: Scalable meta-reinforcement learning environments in JAX.\" *arXiv preprint arXiv:2312.12044* (2023).  \n\\[1b\\] Pignatelli, Eduardo, et al. \"NAVIX: Scaling MiniGrid Environments with JAX.\" *arXiv preprint arXiv:2407.19396* (2024).", "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": "HyARaqRLNX", "reviewer": "Reviewer_SKBz", "summary": "1. The paper introduces NeuralThink, a novel deep thinking architecture designed to efficiently and consistently extrapolate learned algorithms from smaller problems to larger ones.\n2. Unlike previous deep thinking methods, NeuralThink can be applied to both same-size problems (where input and output sizes are the same) and different-size problems (where input and output sizes differ).\n3. The architecture consists of three main components:\n   - A recurrent module that iteratively processes input information at different scales\n   - A processing module that aggregates previously processed information\n   - A curriculum-based training scheme to improve extrapolation performance\n4. NeuralThink outperforms prior state-of-the-art deep thinking approaches in:\n   - Extrapolation capability: Consistently executing learned algorithms on larger problems\n   - Training efficiency: Learning algorithms from smaller problems\n   - Parameter efficiency: Requiring fewer parameters than other approaches\n5. The authors introduce a set of novel different-size tasks to evaluate their method.", "questions": "See Weaknesses. \nCan you provide a more detail comparison between DeepThink and your model. What happens generally for both of these models if you scale them? It would be nice to see when NeuralThink breaks? Afaiu DeepThink was going for the recurrent module to see if larger inputs could be solved by running the network longer in depth. Your model works quite differently afaik and I would again highlight these differences in detail.", "limitations": "I would welcome a dedicated limitation section.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is easy to follow and experiments nicely presented. The proposed model seems to outperform (one) related work method and results seem quite impressive.", "weaknesses": "Since the authors introduce novel benchmarks, the very positive experimental results need more ablations. Please provide hyperparameter scan and their results on baselines.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We express our gratitude to the reviewer for their comments and the effort spent on reviewing our paper. We will now tackle the weaknesses (W), questions (Q), and limitations (L) mentioned:\n\n+ **(W1a) Ablations present in the paper:** We provide in the original paper additional ablation studies in Appendix C, due to the lack of space in the main body of the text. In particular we show:  \n  + Ablation on the progressive loss parameter of DeepThink: Appendix C.3.  \n  + Ablation on the model size of NeuralThink: Appendix C.4.  \n  + Ablation on different types of recurrent networks for NeuralThinkL Appendix C.5.  \n  + Ablation on layernorm, dropout and projection head: Appendix C.6. \n\n  In the updated version of the paper (beginning of Section 5.3) we have emphasized the existence of these additional ablation studies.\n\n+ **(W1b) Hyperparameter scan details on baselines**: In the general rebuttal comment document we show the results of a hyperparameter scan in the GoTo environment for the NeuralThink training hyperparameters. The results show that NeuralThink is robust to hyperparameter changes, with exception of the gradient clip that requires proper tuning. The results further show the significance of using weight decay and both forms of dropout regularization together. In this work we tried to use our model with minimum hyperparameter tuning, keeping the same hyperparameters for the majority of the different-sized environments and using the same hyperparameters for the same-sized tasks proposed by the DeepThink model, as shown in Appendix B.3.  \n+ **(Q1, Q4) Comparison between DeepThink and our model:** We provide in Appendix B.6. of the original paper an extended discussion on the differences between our model and DeepThink. We have clarified the existence of this comparison in line 129 of the updated version of the paper. Moreover, we also highlight in Appendix C.7. how our model outperforms DeepThink on extrapolation to very-large input sizes (256x256 and 512x512).  \n+ **(Q2) What happens if you scale the models:** We provide in Appendix C.4. of the original paper an ablation study on the effect of the model size for the performance of NeuralThink. The results show that decreasing the model size results in a decrease in the performance of NeuralThink, yet still outperforms the previous state-of-the-art. In the updated version of the paper (beginning of Section 5.3) we have emphasized the existence of this ablation study.  \n+ **(Q3) When does NeuralThink break?** In Figures 5 and 6 we highlight that our model struggles to extrapolate when we significantly reduce the size of the input data provided to the model during training. We believe that learning to extrapolate from very small input sizes is a very interesting direction for future work (as mentioned in line 282).  \n+ **(L1) Dedicated limitations section:** Due to lack of space in the main body, we have created a dedicated Appendix section for limitations, merging it with the previous Appendix B.2. “Changes We Tried That Did Not Improve NeuralThink” and mentioned the existence of this Appendix in line 282 of the updated version of the paper.", "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": "U6lgvGGXm4", "reviewer": "Reviewer_dmPv", "summary": "The paper proposes a method called \"NeuralThink\", which is designed to improve the same-size task generalization and different-size task extrapolation performances. The proposed algorithm composes of three components: (1) A recurrent module utilizing a LSTM network to process inputs of different scale, (2) A processing module serving as the aggregation layer, and (3) A curriculum learning training scheme  that gradually increases the dimensionality of the observation. Experiments on an algorithm learning benchmark showed that the proposed approach surpassed the baselines being compared, especially in terms of the different-size task extrapolation performance. It is also demonstrated that the proposed approach has a higher efficiency both in terms of training size and parameter size.", "questions": "See above.", "limitations": "N/A", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The paper addresses a meaningful problem space: Algorithm learning and extrapolation ability. The proposed method achieves oracle-level performance on different-size tasks, showcasing the improved generalization and extrapolation power. \n- The proposed approach outperforms the selected baselines on an algorithm learning benchmark, proving the effectiveness to a certain extent. Further experiments also indicates improved training size and parameter efficiency. The ablation study demonstrates that all components in the proposed approach contribute positively to the final result.\n- The intro section is well-written. Figure 1-3 are well-illustrated.", "weaknesses": "- It's not immediately clear why the proposed approach achieves improved performance above prior arts. Explanation on (1) the fundamental difference compared to the previous approach, and (2) why this should work better would be helpful.\n- The result being presented in the experiment section is a bit counter-intuitive. Table 3 shows that the proposed approach achieves oracle-level performance on different-size task extrapolation, yet Table 1 shows that the performance on same-size tasks is imperfect. Intuitively one would imagine different-size tasks to be more challenging for the algorithm, yet the result suggests otherwise. An analysis on the intuition behind would be helpful. Besides, more case studies to compare the performance gap between the proposed approach v.s. prior baselines would help solidify the understanding. \n- The number of baselines being compared to is limited. It would be helpful to know how does the proposed approach compare to other recently published methods, e.g. [1], or if they are not directly comparable, what is the reason behind.\n- nit: The presentation of some figures need to be improved for better clarity. e.g. Figure 5 it's not immediately clear what does the notion \"x/y\" mean in the legend of the bar chart. \n\n[1] Adaptive recurrent vision performs zero-shot computation scaling to unseen difficulty levels. NeurIPS 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful comments and the time spent reviewing our paper. We now address the weaknesses (W) identified by the reviewer:\n\n+ **(W1) Differences between NeuralThink and the previous work/why NeuralThink achieves improved performance.** We would like to highlight the following fundamental differences between our work and previous work (and why):  \n  1. **Recurrent design:** Previous work uses the simplest form of recurrence by using a ResNet with weight sharing across residual blocks. In our work, we started by evaluating the extrapolation capabilities of DeepThink and found its performance to be lacking in different-sized tasks (as we show in Table 3). To overcome this gap, we explored other more advanced recurrent designs for algorithmic extrapolation, in particular LSTM-based convolutional designs. Our results show that this choice of recurrent design results in better performance of the method, supporting prior work \\[3\\].  \n  1. **Aggregation layer in different size tasks:** Previous DeepThinking methods do not work in this class of problems, as their models only support same size input and output problems. To overcome this issue, we introduced the use of a pooling layer in the design of our model.   \n  1. **Curriculum-based training scheme:** The previous differences are insufficient to allow consistent extrapolation in different-size tasks. To the best of our knowledge, we are the first work that introduces a curriculum-based training scheme for DeepThinking architectures. In Table 4 we highlight the role of our curriculum-based training scheme in the overall performance of the method.\n\n   We provide in Appendix B.6. of the original paper an extended discussion on the differences between our model and DeepThink. We have clarified the existence of this comparison in line 129 of the updated version of the paper. \n\n+ **(W2) Imperfect results in same-sized tasks conflict with oracle-level performance on different-size task extrapolation:** We would like to point out that, while the result for Thin-Maze is not at oracle-level accuracy, the results are still inside of the experimental uncertainty of the other results. On the other hand, the Chess task is fundamentally different from the other tasks, since it is not a spatial extrapolation task but a symbolic extrapolation level, proving to be a harder task for our method. However, despite not achieving oracle-level performance, even in this challenging task our model is able to outperform the previous state-of-the-art.\n\n+ **(W2,1) The number of used tasks is small:** We would like to point out that not only are we using all the benchmarks proposed by \\[2\\] to evaluate DeepThinking models, but we also contribute a novel set of different-sized tasks, due to the absence of an existing adequate benchmark for this type of tasks. Furthermore, we are the first work to explore sequential decision-making scenarios for the application of DeepThinking models.\n\n+ **(W3) Limited Baselines/AdRNNs**: We have updated our related section with \\[1\\]. We have compared directly and indirectly the performance of AdRNNs to NeuralThink.  \n  + **AdRNNs have a worse performance than DeepThink (our baseline) on same-sized tasks.** \\[1\\] uses the same training methodology as NeuralThink on the same-sized Maze tasks (train on mazes of size 24x24 pixels) but test on much smaller mazes (test on mazes of size 44x44 and 56x56 pixels) than in our paper. The best method presented in \\[1\\] (LocRNN) achieves 50% (best performance) in test scenarios that are half the size of the ones presented in our work: DeepThink achieves 91% test accuracy on mazes of size 124x124. As such, DeepThink is a stronger baseline than AdRNNs for algorithmic extrapolation and, thus, we outperform them as well.  \n  + **AdRNNs have a worse performance than NeuralThink on different-sized tasks.** Despite the fact that no public code implementation of \\[1\\] is available, we implemented the proposed LocRNN without the learnable iteration-halting mechanism (as we always select the best performing iteration for the results) and compared its performance against the ConvLSTM used in NeuralThink on the different-sized tasks. The results are shown in the below table, which highlights that LocRNN has a worse performance than our method, while having more parameters than NeuralThink and being more computationally intensive to execute. We added this additional comparison to Appendix C.5.\n\n| Extrapolation accuracy (%) | 1S-Maze | GoTo | Pong | Doorkey |\n| :---- | :---- | :---- | :---- | :---- |\n| NeuralThink | 100.00 \\+- 0.00  | 100.00 \\+- 0.00  | 100.00 \\+- 0.00  | 100.00 \\+- 0.00  |\n| LocRNN \\[1\\] | 87.65 \\+- 9.13 | 82.56 \\+- 17.73 | 94.08 \\+- 8.81 | 86.02 \\+- 11.75 |\n\n\n| Number of parameters (in Millions) | 1S-Maze | GoTo | Pong | Doorkey |\n| :---- | :---- | :---- | :---- | :---- |\n| NeuralThink | 0.231 | 0.231 | 0.230 | 0.231 |\n| LocRNN \\[1\\] | 0.236 | 0.236 | 0.236 | 0.236 |\n\n| Computational complexity (in gigaMACs) | 1S-Maze | GoTo | Pong | Doorkey |\n| :---- | :---- | :---- | :---- | :---- |\n| NeuralThink | 4.12 | 9.76 | 9.76 | 9.76 |\n| LocRNN \\[1\\] | 4.33 | 10.25 | 10.25 | 10.25 |\n \n+ **(W4) “x/y” color labels in Figure 5 are not clear.**  \n  We changed the labels of the colors with very small, small, medium and large, and moved details of the training sizes used for each task to the Appendix.\n\n\\[1\\] Adaptive recurrent vision performs zero-shot computation scaling to unseen difficulty levels. NeurIPS 2023\\.  \n\\[2\\] Bansal, Arpit, et al. \"End-to-end algorithm synthesis with recurrent networks: Extrapolation without overthinking.\" Advances in Neural Information Processing Systems 35 (2022)  \n\\[3\\] Eric Price, Wojciech Zaremba, and Ilya Sutskever. Extensions and limitations of the neural GPU. CoRR, 337 abs/1611.00736, 2016\\. doi: 10.48550/arxiv.1611.00736.", "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:
Dear Reviewers,

We thank all the reviewers for the constructive and interesting questions and suggestions. 

We have added a pdf with the additional hyperparameter scan requested by Reviewer **SKBz**, that highlights the robustness of our method to changes in hyperparameters. 

Additionally, we replicate here the results of the comparison between our method and LocRNN \[1\], requested by Reviewer **dmPv**, that shows that our method significantly outperforms LocRNN in both extrapolation performance and efficiency.

| Extrapolation accuracy (%) | 1S-Maze | GoTo | Pong | Doorkey |
| :---- | :---- | :---- | :---- | :---- |
| NeuralThink | 100.00 \+- 0.00  | 100.00 \+- 0.00  | 100.00 \+- 0.00  | 100.00 \+- 0.00  |
| LocRNN \[1\] | 87.65 \+- 9.13 | 82.56 \+- 17.73 | 94.08 \+- 8.81 | 86.02 \+- 11.75 |

| Number of parameters (in Millions) | 1S-Maze | GoTo | Pong | Doorkey |
| :---- | :---- | :---- | :---- | :---- |
| NeuralThink | 0.231 | 0.231 | 0.230 | 0.231 |
| LocRNN \[1\] | 0.236 | 0.236 | 0.236 | 0.236 |

| Computational complexity (in gigaMACs) | 1S-Maze | GoTo | Pong | Doorkey |
| :---- | :---- | :---- | :---- | :---- |
| NeuralThink | 4.12 | 9.76 | 9.76 | 9.76 |
| LocRNN \[1\] | 4.33 | 10.25 | 10.25 | 10.25 |

   
Please let us know if our comments address all the weaknesses and questions pointed out by the reviews or if you require further clarification.

\[1\] Adaptive recurrent vision performs zero-shot computation scaling to unseen difficulty levels. NeurIPS 2023\.

Meta-review: This paper introduces a deep learning approach that generalizes across reasoning-related tasks of different sizes, including asymmetrical tasks. The reviewers agree this is a useful contribution and should be accepted. I accordingly recommendation acceptance. However, as a requirement for acceptance, the term "deep thinking", which represents a gross anthropomorphism of deep learning, **must be removed from this paper in every place it occurs** and both DeepThink and NeuralThink **must be retitled** (the AC will check the camera-ready to make sure that these edits are made). The term "overthinking" is permissible since it is a term for which there is no easy synonym, but other uses of "think" should be eliminated. While the term "deep thinking" was used sparingly in a prior paper, that is not an excuse to perpetuate it beyond acknowledging the prior paper - there are plenty of other terminology options that can be introduced without violating best practices in deep learning communication.
Decision: Accept (poster)