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

We thank the reviewer for their insightful comments and the time spent reviewing our paper. We now address the weaknesses (W) identified by the reviewer:

• (W1) Differences between NeuralThink and the previous work/why NeuralThink achieves improved performance. We would like to highlight the following fundamental differences between our work and previous work (and why):

  1. Recurrent design: Previous work uses the simplest form of recurrence by using a ResNet with weight sharing across residual blocks. In our work, we started by evaluating the extrapolation capabilities of DeepThink and found its performance to be lacking in different-sized tasks (as we show in Table 3). To overcome this gap, we explored other more advanced recurrent designs for algorithmic extrapolation, in particular LSTM-based convolutional designs. Our results show that this choice of recurrent design results in better performance of the method, supporting prior work $3$ .
  2. Aggregation layer in different size tasks: Previous DeepThinking methods do not work in this class of problems, as their models only support same size input and output problems. To overcome this issue, we introduced the use of a pooling layer in the design of our model.
  3. Curriculum-based training scheme: The previous differences are insufficient to allow consistent extrapolation in different-size tasks. To the best of our knowledge, we are the first work that introduces a curriculum-based training scheme for DeepThinking architectures. In Table 4 we highlight the role of our curriculum-based training scheme in the overall performance of the method.

  We provide in Appendix B.6. of the original paper an extended discussion on the differences between our model and DeepThink. We have clarified the existence of this comparison in line 129 of the updated version of the paper.

• (W2) Imperfect results in same-sized tasks conflict with oracle-level performance on different-size task extrapolation: We would like to point out that, while the result for Thin-Maze is not at oracle-level accuracy, the results are still inside of the experimental uncertainty of the other results. On the other hand, the Chess task is fundamentally different from the other tasks, since it is not a spatial extrapolation task but a symbolic extrapolation level, proving to be a harder task for our method. However, despite not achieving oracle-level performance, even in this challenging task our model is able to outperform the previous state-of-the-art.

• (W2,1) The number of used tasks is small: We would like to point out that not only are we using all the benchmarks proposed by $2$ to evaluate DeepThinking models, but we also contribute a novel set of different-sized tasks, due to the absence of an existing adequate benchmark for this type of tasks. Furthermore, we are the first work to explore sequential decision-making scenarios for the application of DeepThinking models.

• (W3) Limited Baselines/AdRNNs: We have updated our related section with $1$ . We have compared directly and indirectly the performance of AdRNNs to NeuralThink.

  • AdRNNs have a worse performance than DeepThink (our baseline) on same-sized tasks. $1$ uses the same training methodology as NeuralThink on the same-sized Maze tasks (train on mazes of size 24x24 pixels) but test on much smaller mazes (test on mazes of size 44x44 and 56x56 pixels) than in our paper. The best method presented in $1$ (LocRNN) achieves 50% (best performance) in test scenarios that are half the size of the ones presented in our work: DeepThink achieves 91% test accuracy on mazes of size 124x124. As such, DeepThink is a stronger baseline than AdRNNs for algorithmic extrapolation and, thus, we outperform them as well.
  • AdRNNs have a worse performance than NeuralThink on different-sized tasks. Despite the fact that no public code implementation of $1$ is available, we implemented the proposed LocRNN without the learnable iteration-halting mechanism (as we always select the best performing iteration for the results) and compared its performance against the ConvLSTM used in NeuralThink on the different-sized tasks. The results are shown in the below table, which highlights that LocRNN has a worse performance than our method, while having more parameters than NeuralThink and being more computationally intensive to execute. We added this additional comparison to Appendix C.5.

• (W4) “x/y” color labels in Figure 5 are not clear.
  We changed the labels of the colors with very small, small, medium and large, and moved details of the training sizes used for each task to the Appendix.

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

$2$ Bansal, Arpit, et al. "End-to-end algorithm synthesis with recurrent networks: Extrapolation without overthinking." Advances in Neural Information Processing Systems 35 (2022)

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

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.

We express our gratitude to the reviewer for their comments and the effort spent on reviewing our paper. We will now tackle the weaknesses (W), questions (Q), and limitations (L) mentioned:

• (W1a) Ablations present in the paper: We provide in the original paper additional ablation studies in Appendix C, due to the lack of space in the main body of the text. In particular we show:

  • Ablation on the progressive loss parameter of DeepThink: Appendix C.3.
  • Ablation on the model size of NeuralThink: Appendix C.4.
  • Ablation on different types of recurrent networks for NeuralThinkL Appendix C.5.
  • Ablation on layernorm, dropout and projection head: Appendix C.6.

  In the updated version of the paper (beginning of Section 5.3) we have emphasized the existence of these additional ablation studies.

• (W1b) Hyperparameter scan details on baselines: In the general rebuttal comment document we show the results of a hyperparameter scan in the GoTo environment for the NeuralThink training hyperparameters. The results show that NeuralThink is robust to hyperparameter changes, with exception of the gradient clip that requires proper tuning. The results further show the significance of using weight decay and both forms of dropout regularization together. In this work we tried to use our model with minimum hyperparameter tuning, keeping the same hyperparameters for the majority of the different-sized environments and using the same hyperparameters for the same-sized tasks proposed by the DeepThink model, as shown in Appendix B.3.

• (Q1, Q4) Comparison between DeepThink and our model: We provide in Appendix B.6. of the original paper an extended discussion on the differences between our model and DeepThink. We have clarified the existence of this comparison in line 129 of the updated version of the paper. Moreover, we also highlight in Appendix C.7. how our model outperforms DeepThink on extrapolation to very-large input sizes (256x256 and 512x512).

• (Q2) What happens if you scale the models: We provide in Appendix C.4. of the original paper an ablation study on the effect of the model size for the performance of NeuralThink. The results show that decreasing the model size results in a decrease in the performance of NeuralThink, yet still outperforms the previous state-of-the-art. In the updated version of the paper (beginning of Section 5.3) we have emphasized the existence of this ablation study.

• (Q3) When does NeuralThink break? In Figures 5 and 6 we highlight that our model struggles to extrapolate when we significantly reduce the size of the input data provided to the model during training. We believe that learning to extrapolate from very small input sizes is a very interesting direction for future work (as mentioned in line 282).

• (L1) Dedicated limitations section: Due to lack of space in the main body, we have created a dedicated Appendix section for limitations, merging it with the previous Appendix B.2. “Changes We Tried That Did Not Improve NeuralThink” and mentioned the existence of this Appendix in line 282 of the updated version of the paper.

We thank the reviewer for the comments and for the time spent reviewing our paper. We now address the weaknesses (W), questions (Q) and limitations (L) pointed out by the reviewer:

• (W1) Authors do not clearly state what is novel in their architecture in the main text. Due to lack of space in the main body, we have an extended discussion on the differences between NeuralThink and DeepThink in Appendix B.6. Nonetheless, we have improved the description of the novelty of our model in line 100 of the updated version of the paper.
• (W2) Color labels of Fig. 5 and Fig. 6 are confusing: To improve our figures we modified the labels of the colors to “very small”, “small”, “medium” and “large”, and defined these training sizes for each task in Appendix.
• (W3) The exact details of the architecture are not entirely clear: We can describe our architecture as follows: for same size tasks, given an input image $o$, we pass the input through a single layer ConvLSTM ($q$) to obtain the next hidden state ($h\_1$ and cell state ($c\_1$), h\_1, c\_1 \= q(o,h\_0,c\_0) . The initial hidden ($h\_0$) and cell ($c\_0$) states are initialized at zeros. The recurrent process is given by feeding back the previous hidden and cell states to the LSTM, h\_t, c\_t \= q(o,h\_{t-1},c\_{t-1}). After a fixed size of recurrent iterations, we can obtain the model prediction at timestep $t$ by passing the hidden state $h\_t$ through the processing module. The processing module is composed of three convolutional layers with 3x3 kernels denoted by $W\_1,W\_2,W\_3$ and ReLU activation ($\\gamma$). As such, we perform p\_t \= W\_3 \\ast \\gamma(W\_2 \\ast \\gamma(W\_1 \\ast h\_t)) , and a final softmax activation ($\\sigma$), \\hat{y}\_t \= \\sigma(p\_t). For different size tasks, before the softmax we perform a global maxpool operation ($G$) that reduces the processing output height and width to 1, \\hat{y}\_t \= \\sigma(G(p\_t)). In the updated version of the paper we added this description to Appendix B.6.
• (W4) The authors use LSTMs but do not cite the original article: We have updated the paper to include the citation of the original LSTM paper.
• (Q1) It would be nice to see an ablation on the layernorms in Tab. 4. In the original version of the paper we included the ablation on layernorm in Table 15 of Appendix C6, due to lack of space in the main body of the paper. In the updated version of the paper we have emphasized the existence of additional ablation studies available in Appendix, in the beginning of Section 5.3.
• (Q2a) Is the best accuracy obtained by the models at any iteration used in NeuralThink? Yes, we use this evaluation methodology for all methods.
• (Q2b) Does NeuralThink suffer from overthinking? In Appendix B.1. we empirically show that NeuralThink does not suffer from overthinking, across all evaluation tasks. We have updated the text in line 213 to clarify this point.
• (Q3) Why is the thin maze more difficult than the thick one, especially for the baselines? We believe that the difficulty arises from the size of the receptor field in the convolutional layers required to perceive the environment. The original DeepThink architecture uses 5 convolutional layers with a kernel size of 3 for its recurrent module. This means that at each recurrent step, the network is capturing information from 5 adjacent pixels in each direction (i.e., a square of 11x11 pixels). This fact makes it harder for the model to generalize from smaller problems as the model has to learn an iterative algorithm from a square of 11 pixels at each iterative step. On the other hand, the ConvLSTM used by NeuralThink only captures information from the next adjacent input pixel (or squares 3x3 pixels). This design choice allows NeuralThink to learn general iterative algorithms from smaller sizes.
• (L1) It is not clear how NeuralThink can work on real-world applications: We have evaluated NeuralThink on both existing benchmarks of DeepThinking models as well as on novel benchmark scenarios that we contribute for different-sized tasks. We have also explored the use of our model in sequential decision-making tasks, such as the DoorKey environment of the MiniGrid suite in Section 5.4, which is still a standard environment for the evaluation of modern reinforcement learning agents (see $1a-1b$ ).

References:

$1a$ Nikulin, Alexander, et al. "XLand-minigrid: Scalable meta-reinforcement learning environments in JAX." arXiv preprint arXiv:2312.12044 (2023).

$1b$ Pignatelli, Eduardo, et al. "NAVIX: Scaling MiniGrid Environments with JAX." arXiv preprint arXiv:2407.19396 (2024).