Title: NeuralSolver: Learning Algorithms For Consistent and Efficient Extrapolation Across General Tasks
Venue: NeurIPS 2024
Accepted: yes
Keywords: Deep learning

algorithm synthesis

recurrent networks

algorithmic reasoning

sequential decision making

extrapolation
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

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 differentsize 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-ofthe-art recurrent solvers in extrapolating to larger problems, considering smaller training problems and requiring less parameters than other approaches. Code available at https://github.com/esteveste/NeuralSolver

1 Introduction

Humans can solve complex reasoning tasks by extrapolating : employing and combining elementary logical components to build more elaborate strategies. Machine learning models excel at pattern recognition, often outperforming humans in classification [1, 2], control [3] and prediction tasks [4]. However, these models still struggle at reasoning , which affects their ability to 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 a maze (Figure 1), becomes increasingly harder for larger problems.

We focus on designing learning algorithms that are able to efficiently extrapolate, i.e., are able to learn algorithms to solve tasks over small problems (in terms of 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, 6]. Compared to 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 recurrent 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 in a wide range of tasks, which we call different-size tasks, where the input and output dimensions are different (e.g., classification and decision-making tasks).

Instituto Superior Técnico, Universidade de Lisboa

25

We address this limitation and propose a novel recurrent solver that is able to consistently and efficiently extrapolate regardless of the size of the output and the type of problem. We name our approach NeuralSolver , a 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 in the input image, and a 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 NeuralSolver on small observations (e.g., 15 × 15 images) using standard supervised learning techniques and apply the learned algorithm at test time with arbitrarily large observations (e.g., 256 × 256 images, or larger), with minimal loss in performance.

We evaluate NeuralSolver against prior recurrent solver in literature-standard same-size tasks and in a novel set of different-size tasks. We show across all tasks that our model significantly outperforms previous approaches in extrapolation capability , being able to 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 a novel recurrent solver that uses a recurrent convolutional module, a processing module (with an aggregation function) and curriculum-learning to train algorithms in 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 outperforms the prior recurrent solvers regarding extrapolation capabilities, training efficiency and parameter efficiency.

2 Related Work

Schwarzschild et al. [5] introduce a class of recurrent neural networks that can execute at test time more iterations than the number of iterations used during training, to solve problems more complex than the ones seen during training, which we denote by recurrent-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 al. [6] introduces two new components to the previous architecture: (1) a 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) a 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

Observation

Recurrent

Module

Processing

Module

Output

Same-size

tasks

Observation

Recurrent

Module

Processing

Module

Output

Different-size

tasks

A

0.1

0.9

0.0

0.0

more than 95% accuracy on the same-size benchmark. We use this method as a baseline against our method. To address the problems of the computational efficiency and hyperparameter tuning of the progressive loss, [7] propose replacing the progressive loss with a 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 to counter the overthinking problem (as shown in Appendix B.1). NeuralSolver is the first recurrent solver able to consistently and efficiently extrapolate learned algorithms in both same-size and different-size tasks.

3 The NeuralSolver Architecture

We address the challenge of designing a model that is able to efficiently learn algorithms from small problems that consistently extrapolate to larger problems. Contrary to previous recurrent solvers [57] we want to apply our 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, we contribute NeuralSolver , depicted in Figure 2.

3.1 Model Architecture

We design our model with three components: (i) we extract information from the input data using a recurrent module , responsible for iteratively processing the input observation at different scales, allowing our model to cope with arbitrarily large input data; (ii) information is then sent to a processing module , responsible for aggregating (if necessary) the processed information and generating the output of the network, allowing our model to be used in both same-size and different-size tasks; (iii) to train our model we employ a curriculum-based scheme, that improves the extrapolation performance of NeuralSolver to larger problems.

Recurrent Module : This module consists of a layer normalized convolutional LSTM [9-11], a LSTM that uses convolutional layers instead of linear layers in its 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 to perform multiple iterations over the data, allowing our model to process arbitrarily large observations. We show in Appendix B.1 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 of different choices of recurrent architectures in the performance of our model.

Processing Module : This module consists of a block of three convolutional layers (following [6]) whose input is the last hidden state of the recurrent module. Their purpose is to reduce the number of channels to a desired output number of channels. Additionally, we have an optional aggregation layer (in our case a global max-pooling layer) that can be employed in different-size tasks to reduce the variable input size of the network to a fixed output size. For example, for a two-dimensional input observation, we first reduce the input to a 1 × 1 × o tensor, where o is the desired output number of channels, and subsequently flatten the output tensor. This design allows our architecture to be used

配配配配配

iteration:0

iteration:10

iteration:15

iteration:20

R

DL

U

R

R

DLU

R

D

LU

RD

LU

both in same-size and different-size tasks, with minimum changes. In Section 5.3 we compare the performance of different aggregation layers.

Curriculum Learning Training : We train our model on a set of smaller-size observations and test the performance of the model with larger dimensional observations (all belonging to the same task). Additionally, we employ a curriculum learning approach in different-size tasks to counter the effect of the reduced training signal (due to the smaller output dimensionality) [12, 13]: we initially train the models on lower-dimensionality observations, and then gradually increase the dimensionality of the observations every N epochs. To reduce the risk of catastrophic forgetting, we sample a minibatch of observations with a 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-based 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 a maze-like environment at 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 of 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 a 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 convolution : for example, a 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 a convolutional recurrent neural networks we 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 of the correct label (in this case, "Down"). We provide more examples of learned algorithms in Appendix E.

4 Evaluation

We evaluate NeuralSolver to demonstrate how it ourperforms previous state-of-the-art recurrent solvers accordingly to three main criteria: (i) extrapolation capability , the ability to execute learned

(a) GoTo

(b) 1S-Maze

(c)

Pong

(d) Doorkey

algorithms in larger problems; (ii) training efficiency , the ability to learn algorithms from smaller problems and still maintain a high-level of extrapolation capability; and (iii) the higher parameter efficiency of the model. We compare our approach against previous recurrent solver baselines (Section 4.3), considering both same-size and different-size problems (Section 4.2). In Appendix B.3 we present the architecture of our model and training hyperparameters used in our evaluation.

4.1 Methods

Model Training : We apply NeuralSolver in supervised learning scenarios with datasets of input images/features and output images/features, in the case of same-size tasks, or 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 our 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 we select the later model checkpoint.

Evaluation Metrics : In 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, 16], using a significance level of α = 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. [5], consisting of three tasks: a logical toy task ( Prefix-Sum ), a 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 a detailed description of these tasks we refer the reader to the original paper.

Due to the lack of a literature-standard 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 in Appendix A):

GoTo (Figure 4a): Inspired by the Minigrid environment [18], the goal of the task is to select the action that moves the agent (in green) closer to the exit position (in red). The classification target is a four-dimensional 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 a four-dimensional one-hot vector, corresponding to the available actions (up, down, left, right). The name is an abbreviation of '1 Step Maze'.

Pong (Figure 4c): Inspired by the Atari Pong game, the goal of the task is to select the action that moves the paddle horizontally to be underneath the ball. The classification target is a threedimensional one-hot vector, corresponding to the available actions (left, right, no action).

Very Small

Small

Medium

Large

48

128

256

512

0%

50%

100%

48

128

256

512

32

64

96

128

32

64

96

128

(a)

NeuralSolver

Prefix-Sum

Prefix-Sum

(c)

NeuralSolver

Mazes

Mazes

Doorkey (Figure 4d): Inspired by the Doorkey environment in Minigrid [18], the goal of this multistep task is to reach the goal position (red pixel on the bottom right) by first picking-up a key and using it to open a lock door. The classification target is a four-dimensional one-hot vector, corresponding to the available actions (forward, rotate right, grab, toggle).

4.3 Baselines

Bansal et al. [6]: The current state-of-the-art recurrent solver architecture, that combines a recurrent Resnet network with a recall module and progressive loss. We employ the original source code from the authors and suggested training hyperparameters (when available).

FeedForward : A 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 is 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 introducing a global pooling layer similar to that of our model.

5 Results

5.1 Same-size Tasks

Extrapolation Capability : In Table 1 we present the extrapolation accuracy of NeuralSolver against the baselines in same-size tasks [5]. NeuralSolver achieves state-of-the-art extrapolation

performance on same-size tasks. In particular, our approach is able to achieve a 100% accuracy on the Prefix-Sum and Maze tasks, when evaluated on a scenario with observations 16 and 5 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 5 we show the accuracy performance of the learned algorithm with our model when trained with different problem sizes. The results show that NeuralSolver 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-Sum ) and 16dimensional ( Mazes ) input observatons. The results for the other environments are in Appendix C.1.

Parameter Efficiency : In Table 2 we present the total parameter count of NeuralSolver against the baselines in same-size tasks [5]. The results show that our model requires less than 10% of the total parameters of the baselines.

5.2 Different-size Tasks

Extrapolation capabilities : In Table 3 we present the extrapolation performance of NeuralSolver against the baselines in different-size tasks. Our model also outperforms prior state-of-the-art recurrent solvers across all different-size tasks: our approach is able to achieve upper-bound performance when extrapolating to observation sizes 9 ( 1S-Maze ) and 6 ( GoTo , Pong , Doorkey ) times larger than the ones provided during training. In Appendix C.7, we explore the setting of extreme extrapolation, evaluating our model on very large image sizes ( 256 × 256 and 512 × 512 ). The results show that, even 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 our model against the Bansal et al. [6] model with different training sizes. The results show that, while the previous state-of-the-art 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 to upper-bound 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 a 92% reduction in the number of parameters

Model

1S-Maze

GoTo

Pong

DoorKey

NeuralSolver

100.00

±

0

.

00

100.00

±

0

.

00

100.00

±

0

.

00

100.00

±

0

.

00

Use AvgPool

81.96

±

28

.

50

†

97.00

±

3

.

87

†

92.86

±

14

.

66

†

51.86

±

38

.

98

†

Use 5L

78.00

±

5

.

73

†

84.16

±

16

.

39

†

100.00

±

0

.

00

92.99

±

7

.

62

†

No CL

95.26

±

7

.

29

†

64.12

±

23

.

50

†

100.00

±

0

.

00

97.23

±

4

.

73

*

No LSTM

81.87

±

10

.

09

†

67.28

±

14

.

51

†

76.12

±

17

.

63

†

93.68

±

7

.

88

†

Very Small

Small

Medium

Large

33

65

97

129

0%

50%

100%

32

64

96

128

32

64

96

128

32

64

96

128

NeuralSolver

33

65

97

129

0%

50%

100%

(a)

1S-Maze

32

64

96

128

(b)

G

oTo

32

64

96

128

(c)

P

ong

32

64

96

128

(d)

D

oorkey

against Bansal et al. [6] (3.15 million parameters) and a reduction of 99.6% against the FeedForward baseline (71.57 million parameters).

5.3 Ablation Study

In Table 4 we perform an ablation study on the components of NeuralSolver , particularly on the role of: (i) the aggregation function, (ii) the depth of 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 of 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), we replace the LSTM with a ResNet block (No LSTM).

The results show that every component of NeuralSolver contributes to 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 a significant improvement in performance. This result is aligned with previous works that have shown that gated-based recurrent neural networks empirically learn and generalize better than recurrent ResNets [19]. The use of max pooling as the aggregation function results in better performance than average pooling. Removing curriculum learning also results in a 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.

(a) Oracle

(b)

NeuralSolver

(d) FeedForward

We also observe a performance gain in using a single convolutional layer over a 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 on Sequential Decision-Making Tasks

We highlight the versatility of NeuralSolver by exploring visual imitation learning problems, in which, at each time-step, the algorithm is provided with an image of arbitrary size, and needs to output an action. These tasks can be solved through modern reinforcement learning methods [14, 20-26]. However, these methods lack the ability to extrapolate , i.e., are unable to learn algorithms (policies) on small problems, with lower-dimensional observations, and execute on larger problems, with higher-dimensional observations. We consider the DoorKey environment and train an imitationlearning policy algorithm using NeuralSolver with a dataset collected by a pretrained oracle agent. Subsequently, we employ the learned algorithm in the Minigrid DoorKey environment [18], directly as the agent'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 5 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 to the accumulation of errors during task-execution, resulting in sub-optimal trajectories as shown in Figure 7.

6 Conclusion

We proposed NeuralSolver , a simple architecture that learns algorithms that can perform extrapolation across general tasks. We showed that our architecture consistently outperforms prior state-of-the-art recurrent solvers in regard to extrapolation performance, training efficiency and parameter efficiency. In future work, we plan on further improving the training efficiency of our method and explore it as an architecture for reinforcement learning to learn to perform sequential decision-making tasks in an online manner, while maintaining the ability to extrapolate.

References

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: 10.1109/ICCV.2015.123.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 770-778, 2016.

V. Mnih, K. Kavukcuoglu, D. Silver, A. Graves, I. Antonoglou, D. Wierstra, and M. Riedmiller. Playing Atari with deep reinforcement learning. CoRR , abs/1312.5602, 2013.

John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature , 596(7873):583-589, 2021.

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 Marc'Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual , pages 6695-6706, 2021.

Arpit Bansal, Avi Schwarzschild, Eitan Borgnia, Zeyad Emam, Furong Huang, Micah Goldblum, and Tom Goldstein. End-to-end algorithm synthesis with recurrent networks: Extrapolation without overthinking. In NeurIPS , 2022.

Alexandre Salle and Shervin Malmasi. A simple loss function for convergent algorithm synthesis using rnns. 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.

Alexandre Salle and Marcelo O. R. Prates. Think again networks and the delta loss. CoRR , abs/1904.11816, 2019.

Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation , 9(8):1735-1780, 1997. ISSN 0899-7667. doi: 10.1162/neco.1997.9.8.1735.

Xingjian Shi, Zhourong Chen, Hao Wang, Dit-Yan Yeung, Wai-Kin Wong, and Wang-chun Woo. Convolutional LSTM network: A machine learning approach for precipitation nowcasting. In Corinna Cortes, Neil D. Lawrence, Daniel D. Lee, Masashi Sugiyama, and Roman Garnett, editors, Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada , pages 802-810, 2015.

Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. CoRR , abs/1607.06450, 2016. doi: 10.48550/arxiv.1607.06450.

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.

Wojciech Zaremba and Ilya Sutskever. Learning to execute. CoRR , abs/1410.4615, 2014. doi: 10.48550/ arxiv.1410.4615.

R. Sutton and A. Barto. Reinforcement Learning: An Introduction . MIT press, 2018.

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 Màrquez, editors, Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume 1: Long Papers , pages 2773-2785. Association for Computational Linguistics, 2019. doi: 10.18653/V1/P19-1266.

Eustasio Del Barrio, Juan A Cuesta-Albertos, 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.

C.E. Bonferroni. Teoria statistica delle classi e calcolo delle probabilità . Pubblicazioni del R. Istituto superiore di scienze economiche e commerciali di Firenze. Seeber, 1936.

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 environments for goal-oriented tasks. CoRR , abs/2306.13831, 2023.

Eric Price, Wojciech Zaremba, and Ilya Sutskever. Extensions and limitations of the neural GPU. CoRR , abs/1611.00736, 2016. doi: 10.48550/arxiv.1611.00736.

Peter R Wurman, Samuel Barrett, Kenta Kawamoto, James MacGlashan, Kaushik Subramanian, Thomas J Walsh, Roberto Capobianco, Alisa Devlic, Franziska Eckert, Florian Fuchs, et al. Outracing champion gran turismo drivers with deep reinforcement learning. Nature , 602(7896):223-228, 2022.

Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and applications. arXiv preprint arXiv:1812.05905 , 2018.

Volodymyr Mnih et al. reinforcement learning." nature 518.7540 (2015): 529-533. Nature , 518:529-533, 2015.

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 , 588(7839):604-609, 2020.

Xu Wang, Sen Wang, Xingxing Liang, Dawei Zhao, Jincai Huang, Xin Xu, Bin Dai, and Qiguang Miao. Deep reinforcement learning: A survey. IEEE Transactions on Neural Networks and Learning Systems , 2022.

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 neural information processing systems , 34:15084-15097, 2021.

Yueh-Hua Wu, Xiaolong Wang, and Masashi Hamaya. Elastic decision transformer. Advances in Neural Information Processing Systems , 36, 2024.

Yarin Gal and Zoubin Ghahramani. A theoretically 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.

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 arXiv:1406.1078 , 2014.

Vijay Veerabadran, Srinivas Ravishankar, Yuan Tang, Ritik Raina, and Virginia de Sa. Adaptive recurrent vision performs zero-shot computation scaling to 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. URL http://papers.nips.cc/paper_ files/paper/2023/hash/3a40e042c66e84659249f3254460c123-Abstract-Conference.html .

Dennis Ulmer, Christian Hardmeier, and Jes Frellsen. deep-significance - easy and meaningful statistical significance testing in the age of neural networks. CoRR , abs/2204.06815, 2022. doi: 10.48550/ARXIV. 2204.06815.

A Additional Details on the Evaluation Scenarios

We evaluate NeuralSolver on the same-size tasks benchmark introduced in Schwarzschild et al. [5] and on a novel set of different-size tasks. In Table 6 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 to the lack of a benchmark for different-size tasks, we introduce a set of four classification tasks with arbitrary input (image) observation dimensionality and fixed output size (classification labels). For each task we generate a 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 [18], the GoTo task is a simple grid setting, with a 1-pixel white border, and 1-pixel green player, and a 1-pixel red goal, as shown in Figure 4a. The classification target is a one-hot vector of 4 dimensions, corresponding to the next-step actions available to the agent: up, down, left, right. To generate the dataset we use a simple path-finding 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.

1S-Maze : This task is based on the Maze task from [5]. The input is a 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 we made the green player have a random start position so that it does not become trivial to solve. We also changed the thickness of the walls and the border to 1 pixel, as shown in Figure 4b. The classification target is a 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 a depth-first search algorithm, where the label is the next action to take to reach the goal.

Pong : Inspired by the Atari game, the Pong task is a simplified version of the game, where the objective is to move the paddle horizontally to the ball position, as seen in Figure 4c. The classification target is a 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.

Doorkey : This task is inspired by the Minigrid Doorkey environment [18]. The goal is to move the player (red pixel on top) to the goal (red pixel on bottom right) through a sequence of steps, as seen in Figure 4d: 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 4 dimensions, corresponding to the next-step actions available to the agent: forward, rotate right, grab, toggle. We generate the targets using a simple oracle agent that solves the task. In the

NeuralSolver

0

200

400

0

50

100

0

500

1000

0

500

1000

0

100

200

(a)

Prefix-Sum

(b)

Maze

(c)

Thin-Maze

(d)

Chess

original reinforcement learning environment, upon reaching the goal the agent receives a reward of r = 1 -0 . 9 ∗ ( t T ) , where t is the number of steps of the agent and T is the maximum episode length (if t > T 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 task benchmark introduced in [5]. We present a short description of the tasks used in the benchmark. For a more detailed description of the tasks, please refer to the original paper.

Prefix-Sum : This task consists of computing the prefix-sum module two of an input binary array, where the i n bit of the output is the module two sum of the current and the previous n bits of the input.

Maze : This is a maze-solving task where the output is the path from the player to the goal. The training is done on mazes of size 24 × 24 pixels, and tested on mazes of size 124 × 124. While on the original paper, it is referred that the mazes have a training size of 9 and a test size of 59, each maze has a border of 3 pixels and the paths have a width of 2 pixels, thus 9 × 2 + 3 × 2 = 24 and 59 × 2 + 3 × 2 = 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 1 pixel. Thus the training and test sizes are 11 × 11 and 61 × 61, respectively.

Chess : The chess task comprises of chess puzzles, where the input is an 8 × 8 × 12 array indicating the position of the pieces on the board, and the output is an 8 × 8 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 of 1,385, while testing employs 100,000 examples above this threshold.

A.3 Training Efficiency Task Sizes

To simplify the interpretation, we omitted the training sizes of each task used in Figure 5 and Figure 6. In Table 7 we detail the training sizes of each task.

NeuralSolver

0

1000

2000

25

50

75

100

0

500

1000

0

200

400

0

500

1000

(a)

1S-Maze

(b)

GoTo

(c)

Pong

(d)

Doorkey

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 8) and different-size (Figure 9) tasks as a function of 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 or 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 1 × 1 filters and the remaining ones with 3 × 3 filters;

Reducing the amount of ResNet blocks in the original Bansal et al. [6] architecture from two to one;

Removing the ResNet block, and using just a single recall convolution;

Using an Tanh regularization in the output of the recurrent module, to 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 we present the model implementation used across all tasks. In Table 9 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 150 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 of 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 is more computationally efficient than Bansal et al. [6], requiring less than 30% of the operations per training example.

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 1 day to run in all tasks, with the exception of the chess task, where each run takes around 58 hours for the Bansal et al. [6] model and 26 hours for NeuralSolver .

B.6 Comparison with Bansal et al. [6]

In this section we highlight the major differences between NeuralSolver and Bansal et al. [6], shown in 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).

Recurrent Module : The recurrent module in NeuralSolver is a layernorm convolutional LSTM (Figure 11) where x is the input, h is the hidden state, c 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 LSTM 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.

Removal of Projection Layer : We observed a small negative impact in extrapolation when using an initial projection layer for the recurrent module, and as such we removed it from the architecture. This results can be seen in Table 18 of Section C.6.

Processing Module : The processing module consists of three convolutional layers with widths specified in Table 8, 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 an aggregation layer at the end of the model,

x

Recurrent

Block

(ResNet)

Processing

Block

y

Projection

Layer

x

Recurrent

Block

(LSTM)

Processing

Block

y

Aggregation

Layer

NeuralSolver

used for tasks with a different a 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.

A complete description of the final architecture can be described 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 γ . As such, we perform p t = W 3 ∗ γ ( W 2 ∗ γ ( W 1 ∗ h t )) , and a final softmax activation ( σ ), ˆ y t = σ ( 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, ˆ y t = σ ( G ( p t )) .

C Additional Results

C.1 Extended Training Efficiency on Same-size Tasks

In Figure 12 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. [6] being slightly more performant on easier puzzle data, and NeuralSolver being slightly better with more data.

C.2 Training and Validation Accuracy on Different-size Tasks

In Table 11 and Table 12 we show the training and validation accuracy for all models in the differentsize tasks. The results show that all models are trained to optimal performance.

C.3 Fine-tuning of the Progressive Loss Parameter

The Bansal et al. [6] model introduces a progressive loss term to help the model extrapolate. However, this term has an alpha ( α ) parameter that needs to be tuned to each individual task. In Table 13 we show the results of the fine tuning of the alpha parameter for the 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 α = 0 and the α = 0 . 1 for the same-size benchmark comparison. For the remaining same-size tasks, we employ the author's suggested values for α .

x

Convolution

Layernorm

Convolution

Layernorm

pre-computed

+

LSTM

gate

logic

Layernorm

h t

c t

Dropout

c t+1

h t+1

Processing

Block

7/100k

9/200k

11/400k

13/600k

33

65

97

129

0%

50%

100%

33

65

97

129

700000

700000

(a)

NeuralSolver

Thin-Maze

Thin

Maze

(c)

NeuralSolver

Chess

Chess

C.4 Ablation on Model Size of NeuralSolver

We conducted an additional evaluation of the extrapolation performance of our method, as done in Section 5.2, for different sizes of the LSTM's output channel in the different-size tasks. As expected, we can observe a 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 on the Recurrent Module for NeuralSolver

We compare the performance of our model with two different recurrent modules: LSTM (our default method), GRU [28] and LocRNN [29]. In Table 15, we 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 16 and Table 17, respectively.

C.6 Extended Ablation Results

In Table 18, we performed additional ablation on the extrapolation abilities when we removed the layer norm component of the recurrent neural network (No LN) and when we disabled the standard PyTorch 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. [6] on extrapolating to very highdimensional input observations. In particular, we evaluate the performance of the methods on

extrapolating to 256 × 256 (Table 19) and 512 × 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 2 The results show that only NeuralSolver is able to maintain perfect accuracy across all tasks. Moreover, our model does not suffer from overthinking, as shown in Figure 13.

C.8 Hyperparameter Scan

In Fig. 14 we perform a small hyperparameter scan on the choosen default values in Table 9.

GoTo

1S-Maze

Pong

DoorKey

0

500

1000

1500

2000

0

20

40

60

80

100

0

500

1000

1500

2000

0

20

40

60

80

100

(a)

NeuralSolver

AdaGrad

Adam

AdamW RMSProp

SGD

60

70

80

90

100

Mean Value

Optimizer

2.5e-4

1e-3

3e-3

1e-2

70

75

80

85

90

95

100

Learning Rate

0

3

10

70

80

90

100

Warm-Up

25

50

100

150

99.90

99.92

99.94

99.96

99.98

100.00

Mean Value

Epochs

1

2

3

5

10

100

no clip

70

80

90

100

Clip

joint train

2

4

8

80

85

90

95

100

Curriculum Learning Epochs

0

1e-4

2e-4

1e-3

1e-2

85

90

95

100

Mean Value

Weight Decay

0

0.1

0.3

0.5

0.75

97

98

99

100

Standard Pytorch Dropout

0

0.1

0.2

0.4

0.5

0.75

97.5

98.0

98.5

99.0

99.5

100.0

Gal Dropout

FeedForward

NeuralSolver

FeedForward

NeuralSolver

1.00

1.00

1.00

1.00

0.00

1.00

0.51

1.00

0.00

1.00

1.00

1.00

0.00

0.41

0.28

1.00

Prefix-Sum

FeedForward

NeuralSolver

FeedForward

NeuralSolver

1.00

1.00

1.00

1.00

0.00

1.00

0.85

1.00

0.00

1.00

1.00

1.00

0.00

0.03

0.03

1.00

Maze

0.0

0.2

0.4

0.6

0.8

1.0

FeedForward

NeuralSolver

FeedForward

NeuralSolver

1.00

1.00

1.00

1.00

0.00

1.00

1.00

1.00

0.00

0.78

1.00

1.00

0.00

0.00

0.00

1.00

Thin-Maze

FeedForward

NeuralSolver

FeedForward

NeuralSolver

1.00

1.00

1.00

1.00

0.00

1.00

0.03

1.00

0.35

1.00

1.00

1.00

0.00

0.00

0.00

1.00

Chess

0.0

0.2

0.4

0.6

0.8

1.0

D Statistical Significance

The Almost Stochastic Order test (ASO; [15, 16]) was recently introduced to assess statistical significance in Deep Neural Networks across multiple runs. The ASO test measures the presence of a stochastic order between two models or algorithms based on their respective sets of evaluation scores. Using the individual scores of algorithm A and B across various random seeds, the method calculates a test-specific value ( ϵ min), representing how far algorithm A is from being significantly superior to algorithm B. The lower the value of ϵ min, the more likely it is that A is better than B. For ϵ min < 0 . 5 , A is said to be almost stochastically dominant over B in most cases than vice versa, and thus A could be considered superior to B, although with less confidence. We claim that A is stochastically dominant over B with a predefined significance level when ϵ min = 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 a confidence level of α = 0 . 05 (before adjusting for all pair-wise comparisons using the Bonferroni correction [17]).

Random

FeedForward

NeuralSolver

Random

FeedForward

NeuralSolver

1.00

1.00

1.00

1.00

0.00

1.00

1.00

1.00

0.00

0.38

1.00

1.00

0.00

0.00

0.00

1.00

1S-Maze

Random

FeedForward

NeuralSolver

Random

FeedForward

NeuralSolver

1.00

1.00

1.00

1.00

0.00

1.00

1.00

1.00

0.00

0.19

1.00

1.00

0.00

0.00

0.00

1.00

GoTo

0.0

0.2

0.4

0.6

0.8

1.0

Random

FeedForward

NeuralSolver

Random

FeedForward

NeuralSolver

1.00

1.00

1.00

1.00

0.44

1.00

1.00

1.00

0.00

0.00

1.00

1.00

0.00

0.00

0.00

1.00

Pong

Random

FeedForward

NeuralSolver

Random

FeedForward

NeuralSolver

1.00

1.00

1.00

1.00

0.00

1.00

1.00

1.00

0.00

0.00

1.00

1.00

0.00

0.00

0.00

1.00

Doorkey

0.0

0.2

0.4

0.6

0.8

1.0

No LSTM

No CL

Use 5L

Use AvgPool

NeuralSolver

No LSTM

No CL

Use 5L

Use AvgPool

NeuralSolver

1.00

1.00

0.50

0.75

1.00

0.06

1.00

0.00

0.37

1.00

1.00

1.00

1.00

0.81

1.00

1.00

1.00

1.00

1.00

1.00

0.00

0.00

0.00

0.00

1.00

1S-Maze

No LSTM

No CL

Use 5L

Use AvgPool

NeuralSolver

No LSTM

No CL

Use 5L

Use AvgPool

NeuralSolver

1.00

1.00

1.00

1.00

1.00

1.00

1.00

1.00

1.00

1.00

0.14

0.22

1.00

1.00

1.00

0.00

0.00

0.08

1.00

1.00

0.00

0.00

0.00

0.00

1.00

GoTo

0.0

0.2

0.4

0.6

0.8

1.0

No LSTM

No CL

Use 5L

Use AvgPool

NeuralSolver

No LSTM

No CL

Use 5L

Use AvgPool

NeuralSolver

1.00

1.00

1.00

1.00

1.00

0.00

1.00

0.50

0.07

0.50

0.00

0.50

1.00

0.07

0.50

0.25

0.97

0.97

1.00

1.00

0.00

0.50

0.50

0.00

1.00

Pong

No LSTM

No CL

Use 5L

Use AvgPool

NeuralSolver

No LSTM

No CL

Use 5L

Use AvgPool

NeuralSolver

1.00

1.00

0.87

0.02

1.00

0.59

1.00

0.36

0.02

1.00

1.00

1.00

1.00

0.02

1.00

1.00

1.00

1.00

1.00

1.00

0.00

0.03

0.00

0.00

1.00

Doorkey

0.0

0.2

0.4

0.6

0.8

1.0

FeedForward

NeuralSolver

Oracle

FeedForward

NeuralSolver

Oracle

1.00

1.00

1.00

1.00

1.00

1.00

1.00

1.00

0.80

0.49

1.00

1.00

0.06

0.02

0.06

1.00

FeedForward

NeuralSolver

Oracle

FeedForward

NeuralSolver

Oracle

1.00

1.00

1.00

1.00

0.00

1.00

1.00

1.00

0.00

0.30

1.00

1.00

0.00

0.02

0.06

1.00

FeedForward

NeuralSolver

Oracle

FeedForward

NeuralSolver

Oracle

1.00

1.00

1.00

1.00

0.55

1.00

1.00

1.00

0.00

0.00

1.00

1.00

0.00

0.00

0.36

1.00

0.0

0.2

0.4

0.6

0.8

1.0

FeedForward

NeuralSolver

Oracle

FeedForward

NeuralSolver

Oracle

1.00

1.00

1.00

1.00

0.35

1.00

1.00

1.00

0.00

0.00

1.00

1.00

0.00

0.00

0.22

1.00

FeedForward

NeuralSolver

Oracle

FeedForward

NeuralSolver

Oracle

1.00

1.00

1.00

1.00

0.29

1.00

1.00

1.00

0.00

0.00

1.00

1.00

0.00

0.00

0.00

1.00

0.0

0.2

0.4

0.6

0.8

1.0

1.00

0.50

0.10

0.01

0.00

1.00

0.50

0.10

0.01

0.00

1.00

1.00

1.00

0.71

1.00

0.23

1.00

0.80

0.06

0.02

0.44

1.00

1.00

0.24

0.36

1.00

1.00

1.00

1.00

1.00

0.67

1.00

1.00

0.44

1.00

1S-Maze

1.00

0.50

0.10

0.01

0.00

1.00

0.50

0.10

0.01

0.00

1.00

0.66

1.00

1.00

1.00

1.00

1.00

1.00

1.00

1.00

0.74

0.60

1.00

1.00

1.00

0.90

0.30

1.00

1.00

1.00

0.70

0.35

1.00

0.98

1.00

Goto

1.00

0.50

0.10

0.01

0.00

1.00

0.50

0.10

0.01

0.00

1.00

0.65

0.34

1.00

0.70

1.00

1.00

0.49

1.00

1.00

1.00

1.00

1.00

1.00

1.00

1.00

0.58

0.28

1.00

0.58

1.00

1.00

0.38

1.00

1.00

Pong

0.0

0.2

0.4

0.6

0.8

1.0

1.00

0.50

0.10

0.01

0.00

1.00

0.50

0.10

0.01

0.00

1.00

0.75

0.59

1.00

0.76

1.00

1.00

1.00

1.00

1.00

1.00

0.83

1.00

1.00

1.00

0.09

0.21

0.07

1.00

0.35

1.00

0.88

1.00

1.00

1.00

Doorkey

1.00

0.50

0.10

0.01

0.00

1.00

0.50

0.10

0.01

0.00

1.00

0.10

1.00

1.00

1.00

1.00

1.00

1.00

1.00

1.00

0.42

0.16

1.00

1.00

1.00

0.01

0.00

0.43

1.00

1.00

0.01

0.00

0.39

0.76

1.00

Thin-Maze

0.0

0.2

0.4

0.6

0.8

1.0

Iteration: 0

Iteration: 5

Iteration: 15

Iteration: 30

Iteration: 50

Iteration: 70

D

U

Iteration:80

D

U

R

Iteration: 90

D

Iteration: 100

D

Iteration:150

D

E Learned Algorithms with NeuralSolver

Similar to Section 3.2, we 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 over the input image, while the one below shows the predicted action probabilities outputted by the model.

With these visualizations we notice how the information propagates through the input problems, giving a hint of the algorithm that the model is 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 a fixed state around the player position.

E.1 1S-Maze

By looking at how the information propagates through the maze paths in Figure 20 and 21, the model appears to 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 the optimal path between the player and the goal. The algorithm converges to the optimal task after the difference of the state converges near the player. Figure 21 shows that finding the optimal path is not necessary to predict the next action.

Iteration: 0

Iteration: 5

Iteration: 15

Iteration: 50

Iteration:70

R

D

U

R

D

R

D

U

D

U

R

D

U

Iteration:100

Iteration:150

Iteration:200

Iteration:250

Iteration:300

D

R

D

R

R

Iteration:0

Iteration:10

Iteration:20

Iteration:40

Iteration: 60

L

U

R

D

L

U

R

D

L

U

R

D

L

U

R

D

U

R

D

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 a signal is sent from the goal position to the positions above, propagating in a circular/oval shape. Thus the model appears to learn an algorithm that attempts to signal the player if it should go down. Similarly, in Figure 23 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 on the left of the goal. In this setting, the horizontal line with a 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.

Iteration:0

Iteration:10

Iteration:20

Iteration:40

Iteration:60

L

U

R

D

U

R

D

L

U

R

D

L

U

R

D

U

R

D

Iteration:0

Iteration:5

Iteration:10

Iteration: 15

Iteration:30

L

U

R

D

L

U

R

D

U

R

D

U

R

D

U

R

D

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 a 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 to the left action. In Figure 26, we see that this propagation of information does not reach the paddle, and as such the action remains unchanged.

We can also see a line with a 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 simultaneously.

Iteration:0

Iteration: 10

Iteration: 20

Iteration: 30

Iteration: 40

S

S

S

R

S

R

S

R

Iteration: 50

Iteration: 60

Iteration:70

Iteration:80

Iteration: 90

E.4 Doorkey

We provide two examples of learned algorithms in Doorkey : in Figure 27 we show a case when the player is almost reaching the goal, and in Figure 28 a 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 a 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 a line that shoots vertically from the door position that is not present in Figure 27. This shows that for more complex tasks, better visualization and interpretation techniques of the recurrent states are required.

Iteration:0

Iteration:10

Iteration:20

Iteration:30

Iteration: 40

S

R

S

R

S

R

S

S

R

Iteration:50

Iteration:60

Iteration:70

Iteration:80

Iteration:90

R

Iteration:0

Iteration:10

Iteration: 20

Iteration:30

Iteration:40

Iteration:50

R

P

Iteration:60

R

P

Iteration:70

R

P

Iteration:80

R

P

Iteration:90

R

P

R

R

R

R

Iteration:0

Iteration:10

Iteration:20

Iteration:30

Iteration:40

Iteration:50

R

P

Iteration:60

R

P

Iteration:70

R

P

Iteration:80

R

P

Iteration:90

R

P

H

H

R

R

F Additional Example Trajectories

Oracle

NeuralSolver

FeedForward

G Examples of tasks with observations of size 512 × 512

We present additional visualizations of examples of tasks with large observations ( 512 × 512 ).

T

页

购

司

品

5

原

园

5

5

2

强

5

原

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 5.

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 or 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.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Even though we did not provide a specific section to discuss the limitations, through the article our experiments show possible improvements in 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 C.1.

Guidelines:

The answer NA means that the paper has no 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 a 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 on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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

Question: For each theoretical result, does the paper provide the full set of assumptions and a 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 crossreferenced.

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 if they appear in the supplemental material, the authors are encouraged to provide a 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 or 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 experimental results of the paper to the extent that it 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 in 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 Section B.6 we highlight the main differences between the Bansal et al. [6] and NeuralSolver implementations that allow also to faithfully implement our model based on the available Bansal et al. [6] descriptions and available code. In Section A we give the details on 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 to 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 a 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 is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to 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.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to 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 instructions to 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 to perform the main experimental results.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on 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. If only a 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) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 a 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

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: All our experiments are accompanied with standard deviation errors and with the ASO[15, 16] significance test described in Section D.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence 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, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to 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.

It 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 a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Section B.4 and Section B.5.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as 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 it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Our research does not involve human participants. 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. Our training tasks and use of previous datasets do not contain any personally identifiable information, do not use or refer to any sort of real people information, as well respect the previous task licenses.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

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 not 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 or unintentional) misuse of the technology.

If there are 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe 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)?

Answer: [NA]

Justification: [NA]

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 to 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 we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: 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?

Answer: [Yes]

Justification: Our models and tasks are based on previous work that are all properly cited and referenced.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.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 be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 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 A we give the details on how to 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 not release new assets.

Researchers should communicate the details of the dataset/code/model as part of 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 is used.

At submission time, remember to 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 as 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 in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in 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]

Justification: [NA]

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on 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.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nSee above.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n1. 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.\n\nSTRENGTHS:\nThe paper is easy to follow and experiments nicely presented. The proposed model seems to outperform (one) related work method and results seem quite impressive.\n\nWEAKNESSES:\nSince the authors introduce novel benchmarks, the very positive experimental results need more ablations. Please provide hyperparameter scan and their results on baselines.\n\nQUESTIONS:\nSee 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.\n\nLIMITATIONS:\nI would welcome a dedicated limitation section.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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\n\nQUESTIONS:\nIt 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?\n\nLIMITATIONS:\nIt 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-08"}