NeuralSolver: Learning Algorithms For Consistent and Efficient Extrapolation Across General Tasks

Authors: bernardo  esteves
miguel  vasco
francisco s melo
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-11 17:42:17

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.

Main paper: 
 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).Figure 1 : Observations of the 1S-Maze environment of sizes 7×7, 11×11, 33×33 and 129×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 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. 
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 We design NeuralSolver with two fundamental components: (i) a recurrent module (purple), 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 our architecture to be used both in same-size and different-size tasks. Additionally, we employ a curriculum-based training scheme to improve the extrapolation performance of our architecture.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 ArchitectureWe 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 [5] [6] [7] 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 .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. 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. 
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] [10] [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 both in same-size and different-size tasks, with minimum changes. In Section 5.3 we compare the performance of different aggregation layers.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.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.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. 
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 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.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). Mazes Figure 5 : Training efficiency of NeuralSolver and Bansal et al. [6] 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, detailed in Appendix A.3 In the dashed line we show the upper-bound on the performance. 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).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.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. 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). 
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 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.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 against Bansal et al. [6] (3.15 million parameters) and a reduction of 99.6% against the FeedForward baseline (71.57 million parameters).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. 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.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] [21] [22] [23] [24] [25] [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 . 
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. 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. 
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). 
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.  
 
 Appendix: 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 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.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.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.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. 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.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)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.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. 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, 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 (σ), ŷ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, ŷt = σ(G(p t )).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.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.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 α. Chess Figure 12 : Training efficiency of NeuralSolver and Bansal et al. [6] 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-bound on the performance.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 .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.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.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 .In Fig. 14 we perform a small hyperparameter scan on the choosen default values in Table 9 .2 Each test on the 1S-Maze task with observations of size 512 × 512 took on average 28 hours on a Nvidia 4090 GPU. Figure 13 : Accuracy results of extrapolation to large observation sizes (512×512). All results are averaged over 10 randomly-selected seeds. In the 1S-Maze environment, the real number of iterations is 10x larger. Higher is better. 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. ASO scores are expressed in ϵ min , with a significance level α = 0.05 that is adjusted accordingly by using the Bonferroni correction [17] . Read from row to column: e.g., NeuralSolver (row) is almost stochastically dominant over Bansal et al. [6] (column) in the Prefix-Sum task with ϵ min of 0.28.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] ). ASO scores are expressed in ϵ min , with a significance level α = 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 ϵ min of 0.00. 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.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. 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. 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. 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. In Figure 29 we present additional trajectories of agents that execute policies learned using NeuralSolver and other baselines, following the same procedure as in Section 5.4. 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. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

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

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

The authors are using LSTMs but do not cite [1].

[1] Hochreiter et al, 1997:  Long Short-Term Memory
questions: questions: It would be nice to see an ablation on the layernorms in Tab. 4.

In 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?

Why is think maze more difficult than the thick one, especially for the baselines?
limitations: 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.
paper_summary: paper_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.
2. 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).
3. The architecture consists of three main components:
   - A recurrent module that iteratively processes input information at different scales
   - A processing module that aggregates previously processed information
   - A curriculum-based training scheme to improve extrapolation performance
4. NeuralThink outperforms prior state-of-the-art deep thinking approaches in:
   - Extrapolation capability: Consistently executing learned algorithms on larger problems
   - Training efficiency: Learning algorithms from smaller problems
   - Parameter efficiency: Requiring fewer parameters than other approaches
5. The authors introduce a set of novel different-size tasks to evaluate their method.
strength_weakness: strength_weakness: Since the authors introduce novel benchmarks, the very positive experimental results need more ablations. Please provide hyperparameter scan and their results on baselines.
questions: questions: See Weaknesses. 
Can 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: limitations: I would welcome a dedicated limitation section.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- 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. 
- 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.
- 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. 

[1] Adaptive recurrent vision performs zero-shot computation scaling to unseen difficulty levels. NeurIPS 2023.
questions: questions: See above.
limitations: limitations: N/A

Comments: title: Thank you!
comment: Thank you for the comments and for raising the score in recognition of our efforts to address them.
Thank you for the discussion and for recognizing our efforts by raising the score.
Thank you for the suggestions (which we will add to the final version of the paper), the overall discussion and for recognizing our efforts with the score raise.
Thanks sincerely for the detailed response from the authors. Raising the score to 5 since the concerns regarding lack of baseline / comparison have been addressed.
I would like to thank the authors for their detailed answers. Given those, I'm raising my score.

I would appreciate it if the authors could include the full set of equations describing their model in the appendix of the final version of the paper. Also, point (W1) from the answer to the reviewer dmPv in the main paper would add clarity to the paper.
Thank you for the additional data. I will raise my score to 6. Many thanks!