Thinking Deeper With Recurrent Networks: Logical Extrapolation Without Overthinking
Authors: Arpit Bansal
Avi Schwarzschild
Eitan Borgnia
Zeyad Emam
Furong Huang
Micah Goldblum
Tom Goldstein
Venue: ICLR 2022 Submitted (2022)
Keywords: Deep learning
recurrent networks
thinking
extrapolation
generalization

Abstract: Classical machine learning systems perform best when they are trained and tested on the same distribution, and they lack a mechanism to increase model power after training is complete. In contrast, recent work has observed that recurrent networks can exhibit logical extrapolation; models trained only on small/simple problem instances can extend their abilities to solve large/complex instances at test time simply by performing more recurrent iterations.  While preliminary results on these ``thinking systems'' are promising, existing recurrent systems, when iterated many times, often collapse rather than improve their performance.  This ``overthinking'' phenomenon has prevented thinking systems from scaling to particularly large and complex problems. In this paper, we design a recall architecture that keeps an explicit copy of the problem instance in memory so that it cannot be forgotten.  We also propose an incremental training routine that prevents the model from learning behaviors that are specific to iteration number and instead pushes it to learn behaviors that can be repeated indefinitely. Together, these design choices encourage models to converge to a steady state solution rather than deteriorate when many iterations are used. These innovations help to tackle the overthinking problem and boost deep thinking behavior on each of the benchmark tasks proposed by  Schwarzschild et al. (2021a).

Under review as a conference paper at ICLR 2022
THINKING DEEPER WITH RECURRENT NETWORKS: LOGICAL EXTRAPOLATION WITHOUT OVERTHINKING
Anonymous authors Paper under double-blind review
ABSTRACT
Classical machine learning systems perform best when they are trained and tested on the same distribution, and they lack mechanism increase model power after training is complete. In contrast, recent work has observed that recurrent networks can exhibit logical extrapolation; models trained only on small/simple problem instances can extend their abilities to solve large/complex instances at test time simply by performing more recurrent iterations. While preliminary results on these "thinking systems" are promising, existing recurrent systems, when iterated many times, often collapse rather than improve their performance. This "overthinking" phenomenon has prevented thinking systems from scaling to particularly large and complex problems. In this paper, we design a recall architecture that keeps explicit copy of the problem instance in memory so that it cannot be forgotten. We also propose an incremental training routine that prevents the model from learning behaviors that are specific iteration number and instead pushes learn behaviors that can be repeated indefinitely. Together, these design choices encourage models converge to steady state solution rather than deteriorate when many iterations are used. These innovations help to tackle the overthinking problem and boost deep thinking behavior on each of the benchmark tasks proposed by [Schwazschild et al. (2021a).
INTRODUCTION
Humans solve complex logical reasoning problems through the process of logical extrapolation they assemble simple logical primitives into complex strategies. For example, a person taught to prove simple lemmas in turn prove more complex theorems simply by expending more cognitive effort. Neural networks have achieved great success at pattern matching tasks, often exceeding human performance, but they lack the ability to solve complex reasoning tasks in scalable way. Recently, deep thinking systems (Schwarzschild et al.) (2221)) have been proposed as a way to represent and learn scalable reasoning processes using recurrent neural networks. The word 'thinking' in this context refers to sequential processing. These systems train recurrent models (networks that recycle parameters between layers) solve reasoning problems. Unlike traditional feed-forward models, which are limited in the complexity of problems they can solve by their finite depth, the effective depth of recurrent models can be expanded after training simply by iterating the recurrent unit for longer. When trained properly, thinking systems learn scalable algorithmic processes for solving problems. After training to solve small/easy problem instances with few recurrent iterations, the
Figure 1: Our improved thinking network trained on 9 x 9 mazes and its solution (left two plots) autonomously synthesizes a scalable algorithm. By running this algorithm for longer, it solves 59 (center) and 201 x 201 (right) mazes without retraining. Standard architectures, and even existing primitive thinking models, fail to tolerate this domain shift.

---

Under review as a conference paper at ICLR 2022
system is then extended to run for more iterations at test time. In doing so, the system achieve logical extrapolation, and solve problems of greater difficulty than those in the training set.
To date, the level of logical extrapolation observed in thinking systems has been quite modest. For example it has been demonstrated that a system trained on 9x9 mazes can extrapolate to solve a 13x 13 maze. These systems fail to achieve greater extrapolation because of the overthinking problem; recurrent systems, when extended too far outside their training regime, often deteriorate and fail to produce interpretable outputs.
In this work, we design purpose-bin buppt neural architectures and specialized training loops to make possible to train systems that do not suffer from overthinking and instead converge fixed point when iterated for thousands of iterations. By doing so, we are able to build thinking systems that exhibit extreme logical extrapolation behaviors, and leap from solving small/simple training problems with tens of iterations to solving large and complex problem instances at test time using hundreds or even thousands of iterations.
We experiment on the benchmark problems made available by Schwarzschild et al. (2021a) for measuring extrapolation behavior. These tasks include computing prefix sums, finding optimal solutions for two-dimensional mazes, and solving chess puzzles. We modify model architectures and propose a novel training routine for recurrent networks that significantly boosts performance on all tasks in the benchmark suite. Additionally, we demonstrate that existing models that have shown positive logical extrapolation on these datasets specifically are susceptible to overthinking, a problem that is largely overcome by new architectures and training loops (Schwarzschild et al. 2021b).
Our contributions can be summarized as follows.
L. We provide recurrent architecture for logical extrapolation in which the problem input is concatenated directly to the feature stack of certain layers in the recurrent thinking module. This prevents the problem instance from being forgotten if deep features become noisy, corrupted, or lossy.
II. We develop a new training routine that incentivizes recurrent networks to make incremental improvements towards solution, improving the feature representation each iteration. This training process removes information about how many times the recurrent module has been applied, which in turns prevents the network from learning iteration-specifie behaviors (e.g., a special behavior for iteration five and another for iteration six) rather than scalable behaviors than can be iterated indefinitely to yield extrapolation.
III. We analyze the overthinking problem and show that our models overcome this phenomenon.
Our improvements in performance on the easy-to-hard benchmark datasets (Schwarzschild al. 2021a) can be categorized in several ways. First, our models yield uniformly higher accuracy across the most difficult tasks used for testing in previous work. Second, we show that our models can extrapolate to much harder/larger examples than are considered in previous work, where the prior methods generalize poorly, if at all. Lastly, we show that our models do not forget solutions in settings where previous models overthink.
2 RELATED WORK
The work in this paper conceptually builds upon the systems recently studied by Schwarzschild et al. (0211)) which showed that simple recurrent architectures, when trained to solve various reasoning problems, can then perform logical extrapolation while their feed-forward counterparts cannot. In this section, we contextualize this approach amongst prior work on algorithm learning, adaptive neural models, and logical extrapolation.
Algorithm learning describes models that learn recurrent processes from data. Early works on this topic study the ability of recurrent neural networks (RNNs) to process input strings of arbitrary lengths (Gers & Schmidhuber 2001 Schmidhuber et al., 2007). More recently, Graves et al. (004)) introduced neural Turing machines designed to mimic programmable computers, and Kaiser &
Sutskever (2015) propose a parallel version inspired by massively parallel graphical processing units. These methods, and improvements to them, show promising results bit string to bit string tasks, including copying inputs and adding integers, and even demonstrate the ability to generalize
from shorter training strings to longer ones at testing (Graves et al. 2014 Kaiser & Sutskever
2

---

Under review as a conference paper at ICLR 2022
2015 Freivalds & Liepins, 2017). These methods based classical RNNs the amount of computation they perform is directly linked to the length of the input string, which prevents them from executing more less computation independently of the input size (or the difficulty of the problem). Moreover, classical RNNs are often trained incrementally to produce one bit at a time, rather than synthesizing algorithm for solving an entire problem end-to-end. This makes it difficult to apply them in situations where the solution to a problem cannot be decomposed into incremental parts (e.g., chess).
Constraint satisfiability problem (CSP) solving networks disentangle the amount of work from the input size (Selsam et al., 2018). Specifically, message passing neural networks can execute more passes to solve harder CSPs (Selsam et al. (018). These systems are specific to the problem of constraint satisfaction for boolean expressions, but they are an early demonstration of scalable algorithmic behavior.
Adaptive neural networks are designed to expend varying amounts of computation on different inputs, thus overcoming the limitation of classical RNNs. Self delimiting neural networks use one neuron to determine when to stop updating the hidden state in RNNs, and in doing so, they perform more or less computation for each token in an input sequence (Schmichuber) 2012). Adaptive compute time (ACT) is an algorithm that provides RNNs with a halting unit which estimates the probability that computation should continue. This algorithm penalizes 'ponder time' during training to encourage the network to solve problems quickly (Graves) 2016). Eyzaguirre & Soto (2020) exhibit strong performance on visual question answering by introducing differentiable version of language models also exist, notably Universal Transformers and Depth-Adaptive Transformers, which utilize ACT to determine the work required for each input (Dehghani et al. 2018: Press et al.] (202T). Similarly, iterative residual networks like NAIS-Nets repeat blocks stages and perform well on image classification (Ciccone al.) [018). All of these works test their methods in-distribution, i.c. where the training and testing data are sampled from the same distribution; logical extrapolation outside the training domain is not considered.
ACT. Adaptive transformer-base
Logical extrapolation denotes the task of generalizing to test sets which comprise more computation- ally complex samples than the training data. NNuamah (202T) claims that "neural network models with end-to-end training pipelines... ... cannot, on their own, successfully perform algorithmic reasoning," and instead proposes a hybrid hand-crafted and learned approach. Similarly, Palm et al.] (2018) propose recurrent relational networks, which operate on graphs by iteratively passing messages. They also claim that classical architectures lack the inductive bias to reason about relationships between objects. Several recent works call these claims into question. [Scwaarzchildd et al. (2021b) employ recurrent networks based on weight-sharing architectures, which can be made deeper at test time independent of the input size (Schwarzschild al., (021c). These systems were shown to exhibit logical extrapolation behavior in several domains. More details on these methods are discussed in Section 3 as our algorithms build on these directly. [Banino et al. (202T) reformulate the halting unit in ACT leading to probabilistic model with improved performance called PonderNet, and importantly their method outperforms ACT on logical extrapolation tasks for prefix sums.
The thinking systems proposed by Schwarzschild et al.) (2021b) depart from classical recurrent networks for (e.g.) text, which learn from step-by-step supervision to produce output tokens one at a time. In contrast, thinking systems autonomously synthesize scalable algorithm end-to-end with no supervision over what each algorithmic step should produce. For this reason, they can potentially be applied to solve complex problems that are difficult or impossible to decompose by hand.
3 METHODS
We begin with some terms and definitions. We study networks that share weights across blocks of layers during training. For example, instead of four distinct residual blocks, a single residual block is repeated four times (see Figure 2 for graphical depiction). At test time, networks trained this way can be made "deeper" to extend their compute budget simply by repeating the block more times. We refer to the number of convolutional layers applied in recurrent network as its "depth", and this quantity grows as the number of recurrent iterations increases. The number of feature maps produced by each convolution is referred to as its "width."
3

---

Under review as a conference paper at ICLR 2022
More formally, let be function representing recurrent block, e.g. ResNet block (He et 2016). and let r" denote n recurrences that function, c.g. r²(x) r(r(x)). Let 6 denote a feature map, or output of r, and let o, r" (z). Finally, consider initial "embedding function" denoted by which projects an input instance into feature space, and also a final "output head" denoted by h, which maps features to outputs. A Deep Thinking (DT) network with m iterations of the recurrent block can then be expressed as follows.
fm(x) == h(r
In our systems comprises a convolutional layer followed by a ReLU, is single four-layer residual block, and h during training and compute gradients for optimization by backpropagating through the unrolled network. Then, m can be increased for testing, allowing these networks increase their processing power and solve larger and harder problems. Below, we consider two modifications to standard approaches for achieving improved extrapolation: modified architectures and modified optimization.
of three convolutional layers with ReLU after the first and second. We fix m
3.1 ADDING RECALL TO THE ARCHITECTURE
When humans think for long time to solve problem, often stop to reread the question or review the task at hand. We modify thinking architectures to periodically recall the input exactly. We incorporate this capability into architectures by concatenating the input problem the features output from each instance of the recurrent block.
LSI
LSI
L
Figure 2: Architecture schematics. Left to right: A feed-forward network, network containing three recurrent blocks (in green) that share weights, and recurrent network with recall.
Popular architectures in computer vision typically incorporate skip connections that similarly pass information from earlier layers forward. In fact, empirical evidence suggests that skip connections. for example in highway networks, ResNets, and DenseNets. stabilize training (Srivastava et al. 2015| He al., 2016 Huang et al. 2017). Our architectural modification driven by the intuition that a noisy training process creates thinking networks that are imperfect and may leak or distort information over time as features are iteratively fed back through the recurrent unit thousands of times. Recall allows the system to look back at the problem at any time and reproduce any missing or damaged features. To formalize this architectural change, adding recall to the network can be expressed using the notation defined above as follows.
fm == aan (p(z), where Frecal(ó,,)) : r((o,)))
(2)
Whereas the input to r at iteration l is usually dx-1, with recall, the input to r at iteration k: is (dx-1,22) or the concatenation of the input with the feature map output by the previous recurrence. We add a single convolutional layer to map the input [Øx-1, to a feature map of the same shape as dk-1. We refer to DT networks with concatenating skip connections as DT-Recall models.
3.2 PROMOTING FORWARD PROGRESS THROUGH OPTIMIZATION
We modify the training objective to encourage the system to incrementally make progress from any starting point. We do this by inputting a problem instance and running the recurrent module for some random number of iterations. We then take the output of this process, re-initialize the network with these features as if iteration had just begun, and train the model to produce the correct solution after a random number of iterations.
This incremental training process has two benefits. First, trains the network to continue improving the quality of partial solutions, even when they contain errors distortions that creep in from
4

---

Under review as a conference paper at ICLR 2022
running many iterations. Second, by choosing features from random iteration to serve as the initial state for the training step, we prevent the network from internally counting iterations and learning iteration-specific behaviors, such behaviors that get executed only on iteration five, for example. Rather, the network is forced to learn ateration-agnostie behaviors that are effective at any stage of the problem solving process.
In our implementation, we randomly sample the number of iterations used to generate a partial solution, n, and the number of training iterations, k. budgeted for the network improve this partial solution. We then update the network's parameters to minimize loss after 7 + k total iterations when it starts with the initial partial solution. This is done by detaching the recurrent module's 's output after iterations from the computation graph before computing the gradient of the loss at iteration + k. The process above is an analog of truncated backpropagation through time (Jaeger] 2002), with a random starting and end point. During training. we ensure that the sum of and A: is less than fixed maximum number of iteration m, which we call the training regime. The incremental loss described above is added to the standard loss computed with a full forward and backward pass through the unrolled m-iteration network.
The training loop and computation of the loss are given in Algorithm [ The incremental progress term is referred to as Pprogressive (or progressive loss), and the contribution to the loss from the fully unrolled network is denoted by Lmax_iters Note that when the weight is equal to 1, there is no contribution from the full forward pass and also that a = 0 corresponds to full backpropagation through all iterations with incremental objective.
Algorithm 1 Incremental Progress Training Algorithm
Input: parameter vector 0, integer m, weight a for batch idx = 1, 2, do Sample inputs x and targets y.
Choose ~ U10, m - 1} and k ~~ (11 mmn) n)
Compute d with forward pass through n iterations without tracking gradients Compute Uprog with forward pass of additional & iterations
Compute Um with new forward pass of m iterations
Compute Limax sters with D... and Erougressvve with Uprog.
Compute (1 a)- Lamax jitrs + a Lprogressive
Compute VoL and update A
end for
4 EXPERIMENTS
We evaluate our improvements on the benchmark problems available in the Python package casy-to-harddaan The three problems considered are computing prefix sums, finding the optimal path in a two-dimensional maze, and solving chess puzzles. We briefly review the input and output structures for each problem and refer the reader to /Schwarzschild et al.) (2021a) for more detail, including the data generation process. Note that the architectures we consider are fully convolutional, and produce outputs of the same dimension as their inputs. Furthermore, the training and testing datasets we conside have labels of the same dimension as their inputs. Therefore, network trained on inputs of one size can then trivially be applied to inputs of a different size.
We begin with the toy problem of computing prefix sums modulo two. The inputs and targets for this problem consist of bit strings. The 2th bit of the target is the mod 2 sum of all bits prior and including the 7th bit of the input. We control the "difficulty" of the problem by changing the length of the bit string. Note that computing prefix sums of greater length is known to require a greater number of sequential operations (Ladner & Fischer] 1980). All of our training is done on 32-bit strings and we explore the behavior of our models on longer strings, even showing strong performance on 512-bits.
The mazes consider are two-dimensional square images where the walls are black, the permissible paths are white, and the start and end positions are denoted with red and green squares. The targets for this problem path and zeros elsewhere. We make more challenging datasets by increasing the size of the mazes. All training in work is done on 9 mazes. Despite this small training size, we benchmark
maps of the same dimension the input maze, but with ones on the optimal
5

---

Under review as a conference paper at ICLR 2022
extrapolation performance on mazes of size 59 x 59, and observe models solving mazes large as 201 x 201 with high success rates.
The chess puzzles consider are mid--amee boards represented by twelve 8 x 8 planes indicating the positions of all 12 types of pieces. (There are six distinct piece classes and two colors on a chess board.) The goal of this task is to find the best next move, and each target encodes this information in an 8 x 8 array with zeros everywhere except the origin and destination of the piece to be moved, which are populated with ones. The chess dataset is sorted by difficulty rating, as determined by Elo scores computed via human trials on Lichess.or org. Problems are sorted by difficulty and we train on the first 600K easiest puzzles and test our models on indices ranging from 700K 1.IM.
In each problem domain, we show that architectures with recall trained with the progressive loss perform best. We include ablation studies lending weight to our claim that both the architecture and the loss modifications are instrumental in boosting logical extrapolation power. For discussion on hyperparameter choice and training set-ups for all of our experiments, see Appendix [A.3] Even with tuned hyperparameters, training is sometimes unstable e.g. one out of every five maze-solving models fails to train. Code to reproduce all of our experiments publicly available
Prefix Sums Though Schwarzschild et al. (202Tb) only perform tests on strings with up to 48 bits, we use 48-bit and 5 512-bit testing sets to illustrate the extent of our performance improvements. We train models 32- strings with maximum of only m = 30 recurrent iterations (indicated in the figures as the shaded 'training regime' section), 400 convolutional filters per layer, and incremental loss weight a = 1 when the new loss is used.
Tested on 48-bits
Tested on 512-bits
100
100
Accuracy (%)
DT-Recall
Loss
DT
50
50
FF
FF- PF-Recall
Training Regime
100
200
100
200
300
400
500
Remt-Time Iterations
Test-Time Iterations
Figure 3: Comparison of our models, DT networks, and feed-forward models on 48-bit (left) and 512-bit (right) data. Our models solve more than 99% and 97% of 48-bit and 512-bit problems, respectively. Note, the curves for prior methods coincide with the x-axis on the right. Shaded regions (in all plots like this) denote + one standard error.
In Figure [3] we show that DT networks without our proposed modifications are unable to solve very long binary strings, achieving 0% accuracy on 512-bit data, while DT- models reach 100% after approximately 200 iterations. Notably, our models with recall and progressive loss do not suffer from the sharp decline in accuracy as the number of iterations increases beyond 20. The monotonic increase in accuracy with added iterations in our model is practically useful, as it allows for pre-defining large iteration number at which to terminate the recurrence rather than having to carefully choose a stopping iteration to maximize performance while avoiding degradation.
We compare our models baseline of DT models trained with 30 and 400 filters per layer, which is both wider and deeper than the models studied by SSchwarsschild et al. (20216). However, to provide fair comparison, we the same values for these parameters since we find that all architectures benefit from additional width and depth. We also compare to feed-forwards model without weight sharing of the same effective depth 30-iteration DT network.
To better understand the individual effects of our proposed approach, we perform an ablation study. On both test sets, DT-Recall networks trained without the progressive loss achieve the same final accuracy as models trained with it, but require approximately twice the number of iterations to get there, see Figure Because our proposed objective succeeds in making models iteration agnostic, solutions are often found sooner than DT models that are trained to solve 32-bit problems in 30 iterations specifically. DT models trained with the progressive loss reach approximately 3% higher peak accuracy than DT networks on 48-bit data, which is also sustained for more iterations. However, both these models (without recall) fail on 512-bit data and they suffer from overthinking.
Code he found in the supplementary material.

---

Under review as a conference paper at ICLR 2022
Tested on 48-bits
Tested on 512-bit Strings
100
100
Accuracy (%)
DT-Recall New
DT-Recall
50
50
DT m/ New Loss
DT
Training Regime
100
200
100
200
300
400
500
Test-Time Iterations
Test-Time Iterations
Figure 4: An ablation study showing the value of recall and progressive loss. Our models, which have both, outperform models with neither only one of these improvements.
Mazes As above with the prefix sum problem, we can improve performance on hard mazes by combining incremental training and recall architectures. In particular, while Figure 4 may not convince the reader that our proposed loss is critical since DT-Recall models without it perform very well, with more complex data drastic difference emerges.
We show in Figure that our models trivially solve the test set 13 x 13 mazes considered to be 'hard' in prior work (Schwarzschild et al. (0021b). Moreover, the significantly harder test set of 59 x 59 examples, our models exhibit strong logical extrapolation, while previous methods, both feed-forward and thinking systems, completely fail. Not only do our models achieve higher peak accuracy, but they do not overthink, as can be seen by the flat spans in the red curves in Figure 5 In fact, can push these systems their limits (and the limits of our hardware) and find that our models can solve nine of ten 201 x 201 mazes using more than 2,000 iterations. See Appendix A.9 for an example of an input this large.
Tested on 13x13 Mazes
Tested on 59x59 Mazes
100
100
Accuracy (%)
DT-Recall w/ New Loss
50
50
DT
FF
FF- FF-Recall
Training Regime
25
50
75
100
250
500
750
1000
Test-Time Iterations
Test-Time Iterations
Figure 5: Test accuracy on 13 x 13 (left) and 59 x 59 (right) mazes as a function of test-time iterations. The horizontal dotted lines correspond to ferd-forward models and coincide with the x-axis (0% accuracy) in the plot on the right.
Tested on 13x13 Mazes
Tested on 59x59
100
100
Accuracy (%)
DT-Rtcall w/ New Less
DT-Recall
DT w/ New Less
50
50
Training Regime
25
50
75
100
250
500
750
Test-Time Iterations
Test-Time Iterations
Figure 6: Test accuracy for a variety of models demonstrating the sensitivity the weight used the loss computation and to the presence/absence of recall.
We seek to understand the importance of recall and training with the improved optimization procedure. Since the best models in Figure 5 are DT-Recall models trained with 30 maximum iterations and weighting in the loss of a 0.01, we compare models with other combinations of design elements. In order to see how critical recall is in overcoming the overthinking problem, we show in the left-hand pane of Figure that without recall all models suffer from overthinking. The decay in each dashed curve and the stability with high iteration counts of the solid curves makes this clear. Notice, however, that testing 13 x 13 mazes does reveal large differences between D'T-Recall
7

---

Under review as a conference paper at ICLR 2022
models (with/without progressive loss). For this, shift our attention tests on 59 59 mazes, where DT-Recall models with progressive loss achieve on average 97% accuracy while DT-Recall models without progressive loss only reach an average of 83%. Interestingly, mazes, we find that the best model has a = 0.01 in the loss, much smaller weight than we use for prefix sums. More on finding the optimal value for a: can be found in Appendix A..3
Chess Puzzles We further find that using recall and the progressive loss yields notable improvement on chess puzzles. In Figure these modifications generate 4% accuracy improvement compared vanilla DT networks. Moreover, the accuracy of the DT-Recall networks is preserved as the number of iterations increases, while that of the DT networks drops to nearly 0% after 100 iterations. As with the above benchmark datasets, we justify our hyperparameter choices with ablation studies. In the right-hand plot of Figure 7 we show that either removing recall or training without the incremental progress loss will damage performance. Our best models are D'T-Recall networks with 512 convolutional filters in each layer trained with a maximum of 30 iterations and a weight for the combination of loss terms of a 0.5. Results from tests on harder puzzles are in Appendix A.4
Chess Models Tested 600K-700k
100
100
Accuracy (%)
DT-Recall
New Loss
DT-Recall
BT W. New Loss
DT
Training Regime
25
25
50
125
Test-Time Iterations
Test-Time Ittrationn
Figure 7: Test accuracy on harder puzzles. Left: Our models compared with baselines. Right:
Accuracies of an array models showing the importance of recall and our loss.
5 THE OVERTHINKING PROBLEM
Deep Thinking networks boast impressive capabilities to solve harder problems by thinking deeper, but they are prone to overthinking. In particular, we see in Section that some recurrent networks that can perform logical extrapolation may collapse entirely when performing too many iterations. One observation that can be made from the results above is that the overthinking problem seems to disappear when skip connections are added to provide networks with an uninterrupted view of the input. Another way to describe this is that models seem converge to fxed-poin solution as they iterate rather than becoming unstable. This property is desirable and using representative models trained solve mazes, we explore how robust models when we manipulate the features during the thinking process. In Appendix A.8 we present similar findings for prefix sum models.
5.1 MANIPULATING FEATURE MAPS
First, we investigate sensitivity to adding noise in the feature maps before concatenating the inputs. We examine model behavior when add Gaussian noise with mean zero and standard deviation one to the features after one iteration of maze solving. Models with recall can still solve 13 x 13 mazes when perturb the features, but models without recall cannot. See Appendix A.8 for plots, results, and further discussion.
Next, we ask whether or not the initial feature maps (after the projection layer p) carry any important information. We test this by replacing the initial feature maps with zeros after 50 iterations. this case, our DT-recall models naturally regenerate their features and recover solve the problem again, indicating that the learned algorithm is able to find solution using the input without the initial projection. See the lop-left panel of Figure 20 Appendix A.8
5.2 MANIPULATING INPUTS
With the notion that models may embody convergent process, turn attention to investi- gating how the networks determine when to stop manipulating the representation. Perhaps there is something in the feature maps output by late iterations that tells the network to stop working. We
8

---

Under review as a conference paper at ICLR 2022
test this in two ways. We perturb the input (which concatenate onto the features) after some number of iterations first subtly, then by swapping the maze with an entirely different example.
To start, change the imput by moving the end the maze two steps closer to the start of the maze. We this new input concatenated with features generated after 50 iterations of solving the original maze. In this case, see in the bottom-left pane of Figure 20 (appendix) that models with recall correct solve the new problem in very few iterations.
Test Accuracy Under Perturbation
100
Features sw appped
Accuracy (%)
50
recovered
DTT-Realll w/ New Loss
DT-Recall
final way test the hypothesis that networks are continually comparing their solution to the problem instance is by partly solving Maze "A," and then swapping the features with those ob- tained from 50 iterations of trying to solve Maze "B." In other words: if we concatenate the input problem (A) with the features from iteration 50 corresponding to a different maze (B), which maze will the network solve? Clearly, a system without recall will solve maze B. However, with recall, networks will recover and pull the features back to representing solution for maze A. Figure shows the effect of swapping feature maps.
100 Test-Time Iterations
50
150
Figure 8:
Test accuracy on 13 x 13 mazes when features are swapped after 50 iterations.
5.3 CONVERGING TO FIXED POINT
Changes in the Features
DT-acci Lan BT-Racai
10
Finally, we can study the convergence by mea- suring the change the output at each itera- tion. A decreasing change in the feature maps with each additional iteration suggests that the network manipulates the representation, mov- ing it closer to one that solves the problem, and that will hold onto this representation (or stay nearby) once it is reached. Moreover, we seek to qualitatively categorize each model type (architecture/loss pair) as convergent or non-
N/N
10
10a
10a
10
150
Test-Time Iterations
Figure 9: The change in the features when solving
13 13 mazes. Recall keeps this quantity small. convergent. To do we measure the change in the solution each iteration with norm of the difference. Figure(Jshows (ő(n) == |lón dh-11 1l1.. We see that our models appear to converge, while DT networks without recall explode, providing another view into overthinking.
6 CONCLUSION
In this work, we improve the logical extrapolation power of neural networks. We propose an architecture modification and a new loss that lead to a gigantic generalization leap from easy training data much more complex testing examples. Furthermore, we show that models avoid the overthinking trap. We test logical extrapolation with chess puzzles where the spacial dimension is consistent across difficulties and with maze solving and prefix sum computation where our models can extrapolate to larger problems well. In fact, our models use more than 2,000 iterations, the equivalent of more than 10,000 convolutional layers, to solve the largest mazes we consider.
Learning scalable processes and performing logical extrapolation are difficult tasks for most neural models, but with our architecture and loss, we demonstrate that huge leaps in complexity from training testing data can be handled without overthinking. Importantly, neural networks learn end-to-end how to perform logical extrapolation. Observing neural models that handle this important type of domain shift prompts future inquiry into online stopping conditions, through which models will determine on their own when to stop thinking and into applications to even more complex real world problems.

---

Under review as a conference paper at ICLR 2022
7 ETHICS AND REPRODUCIBILITY
The experiments in our work, on academic logical reasoning problems, indicate that our models can extrapolate their reasoning, but real-world domain shift not limited to logical extrapolation, and our paradigm might fail for numerous unforseen reasons. Thus, we caution against the assumption that our methods will translate to critical applications, which were not lested in this work.
In order to reproduce our findings, refer to the code and the supplementary material. Both publicly available, they provide reference for the exact implementations we use along with specific hyperpa- rameters and training set-ups. The data use is freely accessible part of open-source Python package called easy-to- ---ttto-rrrdl-- derddddattte data, and we use version 1.0.0, released in August 2021.
REFERENCES
Andrea Banino, Jan Balaguer, and Charles Blundell. Pondernet: Learning to ponder. arXiv preprint arXiv:2107.05407, 2021.
Marco Ciccone, Marco Gallieri, Jonathan Masci, Christian Osendorfer, and Faustino Gomez. Nais-net: Stable deep networks from non-autonomous differential equations. arXiv preprini arXiv:1804.07209, 2018.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals. Jakob Uszkoreit, and Eukasz Kaiser. Universal transformers. arXiv preprint 1807vv1807.038/1 2018
Cristobal Eyzaguirre and Alvaro Soto. Differentiable adaptive computation time for visual reasoning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12817-12825, 2020.
Karlis Freivalds and Renars Liepins. Improving the neural gpu architecture for algorithm learning. arXiv preprint arXiv:1702.08722 2017.
Felix A Gers and E Schmidhuber. Lstm recurrent networks learn simple context-free and context- sensitive languages. IEEE Transactions on Neural Networks, 1336)1133- 2001.
Alex Graves. Adaptive computation time for recurrent neural networks. arXiv:1603.08983 2016.
arXiv preprint
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines, 2014.
K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In 20/6 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, 2016.
Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger. Densely connected convolutional networks. Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4700-4708, 2017.
Herbert Jaeger. Tutorial on training recurrent neural networks, covering BPPT, RTRL, EKF and the" echo state network" approach, volume 5. GMDD-Forschungszeetrum Informationstechnik Bonn, 2002.
Lukasz Kaiser and Ilya Sutskever. Neural gpus learn algorithms. arXiv preprint 1rXi0:15 2015.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv: 1412.6982, 2014.
Richard Ladner and Michael Fischer. Parallel prefix computation. Journal of the ACM (JACM), 27(4):831-83) 1980.
Kwwabena Nuamah. Deep algorithmic question answering: Towards a compositionally hybrid AI for algorithmic reasoning. arXiv preprint 1rXiv:2109.08000 2021.
10

---

Under review as a conference paper at ICLR 2022
Rasmus Berg Palm, Ulrich Paquet, and Ole Winther. Recurrent relational networks. In 32nd Confer- ence on Neural Information Processing Systems, pp. 3368 2278. Neural Information Processing Systems Foundation, 2018.
Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108. /2400, 2021.
Jürgen Schmidhuber. Self-delimiting neural networks. arXiv preprint arXiv: 2012.
Jürgen Schmidhuber, Daan Wierstra, Matteo Gagliolo, and Faustino Gomez. Training recurrent networks by evolino. Neural computation, 19(3)7557777, 2007.
Avi Schwarzschild, Eitan Borgnia, Arjun Gupta, Arpit Bansal, Zeyad Emam, Furong Huang, Micah Goldblum, and Tom Goldstein. Datasets for studying generalization from easy to hard examples, 2021a.
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, 2021b.
Avi Schwarzschild, Arjun Gupta, Amin Ghiasi, Micah Goldblum, and Tom Goldstein. The uncanny similarity of recurrence and depth, 2021c
Daniel Selsam, Matthew Lamm, Benedikt Bünz, Percy Liang, Leonardo de Moura, and David Dill. Learning sat solver from single-bit supervision. arXiv preprint
2018.
Rupesh Kumar Srivastava, Klaus Greff, and Jürgen Schmidhuber. Highway networks. arXiv preprint 1055v0150 2015.

---

Under review as a conference paper at ICLR 2022
A APPENDIX
A.1 ARCHITECTURES
The architectures we use in our experiments are very similar to those described by Schwarzschild al. (202Ib), but we provide the details here for convenience.
In all the models, the first layer p is convolutional layer with 3 3 filters (or filters of length three in the 1D case) followed by ReLU non-linearity, which projects the input into feature space. Each filter strides by one, and inputs are padded with one unit in each direction. The number of filters specified per dataset in Table ass the "width" of the network. The internal blocks standard residual blocks with four convolutional layers (followed by ReLUs) and skip connections every two layers (He et al., 2016). These blocks share weights in recurrent networks and are simply repeated with distinct parameters in feed-forward models. Models with recall have additional convolutional layers that map the concatenated inputs and features [o, x| the shape of the feature maps ø. The feature maps have the same spacial dimension as the input and w channels, where denotes the width of the model. The final block widths (specific numbers are in the Table T). and ReLUs after the first two. The third and final convolutional layer in / has two channel outputs used for binary pixel classification.
composed of three convolutional layers with decreasing
Table 1: Model architecures for all main experiments. Note, we perform ablations where we change the width or the maximum number iterations and those parameters are indicated where appropriate.
Dataset
Width
Channels in h layers
Prefix Sums
400
400, 200,
Mazes
128
32, 8, 2
Chess
512
32, 8,
A.2 THE LOSS
Prefix sum problems with bits are input to the model as vector E 10, and the output is denoted by ÿ E R=x2 The target output is y E 10". We compute the loss as follows.
n ((..)) = n
[1-36)
log
(3)
el + c
where [ indexes the output array. Therefore, for batch of B instances, the total loss is
B Σ L = B b=
E(G,, 9b)
(4)
This loss applies to all three problem types we consider. Note that maze represented by m X m input has up there are always 64 pixels.
pixels and the loss can be averaged over those pixels. The same applies to chess, where
A.3 HYPERPARAMETERS
We describe the training set-up for each experiment and discuss hyperparameter choices.
Due to instability training for prefix sum networks, all models are optimized using Adam (Kingmal & Ba 2014) and gradient clipping. Maze models also trained with Adam, but clipping used. Chess models train stably with SGD without gradient clipping. All training is done with a weight decay coefficient of 0.0002 and training with SGD uses momentum coefficient of 0.9.
In each training run, we hold out 20% the training data to use as a validation set, and we save for testing the checkpoint with the highest validation accuracy. Additionally, when compute averages,
12

---

Under review as a conference paper at ICLR 2022
we only include models that trained beyond threshold training accuracy. Our threshold choices for 90% prefix sums and for chess and 99% for mazes.
Coarse experimentation with learning rates and decay schedules is used to determine optimal choices of these hyperparameters for reproducibility and speed of training. Models were trained with exponential warm-up for small number of epochs before the main training routine. All DT models presented in are trained using early stopping to combat overfitting network parameters are saved for model with highest accuracy on hold-out set during training. Feed-forward networks are tested with both early stopping and full epoch training, and the best performing class of model is taken as the baseline.
In Table 2 we present the training hyperparameters shared among all models presented in Section for each problem instance.
Table 2: Training hyperparameters. Dashes indicate that we did not utiltize those options.
Taak
Optimizer
Leammnn Rate
Decay Schodule
Decay Factor
Warm Up Epochs
Epocts
Clip Cocflicient
Prefis Sums
Adam
0.001
[60, 100)
0.01
10
150
100
Mares
Adam
0.001
10
50
Mazes (FF)
[175]
200
Chess Puzzles
SGD
0.010
[100, 110
0.01
120
We perform width and depth experiments to demonstrate that our best prefix sum models are, roughly speaking, on a plateau with respect to performance on 512-bit data in hyperparameter space. We see from the right panel of Figure 10 that performance increases steadily with depth, but no additional performance is achieved by increasing m past 30. From the left panel of Figure TO, we also see performance steadily increases with network width. We choose to use width of 400 filters per layer and m = 30 because this adequately solves 512-bit data, and it is memory inefficient to use any more for the task at hand.
Tested on 512-bit Strings
100
100
Accuracy (%)
Width 400
30
50
Width 300
50
m - 25
Width 200
m 20
Width 100
m 15
Training Regime
Training Regime
100
200
300
400
100
200
300
400
Test-Time Iterations
Test-Time Iterations
Figure 10: Experiments to determine optimal depth and width for prefix sum models trained with
progressive loss (a = 1) and recall.
These experiments are performed on prefix sums, which are relatively memory and computationally inexpensive, to exemplify the general correlation between depth, width, and performance observed on all tasks. The choice of depth and width of maze and chess models is also made to be only as large as necessary to avoid excessive memory and computational costs.
We determine the optimal weight to be used it the combination of the loss terms using a coarse grid search. For prefix sums we test models with a values in [0, 0.25, 0.5, 0.75, 1]. for mazes we test values [0, 0.01, 0.1, 0.5, 1]. and for chess we test values [0, 0.5, 1]. We choose the best weight for each task to present in the main experiments.
For evaluating recurrent models, Schwarzschild et al. 2021b) propose a maximum confidence rule. This algorithm takes the output from the iteration at which the outputs have the highest confidence. This is measured by averaging the per pixel confidence. where confidence is computed as the maximum of the softmax over the output channels. In order to fairly compare our recurrent models to theirs, we use the same evaluation technique. For more detail see Appendix AA.6
A.4 EXTENDED RESULTS
Tables 3 and 5 show the peak accuracy of each of the curves presented in the plots in Section
13

---

Under review as a conference paper at ICLR 2022
Texted on 1824-bits
100
DT Reoll
New Loss
IVT
Training Repitne
20
100
200
300
400
500
600
700
800
969
1000
Tea-TTime Iterations
Figure 11: Prefix sum results on 1024-bit inputs.
Tedee 201x201 Mores
50
Aucunacy (%)
60-
Accuracy
40
Num Pixcle Connect
Training Regime
1000 2000
4000
6000
Tea-Time Berstions
Figure 12: Accuracy curve from a single representative model on 201 x 201 mazes. On this plot, we add curve to indicate what percentage of pixels that are correct to draw attention to the fact that the seemingly low performance of 74% accuracy should understood in context. Of the 100 mazes in the test set, 20 have one pixel (out of 161, ,604) wrong. No single maze has more than seven mistakes.
Chess Models Tested 700K-800K
100
100
Accuracy (%)
DT-Recall
New Loys
DT-Recall
DT W/ New Loss
DT
Training Regime
25
50
100
25
50
100
Test-Time Iterations
Test-Time Iterations
Figure 13: Chess performance when tested puzzles with indices from 700K to 800K.
Chess Models Tested on IM-1.1Mk
100
100
Accuracy (%)
DT-Recall
New Lass
DT-Recall
DT W/ New Loss
DT
Training Regime
25 50
25
50
100
Test-Time Iterations
Test-Time Iterations
Figure 14: Chess performance when tested on puzzles with indices from 1M to 1.1M.
14

---

Under review as a conference paper at ICLR 2022
Table 3:
The peak accuracy and corresponding test-time iteration number for prefix sum solving model performance curves in Figures 3 and
Tested on 48-bit Strings
Tested 512-bit Strings
Model Peak Iter.
Peak Acc. (%)
Model
Peak Iter.
Peak Acc. (%)
DT 0.0
42
94.61 1.19
DT
0.0
0.00 = 0.00
DT
97.73 = 1.80
DT 1.0
171
11.26 = 6.90
DT-Recall 0.0
99.97 = 0.01
DT--ecaall 0.0
466
96.19 = 3.73
DT-Recall 1.0
26
99.96 = 0.02
DT-Recall 1.0
237
97.12 = 1.88
FF 0.0
30
27.15 2.56
FF 0.0
30
0.00 = 0.00
FF-Recall 1.0
99.87 0.04
FF-Recall 1.0
30
0.00 + 0.00
Table 4: The peak accuracy and corresponding test-time iteration number for maze solving model performance curves in Figures 5 and 6
Tested 13
Mazes
Tested on 59 x 59 Mazes
Model
Peak Iter.
Peak Acc. (%)
Model
Peak Iter.
Peak Acc. (%)
DT 0.00
40
85.59 + 2.81
DT 0.00
0.00 0.00
DT 0.01
86.08 3.96
DT 0.01
0.00 0.00
DT-Recall
99.94 0.03
DT--Recall
0.00
999
82.72 = 15.14
DT-Recall 0.01
70
99.88 + 0.05
DT--Realll 0.01
984
97.30 0.68
FF 0.00
38.22 + 5.28
FF 0.00
0.00 = 0.00
FF-Recall 0.01
43.01 5.14
FF-Recall 0.01
0.00 0.00
Figures 13 and [4 show the test performance of chess models on even harder test sets. is interesting
to note that recall and our loss both still help dramatically, but there is apparent ceiling generalizing from the easy puzzles to much harder ones. We leave further investigation into this phenomenon for future work.
A.S IN-DISTRIBUTION RESULTS
Information on how models perform in distribution reveals that each class of model fits the training data and generalizes the classical sense. Specifically, we report the accuracy on held-out validation set from the same problem difficulty the training data. For recurrent models, we evaluate the performance using the maximum number iterations used during training, irrespective of the training objective. See Tables and 8
15

---

Under review as a conference paper at ICLR 2022
Table 5: Peak accuracy and iteration number for chess puzzle performance curves in Figure] 7
Tested puzzles 600K-700K
Model
Peak Iter.
Peak Acc. (%)
DT
0.0
78.81 = 1.04
DT
0.5
43
82.32 = 0.10
DT--Recall
0.0
29
82.12 + 0.69
DT-Recall
0.5
57
82.69 + 0.27
FF
0.0
30
61.17 = 1.76
Table 6: Training accuracy and in-distribution validation accuracy for models presented in Figures
and
Trained on 32-hit Strings
Model
a
Train Acc. (%)
Val Acc. (%)
DT
0.0
99.98 0.01
100.00 0.00
1.0
99.57 = 0.38
97.73 = 1.80
DT-Recall
0.0
99.95 1 0.02
100.00 + 0.00
DT-Recall
1.0
99.98 0.01
100.00 0.00
FF
0.0
99.76 0.14
99.76 0.14
FF-Recall
1.0
100.00
1 0.00
100.00
Table Training accuracy and in-distribution validation accuracy for models presented in Figures
and
Trained on 9 x Mazes
Model
a
Train Acc. (%)
Val Aec. (%)
DT
0.0
100.00 = 0.00
100.00 + 0.00
DT
0.01
99.98 + 0.01
99.98 + 0.01
DT-Recall
0.0
99.91 + 0.07
99.92 + 0.06
De-Recall
0.01
99.77 +
99.74 1
FF
0.0
99.94 1 0.01
99.93 1 0.05
FF-Re
0.01
99.99 = 0.00
99.99 + 0.00
Table 8: Training accuracy and in-distribution validation accuracy for models presented in Figures
Trained on puzzles 0-600K
Model
a
Train Ace. (%)
Val Acc. (%)
DT
0.0
99.98
0.02
92.94 0.34
DT
0.5
99.90 = 0.01
94.16 = 0.02
De-Recall
0.0
99.77 + 0.02
100.00 = 0.00
DT-RRcall
0.5
99.31 +: 0.02
91.18 + 0.07
FF
0.0
96.89 + 0.58
83.21 + 1.22
16

---

Under review as a conference paper at ICLR 2022
A.6
EXIT RULES
When DT network processes a single input, it can produce output at each iteration. When we choose number of test-time iterations M, we are specifying the maximum number of iterations, but we may choose from any of those to extract single output. A naive approach is to take that output iteration M; we call this We call this the defoult exit rule. A slightly more complicated process, and what [Schwazschilld et al. (2021b) do, is to take the output from iteration m² which solves the following maximization problem. Let gi(T) be the output from iteration m, and let (/m)/k, be the (k, entry in the output that corresponds to the confidence that the h pixel is in class j. finally let k E {1, 2, ..., K} and let E {0,1) l).
m arg max Σ m K
1 max(((m))) (k.
(5)
Choosing the output using this maximization technique called the maximum confidence exit rule.
In order to stay faithful to the prior work, all plots in the main body of this paper use the maximum confidence exit rule (Schwarzschild al.] (021b). In this section, compute accuracy at iteration by using the output after n iterations instead of taking the highest confidence output at all iterations up n. We call this the default exit rule.
Tested 48-bits
Tested 512-bits
100
100
Accuracy (%)
DT-Recall FF
50
50
FF-Recall
Training Regime
100
100
200
300
400
500
Test-Time Iterations
Test-Time Iterations
Figure 15: Test accuracy 48-bit (left) and 512-bit (right) prefix sums as function of test-time iterations using the output at each iteration. The horizontal dotted lines correspond to feed- forward models and coincide with the x-axis (0% accuracy) in the plot the right. Shaded regions denote + one standard error.
Tested on 13x13 Mazes
Tested 59x59 Mazes
100
100
Accuracy (%)
DT-RRccll w/ New
50
50
FF FF-Recall Training Regime
25
50
75
100
250
500
750
1000
Test-Time Iterations
Test-Time Iterations
Figure 16: Test accuracy 13 13 (left) and 59 59 (right) mazes function of test-time iterations using the output at each iteration. The horizontal dotted lines correspond to feed-forward models and coincide with the x-axis (0% accuracy) in the plot the right. Shaded regions denote + one standard error.
To better understand the evolution of confidence as the iteration count increases, we plot this in Figure 18 Note the confidences are high in the early iterations even before the models are correct.
A.7
HARD TO BASY
One natural query about our models is how well they perform on test sets that are easier/smaller than the data used for training. Figure 19 shows that recurrent models can generally solve casier/smaller problems in fewer iterations. Even though they solve the easier problems in fewer than the maximum number of iterations used in training, we still see that models without recall suffer from overthinking.
17

---

Under review as a conference paper at ICLR 2022
Chess Models Tested 600k-700k
100
Accuracy (%)
Ours (Max-Conl)
Ours (Default)
50
DT
Training Regime
0
25
50
75
100
125
Test-Time Iterations
Figure 17: Test accuracy on chess puzzles with indices 600K 700K function of test-time iterations. Exit
Prefix Sums Tested 512-bil Strings
Mases Tested on 59s59
Chess Tested 600-70%
100
100
100
Coufidence (7C)
Confidence (Mp)
Confiiencee
60
100
200
300
400
500
150
300
450
600
750
900
25
75
100
125
Tese-Time Iterationns
Tes-TTime Iterations
Test-Time Iterations
Figure 18: Evolution of confidence of representative models for each dataset as iteration count increases. Left: prefix sums, center: mazes, right: chess.
A.8 MORE THE OVERTIINKING PROBLEM
The experiments we present in Section shed light on how some models overcome the overthinking problem. Here, we present additional plots and tables to give a more complete picture of our findings. Table shows the results from several experiments. The first column indicates the average number of iterations to solve maze and the best accuracy (over all the test iterations considered). Then, we examine the behavior when the features are perturbed after the first iteration by adding mean zero and variance one Gaussian noise ("Noise' column) and we find that this completely destroys the models without recall while only slightly slowing down those with recall. Next, we replace the features at the first iteration with arrays of zeros ('Zeros' column). Finally, we swap the features after 50 iterations with those obtained from solving entirely different maze, and again the models without recall cannot recover from this, but those with recall achieve near perfect accuracy, Figure 20 shows accuracy curves for these experiments. The bottom right plot in Figure 20 shows the norm of the difference in features when we input random noise instead of actual mazes. This plot shows that the stable behavior of our models not limited the portion of input space occupied by real mazes, rather the process they learn exhibits convergent behavior even on random inputs.
the same trend when
Table 9: Average iterations to solution, with the peak accuracy in parentheses. We evaluate maze solving models on 13 x 13 mazes, where we intervene in the solving process different ways.
Model
Clean
Noise
Zeros
Swapped
DT
27.68 (87 %)
(0%)
(0%)
(0%)
DT w/ New Loss
16.87 (95 %)
(0%)
(0%)
(0%)
De-Recall
23.54 (100 %)
35.70(100 %)
35.04 (96 %)
33.92 (97 %)
DT-Recall w/ New Loss
22.10 (100 %)
25.10(100 %) 22.48 (100 %)
28.79 (100 %)
Similar experiments on prefix sum computation reveal that recurrent models with recall seem to converge nicely well. Note, the top two plots in Figure 21 come from identical experiments as those defined above. Whereas with mazes we shortened the optimal path, in this setting we flip single bits in the input string. We explore the response of the model to flipping bits at different indices after 50 iterations (when the models have already solved the initial 48-bit problem). In the lower left plot in Figure 21 see that the recovery time function of which bit flipped reveals two things: (i)
18

---

Under review as a conference paper at ICLR 2022
Models Tested 24-bits
100
Accuracy (%)
50
Durs
DT-Rrcall
DT W/ New Loss
Training Regime
0
25
50
75
100
Test-Time Iterations
Figure 19: Accuracy as function of iteration for prefix sum models when generalizing from hard (32-bit strings) easy data (24-bit strings).
Noise Added To Features
Zeros Inserted After Iter. 50
100
100
Accuracy (%)
Recuraa Sereees (%)
DT-Recall
De--le
50
DT-Recall
DT
DT W/ New Loss
DT-Recall
%
50
100
150
0
50
100
150
Test-Time Iterations
Test-Time Iterations
Maze Shortened After Iter. 50
Random Noise Input
100
DTRi Ried
10
Accuracy (%)
hen
10
D7
DT- Recall
50
w/ New Loss
10
10a
0
50
100
150
10"
100
150
Test-Time Iterations
Test-Time Iterations
Figure 20: 13 13 mazes.
our models can recover from single bit flips, and (ii) the recovery time decays linearly with the index of the bit flipped. Higher indices are closer to the end and affect fewer bits in the output than lower ones.
19

---

Under review as a conference paper at ICLR 2022
Noise Added Features
Zeros Inserted After Iter. 50
100
100
Accuracy (%)
Accuracy (%)
D'TRRecall w/ New Loss
DT-Recall
DT-Recall w/ New Loss DT-Recall
DT
%
50
100
150
0
50
100
150
Test-Time Iterations
Test-Time Time Iterations
Iters Recover
Convergence Prefix Sum Models
Avg. Iters Solution
30
10 14
N/Now Tew
DT Roul
10
D7
20
B7 Witten Lan
10
DT-Rocall DT-Recall
10"
20 Index of Perturbed Bit
40
50 Test-Time Time Iterations
100
Figure 21: 13 13 mazes.
20

---

Under review as a conference paper at ICLR 2022
A.9 A FUN MAZE
Figure 22: Fun for the whole family! We leave solving this 201 201 maze as an exercise for the reader. The green starting dot is near the center of the upper right quadrant. The red ending point is on second row from the bottom.
21


Reviews:
[{"review_id": "rxcSl-nCEL5", "reviewer": "Reviewer_ue13", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n - The paper clearly states the problem, and convincingly shows a solution.\n - The paper is pretty well written\n\nWeakness\n - The paper is evaluated on a very new dataset, which seems ideal for the kind of method suggested, and the only real comparison is the work they directly build on. You can always find a dataset for which your method is the best, and this paper has a bit of that feeling. Except for the chess positions the examples are new and toy-like. Why do we need these new benchmarks? Why not evaluate on some more standard benchmarks, e.g. image classification/segmentation, simple QA (bAbI for instance), etc. \nAlso the bit sums and maze path finding are trivially solved with standard hand-coded algorithms. Why would you want to use a neural net on them? If the benchmarks are not interesting problems in their own right, they might still be nice benchmarks because they nicely exemplify some specific problematic environments, but then it must be very clearly explained how this ties back into (better) solving real-world problems. The chess positions task is the nicest example problem, but then again, how does this compare to something like AlphaZero, or any other modern MCTS approach? What's the benefit? Computational? Generalization? This must be explicitly explained.\n\n - The paper doesn't discuss and compare to the very relevant work \"Recurrent Relational Networks\" (RRN) [1]. RRNs are recurrent graph neural networks, and very very similar. Both neworks have the same output(recurrent_with_recall(embed(x))) structure and the RRN measures a loss on every step of recurrent computation similar to the modified loss. The RRN also learns a convergent algorithm and generalizes to harder tasks. I would very much appreciate a comparison to a RRN network in all of the new benchmarks, or conversely evaluating on some of the same benchmarks as in [1]. I suspect the performance will be very similar for the two networks.\n\n - The loss weight alpha hyperparameter is kind of inelegant. Have the authors tried simply measuring a loss on every recurrent iteration like in [1]? This would also encourage a convergent algorithm, and would be simpler.\n\n[1] - Palm, R. B., Paquet, U., & Winther, O. (2018). Recurrent Relational Networks. In 32nd Conference on Neural Information Processing Systems (NeurIPS 2018), Montréal, Canada (Vol. 31, pp. 3368-2278).", "summary_of_the_paper": "The paper proposes two modifications to recurrent neural networks that enable them to extrapolate to larger problems than seen during training.\n\nThe problems (generalizations) are \n1) pathfinding in a maze (larger mazes), \n2) binary prefix sums (larger bit strings)\n3) Evaluate the best chess move given a position (harder positions).\n\nThe modifications are: \n\n1) Add the initial problem features to every step of the recurrent computation, coined \"recall\" in the paper.\n2) train on a combination of a regular loss with m recurrent iterations and with a loss with m=n+k iterations where the gradient is not tracked for the first n iterations.\n\nThe paper shows that the modified recurrent network learns an algorithm that converges to the correct result with more recurrent iterations, thus overcoming the problem of \"overthinking\" (divergence really) as coined in earlier papers. \n\nThe paper further shows that the learned convergent algorithm \"extrapolates\" to larger/harder problems. For instance, the networks are trained on 32-bit strings, and evaluated on 512-bit strings, and trained on 9x9 mazes and evaluated on up to 201x201 mazes.\n\nThe paper compares their method, ablations w.o. \"recall\" and modified loss, MLPs and MLPs with \"recall\", and show that only their proposed network learns a convergent solution that extrapolates to the larger/harder problem instances.", "summary_of_the_review": "Decent paper, but should be evaluated on more common benchmarks / less toy settings and should discuss and compare to [1].", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "cxt6KnbZRO6", "reviewer": "Reviewer_w7Ak", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is well written and easy to read and understand. It builds directly on recent work by Schwarzschild et al. [2021a, 2021b, 2021c], and is essentially about addressing the limitations of that work on the challenging tasks identified in that work that involve learning from \"easy\" instances of a task and generalizing to \"hard\" instances. This is an interesting line of work that could lead to useful insights, and this paper's contributions certainly to make significant progress on the challenges considered.\n\nIn particular, the authors have already conducted various analyses to answer follow-up questions that a curious reader might have (and I had) about the working of the proposed techniques. The results of these experiments will be very useful and instructive to readers interested in these problems.\n\nI would like to point out two \"weaknesses\" of this paper:\n\n1. The first out of the proposed two techniques, the Recall skip connection, appears to be fixing what one may call a mistake made by Schwarzschild et al. in recent work during model design. This is because if one looks at prior literature on models that are stable when unrolled longer than training time, one immediately find this model design, based on theoretical insights from control theory, see Ciccone et al. [A] and the non-autonomous networks therein. This is not to say that the contribution of this paper isn't useful, but it should be put in context of past work that has already established strong theoretical results directly related to the behavior of interest. I find it interesting that in this paper, no conditions on weights were necessary in practice to achieve stability.\n\n2. I'm concerned that the authors are playing fast and loose with the term \"thinking\" in this line of work. This is very loaded metaphor, and in my opinion every paper that decides to use it must at least include a clear note that \"thinking\" is just a fancier term for a particular type of processing here, and these networks are not really thinking in any regular sense of the term. Otherwise, we can say that any classifier or detector is also thinking. I realize that this is not just something unique to this paper and some other papers have used such terms without such clarification, but this is only a review of this paper, and I'd appreciate if the authors think about this issue and try to do better than prior work.\n\n[A] Ciccone, M., Gallieri, M., Masci, J., Osendorfer, C., & Gomez, F. (2018). Nais-net: Stable deep networks from non-autonomous differential equations. arXiv preprint arXiv:1804.07209.", "summary_of_the_paper": "This paper is focused on training networks to solve problems via extrapolation, e.g. solving large mazes by learning to solve small mazes. In order to achieve this, two modifications to recurrent convolutional networks are proposed: a) adding a concatenating skip connection from the input to the recurrent layers that stabilizes extrapolation (Fig. 2), and b) using a modified training method and loss that encourages the network to learn computations independent of current iteration (Alg. 1). Together, these techniques lead to substantial improvements in results on the problems of computing prefix sums, solving 2D mazes and solving chess puzzles.", "summary_of_the_review": "The paper proposes simple techniques that clearly addresses the limitations of recent work on solving tasks in the easy-to-hard benchmark dataset.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "REju5-0RXAs", "reviewer": "Reviewer_Ap3D", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Both the underlying approach, and the proposed extensions, are compelling, and the results are promising. It is good to see results across a range of task domains. The analyses on convergence and overthinking in the latter part of the paper are also very nice. I think there are a few issues that need to be sorted out, but am happy to raise my score if the authors can address these concerns:\n\n- The only baseline considered is a feedforward version of the primary model (i.e. in which parameters are not shared across iterations). This seems like a good comparison, but other baselines should be considered as well. How important is the particular recurrent architecture employed here, in which the output is fed back into the model as input? How might an LSTM, with a recurrent hidden state, or a model with an external memory, perform on the generalization benchmarks that are considered in this paper?\n- The proposed 'progressive loss' seems to be useful primarily for getting these kinds of models to solve the task in the shortest number of iterations of possible. The reason seems to be that a model trained only on a fixed number of iterations has no reason to arrive at the correct answer any sooner than the point at which a loss will be computed. Given that, I wonder whether the proposed method is the simplest or best way to accomplish this. Would it work to simply train models on a randomly sampled number of iterations, or to penalize longer processing (as is done in 'adaptive computation time')? It would be good to compare the proposed approach to these alternatives.\n- In the appendix, the authors state that 'when we compute averages, we only include models that trained beyond a threshold training accuracy.' Does this mean that the reported results are only for a subset of models that reached some training criterion? How many models failed to reach this criterion, and are the results qualitatively different when no such criterion is used?\n- Models are evaluated by taking the iteration with the highest confidence, rather than simply using the final iteration. Do the models still perform just as well at generalization to more complex problems if the final iteration is used? \n- On the prefix sum and maze problems, the 'recall' element seems to be useful primarily for preventing the 'overthinking' phenomenon, and, for those task domains, this element of their proposal appears to be highly effective (as it prevents overthinking even without the 'progressive loss'). However, on the chess problems, the version of the model with recall, but without the 'progressive loss', still suffers from overthinking. What might explain this discrepancy?\n\nMinor issues and additional questions:\n\n- How does confidence evolve over time? Do the models generally become more confident with a greater number of iterations? If so, could confidence be used as a signal for autonomously selecting how many iterations to perform, rather than having to select this by hand?\n- When reading the paper for the first time, it is unclear where the results in Figure 1 come from. I'm assuming these results employ both of the proposed extensions, but it would be helpful to specify that this is the case (i.e. that 'vanilla' thinking systems would not be capable of producing these results).\n- The legend for table 1 of the appendix says 'perfrom' instead of 'perform'.\n- The legend for table 4 in the appendix references the prefix sum task, but the table appears to contain results from the maze task.\n- At the end of the main section with results on the maze task, the reader is directed to section A.3 to see an example of a 201x201 maze, but this is actually in section A.7.\n\n\n", "summary_of_the_paper": "This paper proposes two extensions to the recently proposed recurrent 'thinking systems', in order to enable them to better generalize from training on simple problems to testing on more complex problems. The proposed extensions involve 1) giving the model access to a cure indicating the to-be-solved problem at each time step ('recall') and 2) a method intended to prevent the model from learning behaviors specific to particular iterations (so as to enable generalization to more complex problems via a larger number of iterations). When combined, these extensions enable generalization to significantly more complex problem instances across three separate task domains.", "summary_of_the_review": "The proposed approach is compelling, and the results are promising, but there are a few issues that needed to be sorted out, including some additional baselines, and clarification of the selection criteria for the reported results.\n\nUpdate after discussion period: the authors included a number of supplementary results and informative additional control experiments. I still think that the results would be more convincing if compared to a broader range of competitive baselines, but I think these new results/experiments are a significant enough improvement to merit a score increase.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "-D3FBEmipBO", "reviewer": "Reviewer_ajNt", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n- The paper gives solid extensions to the idea of deep thing networks (Schwarzschild et al., 2021). The writing is generally clear and easy to follow. The modifications are well-motivated,  addressing the weaknesses of the original model.\n- The experiments demonstrate well the effectiveness of the modifications, outperforming the baselines by a significant margin. There are detailed ablation studies and behaviour analyses that give insights into the inner working of the model.  \nCons:\n- The modifications are incremental, using common techniques like residual connection and truncated backpropagation through time. Unsurprisingly, using these techniques will help improve the performance. \n- The tasks are toyish. It would be more persuasive if the method could work for realistic data or some classical benchmarks. \n- The main manuscript misses some details of the method (see questions below).\n\nOverall, I like the idea of this paper, which proposes an interesting way to exhibit logical generalization by allowing recurrent networks to \"think\" more during inference. However, as mentioned above, compared to the original work, the proposed modifications are marginal. It is straightforward that using residual connections is critical to training deep networks (to combat missing input information or gradient vanishing, etc.).  The progressive loss is more interesting, yet it is unclear from the writing that this is a novel contribution to the training of RNNs (e.g., compared to other auxiliary training for RNN as in Trinh et al., (2018), what is the difference and advantage of using the proposed method). \n\nRegarding experimental results, only Chess Puzzle data seems real. But here, the result is not impressive as there is only a 4% improvement in terms of peak accuracy. The performance of the proposed method on the other two synthetic data is promising.  I am curious whether the method only works for these specific tasks or is extendable to other problems: sequential inputs such as Copy, Associative Recall, graph reasoning (Graves et al., 2016) or natural data (text, images).\n\nQuestions and comments:\n\n- It is reasonable to use CNN architecture for image-like inputs. However, for prefix-sum, the input is a sequence of 0 and 1. Did you apply CNN here? \n- What is the loss function? To make the paper self-contained, in Sec. 3 or 4, you should describe the output format and the loss used to train the model (from Appendix A, it seems that the final output is a 2d bitmap?).\n- The procedure of selecting the inference iteration for performance measurement should be mentioned in the main manuscript. \n- The analysis would be more informative if the authors could compare the results of other iterations with that of the peak one (how much difference? Did they converge to the final output as suggested in Fig. 9?)\n- For prefix-sum, did you test with longer sequences? Can your method generalize to 1024-bit? Why there is no performance visualization for maze 201x201?\n- In the paper, the feed-forward model is a weak baseline. The experiments could be stronger with advanced architectures such as Transformer or self-attention model.  \n- Fig. 8, More explanation on why the new loss helped more in the setting will be appreciated.\n- Fig. 9, the y-axis values go up to 10^10. How's that possible? How does the new loss help in this case?", "summary_of_the_paper": "The paper proposes two modifications to recurrent neural networks that help improve generalization on three synthetic tasks. In particular, the authors implement a recall mechanism through residual connections to prevent the recurrent network from losing original input information after many iterations. In addition, they randomly apply truncated backpropagation through time to the computation graph to build an auxiliary loss, which diminishes overfitting to the number of recurrent steps. The trained network can solve significantly harder problems by performing more recurrent iterations, showing strong generalization ability.\n", "summary_of_the_review": "Overall, I like the idea of this paper, However, compared to the original work, the proposed modifications are marginally significant.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This is an interesting work, and I urge the authors to keep pushing this direction of research. Unfortunately, I feel like the manuscript, in its current format is not ready for acceptance.

The research direction is definitely under-explored, which makes the evaluation of the work a bit tricky. Still I think that some of the points raised by the reviewers hold, for e.g. the need of additional baselines (to provide a bit of context for what is going on)I understand that the authors view their work as an improvement of the previously proposed DT network, however that is a recent architecture, not sufficiently established not to require additional baseline for comparisons. This combined with the novely of the dataset makes it really hard to judge the work. 

The write-up might also require a bit of attention. In particular it seems a lot of important details of the work (or clarifications regarding the method) ended up in the appendix. A lot of the smaller things reviewer pointed out the authors rightfully so acknowledged in the rebuttal and propose to fix, however I feel this might end up requiring a bit of re-organization of the manuscript rather that adding things at the end of the appendix. I also highlight (and agree) with the word "thinking" being overloaded in this scenario.

Ablation studies (some done as part of the rebuttal) might be also a key component to get this work over the finish line. E.g. the discussion around the progressive loss. I acknowledge that the authors did run some of those experiments, though I feel a more in depth look at the results and interpretation of them (e.g. not looking just at final performance, but at the behaviour of the system), and integrating them in the main manuscript could also provide considerable additional insight in the proposed architecture. 

My main worry is that in its current format, the paper might not end up having the impact it deserves and any of the changes above will greatly improve the quality and the attention the work will get in the community.
Decision: Reject