End-to-end Algorithm Synthesis with Recurrent Networks: Extrapolation without Overthinking
Authors: Arpit Bansal
Avi Schwarzschild
Eitan Borgnia
Zeyad Emam
Furong Huang
Micah Goldblum
Tom Goldstein
Venue: NeurIPS 2022 Accept (2022)
Keywords: Extrapolations
algorithm synthesis
recurrent networks
algorithmic reasoning

Abstract: Machine learning systems perform well on pattern matching tasks, but their ability to perform algorithmic or logical reasoning is not well understood. One important reasoning capability is algorithmic extrapolation, in which models trained only on small/simple reasoning problems can synthesize complex strategies for large/complex problems at test time. Algorithmic extrapolation can be achieved through recurrent systems, which can be iterated many times to solve difficult reasoning problems. We observe that this approach fails to scale to highly complex problems because behavior degenerates when many iterations are applied -- an issue we refer to as "overthinking." We propose a recall architecture that keeps an explicit copy of the problem instance in memory so that it cannot be forgotten. We also employ a progressive 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. These innovations prevent the overthinking problem, and enable recurrent systems to solve extremely hard extrapolation tasks.

End-to-end Algorithm Synthesis with Recurrent Net- works: Logical Extrapolation Without Overthinking
Arpit Bansal" University of Maryland bansal01@umd, edu
Avi Schwarzschild" University of Maryland tvi1@umd. edu
Eitan Borgnia University of Chicago
mnorgnia@uchicago edu
Zeyad Emam University of Maryland
Furong Huang University of Maryland
Micah Goldblum New York University
Tom Goldstein University of Maryland
Abstract
Machine learning systems perform well on pattern matching tasks, but their ability perform algorithmic logical reasoning not well understood. One impor- tant reasoning capability is algorithmic extrapolation, in which models trained only small/simple reasoning problems can synthesize complex strategies for large/complex problems at test time. Algorithmic extrapolation can be achieved through recurrent systems, which can be iterated many times to solve difficult reasoning problems. We observe that this approach fails to scale to highly complex problems because behavior degenerates when many iterations are applied an issue we refer to as "overthinking. We propose recall architecture that keeps an explicit copy of the problem instance in memory so that cannot be forgotten. We also employ progressive training routine that prevents the model from learning behaviors that are specific to iteration number and instead pushes iors that can be repeated indefinitely. These innovations prevent the overthinking problem, and enable recurrent systems to solve extremely hard extrapolation tasks.
learn behav-
LCS
Figure 'thinking' network trained on 9 9 mazes and their solutions (left) autonomously synthesizes scalable algorithm. By rumning this algorithm for longer, it reliably solves problems of size 59 59 (middle). 201 x 201 (right), and much larger (appendix) without retraining. Standard architectures, and even existing primitive thinking models, fail to tolerate this domain shift.
Equal contribution.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Introduction
Humans solve complex logical reasoning problems through logical extrapolation they assemble simple logical primitives into complex strategies. For example, a person taught to prove simple lemmas can 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 struggle to solve complex reasoning tasks in scalable, algorithmic way. Recently, Deep Thinking systems have been proposed as a way to represent and learn scalable reasoning processes using recurrent neural networks [22]. The word 'thinking' in this context refers to sequential processing to solve discrete/logical problems. These systems train recurrent models (networks that recycle parameters between layers) to 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 algorithms for solving classes of problems. After training to solve small/easy problem instances with few recurrent iterations, the algorithm is then extended to run for more iterations at test time. In doing so, the system can achieve algorithmic extrapolation, solving problems greater difficulty than those in the training set.
To date, the level of algorithmic extrapolation observed in thinking systems has been quite modest. For example has been demonstrated that a system trained on 9x9 mazes can extrapolate solve a 13x 13 maze. These systems fail to achieve greater extrapolation because of problem we call overthinking; recurrent systems, when extended too far outside their training regime, often deteriorate and fail to produce interpretable outputs.
In this work, we design purpose- built neural architectures and specialized training loops to make possible to train systems that do not suffer from overthinking and instead converge to a fixed point when iterated for thousands of iterations. By doing so, we are able build thinking systems that exhibit extreme algorithmic 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 thousands of iterations.
We experiment on benchmark problems for measuring extrapolation behavior. These tasks include computing prefix sums, finding optimal solutions for two-dimensional mazes, and solving chess puzzles [21]. Our architectures and training routines significantly outperform existing methods on all tasks in the benchmark suite. Additionally, we demonstrate that iterative methods for these reasoning problems are susceptible overthinking, problem that overcome by our new architectures and training loops.
Our contributions can be summarized as follows:
We provide a recurrent architecture for algorithmic extrapolation, in which the problem input concatenated directly to the feature stack of certain layers the recurrent thinking module. This prevents the problem instance from being forgotten if deep features become noisy, corrupted, or lossy.
We develop new training routine that incentivizes recurrent networks to make incremental improvements towards a solution, improving the feature representation after each iteration. This training process removes information about how many times the recurrent module has been applied, preventing the network from learning iteration-specific behaviors and instead allowing models to learn scalable behaviors that can be iterated indefinitely for extrapolation.
We analyze the overthinking problem and show that our models overcome this phenomenon. In some cases, the algorithms learned by the proposed networks appear to be capable of solving problems of arbitrary size, despite training only on very small problem instances.
Our improvements in performance on the tasy-to-hard benchmark datasets 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 models do not forget solutions in settings where previous models overthink.

---

2 Related work
Preliminary work [22] shows that simple recurrent architectures, when trained solve various reasoning problems, can exhibit algorithmic extrapolation while their feed-forward counterparts cannot. this section, we contextualize this approach amongst prior work on algorithm learning, adaptive neural models, and logical extrapolation.
Algorithm learning describes models that learn scalable processes from data. Early works on this topic study the ability of recurrent neural networks (RNNs) to process input strings of arbitrary lengths [7. 20). More recent work by Graves et al. [9] introduces neural Turing machines designed to mimic programmable computers, and Kaiser and Sutskever [13] propose a parallel version inspired by massively parallel graphical processing units. These methods, and various improvements them, show promising results on 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 [9, 13, 6]. Since they are based on classical RNNs, however, the amount of computation they perform directly linked to the length of the input string, which prevents them from executing more or less computation independently of the input size (or corresponding to the difficulty of the problem). Moreover, classical RNNs are often trained incrementally to produce one bit at a time, rather than synthesizing an algorithm for solving an entire problem end-to-end. This makes it difficult to apply them to problems where the solution cannot be decomposed into incremental parts (e.g., chess).
Constraint satisfiability problem (CSP) solving networks disentangle the amount of work from the input size [24]. Specifically, message passing neural networks can execute more passes to solve harder CSPs [24]. 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 [19]. 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 [8]. Eyzaguirre and Soto [5] exhibit strong performance on visual question answering by introducing a differentiable version of ACT. Adaptive transformer-based language models also exist, notably Universal Transformers and Depth-Adaptin Transformers, which utilize ACT to determine the work required for each input [3, 18]. Similarly, iterative residual networks like NAIS- repeat blocks in stages and perform well on image classification
[2]. All of these works test their methods in-distribution, i.e. where the training and testing data are sampled from the same distribution; logical extrapolation outside the training domain is not considered.
Logical extrapolation describes the task of generalizing to test sets which comprise more com- putationally complex samples than the training data. Nuamah [16] claims that "neural network models with end-to-end training pipelines ... cannot, on their own, successfully perform algorithmic reasoning," and instead proposes hybrid hand-crafted and learned approach. Similarly, Palm et al. 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. Schwarzschild et al. [22] employ recurrent networks based on weight-sharing architectures, which can be made deeper at test time independent of the input size. These systems 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. ] reformulate the halting unit in ACT leading to probabilistic RNN models 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. [22] depart from classical recurrent networks for text, which learn from step-by-step supervision to produce output tokens one at a time. In contrast, "thinking" systems autonomously synthesize a scalable algorithm end-to-end with no supervision over what each algorithmic step should do. Even more importantly, they can be applied to solve complex problems that are difficult or impossible to decompose by hand. In other words, lequence-to-sequeee models, even adaptive RNNs, have severe limitation in that the problem needs to be represented as sequential input and that each token in the output is generated by the same function. 'Thinking'
3

---

systems, on the other hand, provide mechanism for domains where this type of decomposition is unnatural, difficult, or even impossible.
3 Methods
LSI
Figure 2: Architecture schematics. Left to right: A feed-forwant network, network containing three recurrent blocks (in green) that share weights, and recurrent network with recall.
We begin with some terms and definitions. We study networks that share weights across blocks of layers during training. For example, instead of three distinct residual blocks, a single residual block repeated three times (see Figure 2 for a graphical depiction). 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 layers applied in a recurrent network as "depth," and this quantity grows as the number of recurrent iterations increases. The number of feature maps produced by each layer (or the number of filters in a convolutional layer) is referred to as its "width."
More formally, let r be a function representing a recurrent block, e.g. ResNet block [10], and let r denote recurrences of that function, e.g. (a)) = r(r(z)). Let denote feature map, or an output of r, and let e = r" (z). We also consider an initial "embedding function" denoted by p, which projects an input instance into feature space, and also a final "output head" denoted by h, which maps features outputs. A Deep Thinking (DT) network with iterations of the recurrent block can then be expressed as follows.
mx) == Mrr (p(z))
(1)
In our systems comprises a convolutional layer followed by a ReLU, is single four-layer residual block, and h is set of three convolutional layers with ReLUs after the first and second. We fix m 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 new approaches for achieving improved extrapolation: recall architectures and modified optimization process.
3.1 Recall architectures
When humans think for a long time to solve a problem, we often stop to reread the question or review the task at hand. We improve DT architectures to periodically recall the input exactly. We incorporate this capability into architectures by concatenating the input problem to the features output from each instance of the recurrent block.
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 highway networks, ResNets, and DenseNets, stabilize training [25, 10, 11]. Our architectural modification is driven by the intuition that 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 reproduce any missing or damaged features, and makes impossible to "forget" the problem being solved. To formalize this architectural change, adding recall to the network be expressed using the notation defined above as follows.
freeal(2: v; m) = hi (P())))) where
Trecall (0) x) == r((ó,,)) x]) and r x) = Treean
m- m (d), x),
Whereas the input to r at iteration is usually dx-1, with recall, the input to r at iteration is [dx-1- .t). or the concatenation of the input with the feature map output by the previous recurrence. We add single convolutional layer to map the input [ xr) to feature map of the same shape as - We refer to DT networks with concatenating skip connections DT-Recall models.

---

3.2 Promoting forward progress through optimization
We propose a training objective to encourage the system to incrementally make progress from any starting point. We do this by inputting problem instance and running the recurrent module for some random number of iterations. We then take the output of this process, restart the recurrence in the network with these features as if iterations had just begun (discarding gradients from the initial iterations), and train the model to produce the solution after a random number of additional iterations.
This incremental training process has two benefits. First, it trains the network to continue improving the quality of partial solutions, even when they contain errors or distortions that creep in from running many iterations. Second, by choosing features from a random iteration to serve as the initial state for the training step. we discourage the network from internally counting iterations and learning iteration-specifee behaviors, such as behaviors that get executed only on iteration five, for example. Rather, the network is encouraged to learn veration-agnosse 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 's parameters to minimize loss after -- k total iterations when it starts with the partial solution. This is done by detaching the recurrent mod- ule's output after n iterations from the computation graph before computing the gradient of the loss at iteration n k. The process above is an analog of truncated backpropagation through time [12], with a random starting and end point. During training, we ensure that the sum of n and k is less than a 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 full forward and backward pass through the unrolled m-iteration network.
Algorithm 1 Incremental Progress Training Algorithm Input: parameter vector 0, integer m., weight a for batch _dx 1, 2, ... do
Choose n ~ U1, m - 1} and lt ~ Uf1, m -nnn n) Compute on with n iterations w/o tracking gradients Compute Drag with additional & iterations
Compute Um with new forward pass of m iterations Compute Limax iters with Do. and Eprogresive with fpon
Compute L (1 -11-on a) Lmax_ites + a Eprogressive
Compute VoL and update 0
end for
The training loop and computation of the loss are given in Algorithm 1. The incremental progress term is referred to as Erougeessive (or progressive loss), and the contribution to the loss from the fully unrolled network is denoted by Lemax Lierss See Appendix A.3 for more details.
3.3 Datasets
We evaluate our methods on the benchmark problems available in the Python package easy-to- ttto-hharddaaat which generates reasoning problems of various difficulties. The three prob- lems 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. [21] for more detail, including the data generation process. Note that the architectures consider are fully convolutional, and produce outputs of the same dimension as their inputs. Furthermore, the training and testing datasets we consider have labels of the same dimension 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 7th bit of the target is the mod two sum of all bits prior to and including the 2t 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 [15]. All of our training is done 32-bit strings and we explore the behavior of our models on longer strings, even showing strong performance on 512-bits.
The mazes we consider are -wo-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 are maps of the same dimension as the input maze, but with ones on the optimal path and zeros elsewhere. We make more challenging datasets by increasing the size of the mazes. All training in our work is done on 9 x 9 mazes. Despite this small training size, we benchmark
5

---

extrapolation performance on mazes of size 59 x 59, and observe models solving mazes as large as 201 x 201 with high success rates.
The chess puzzles we consider are mid-game 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 chess board). The goal of this task is to find the best next move, and each target encodes this information in an 8 array with zeros everywhere except the origin and destination of the piece 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 org. Problems are sorted by difficulty and we train on the first 600K easiest puzzles and test our models on indices ranging from 700K to 1.1M.
Our experiments have two categories: (I) Testing extrapolation at extreme leaps in problem diffi- culty/size and (II) Testing scenarios that emphasize the overthinking trap and how we avoid it. The following two sections describe each of these experiments, respectively.
4 Extreme extrapolation
In each problem domain, we show that our models with recall and trained with progressive loss perform best.² 2 We compare our models to baseline DT models trained with m = 30 and problem- specific widths (detailed below), which are both wider and deeper than the models previously studied. However, to provide a fair comparison, we use consistent values for these parameters across our experiments since we find that all architectures benefit from additional width and depth. We also compare to ceed-forward models without weight sharing of the same effective depth as a 30-iteration DT network, which are ResNets [10]. We include ablation studies lending credence to our claim that both the architecture and the loss modifications are instrumental boosting extrapolation power. For discussion hyperparameter choice and training set-ups for all of our experiments, see Appendix A.4. Also, code to reproduce all of our experiments is publicly available."
Prefix Sums Though prior DT systems show extrapolation behavior on strings with up to 48 bits, we use a 512-bit testing set to illustrate the extent of our performance improvements over prior work. We train our models on 32-bit strings with a maximum of only m = 30 recurrent iterations (indicated in the figures as the shaded train regime section), 400 convolutional filters per layer, and incremental loss weight a = 1 when progressive loss is used.
DTT-Recall Prog,
DT--Relall
DT Prog.
DT-Recall Prog.
DT-Rccall
Prog. Train Regime
FF
Teain Regine
DT
..... FF
100
100
80
80
Avcuracy (%)
60
60
40
20
300
500
0-
250
500
750
Test- Test-Time Iterations
Test-Time Iterations
Figure 3: Prefix sum models trained on 32-bit inputs extrapolate 512-bit data. The value of our recall and progressive loss is clear by how quickly and accurately our models solve this very large problem.
Figure 4: Maze solving models trained x 9 inputs extrapolate to 59 x 59 problems. Mazes this large cannot be solved without recall, and furthermore progressive loss leads to more accurate models.
In Figure 3, we show that without our techniques, DT networks are unable to solve very long binary strings, achieving 0% accuracy on 512-bit data, while DT-Reccll models with progressive loss can solve more than 97% after approximately 200 iterations.
Note that the DT models presented in this section maximum contidence exit rule for fair comparison. This exit rule marginally mitigates overthinking in the regime initially tested by Schwarzschild al. [22]. Models with our new architecture and/or loss do not require this rule and are presented with default outputs. A PoTohch implementation of the code available at gi github. com/ 200m// deep- deepethinneeeee thinking.

---

To better understand the individual effects of our proposed approach, we perform an ablation study. Figure also makes clear that DT-Recall networks trained without progressive loss achieve the same final accuracy as models trained with it; but they require approximately twice the number of iterations to get there. 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 progressive loss outperform vanilla DT networks, however neither of these models (without recall) can solve more than a handful of 512-bit testing examples.
One other study in this domain reveals the importance of randomly setting n in Algorithm To show this, we modify the algorithm slightly, to always use n O, rather than randomly choosing that value. While this change seems minor, and indeed makes a small impact on the results, it does lower the accuracy. Specifically. DT networks with recall on average solve 90.27% of 312-bit test samples when trained with the n = 0 loss, and with randomly sampled values of n, as described in Algorithm 1, they achieve 97. 17% accuracy. From this ablation experiment, we conclude that randomly sampling these values is, in fact, beneficial.
Mazes As with the prefix sum problem, we can improve performance on hard mazes by combining incremental training and recall architectures. In particular, while Figure 3 may not convince the reader that our proposed loss is critical since DT-RReaall models without it perform very well, with more complex data a drastic difference emerges.
We show in Figure 4 that on the significantly harder test set of 59 x 59 mazes, our models exhibit strong algorithmic extrapolation, while previous methods, both feed-forward and DT systems, completely fail. Not only do our models achieve a higher peak accuracy, but they do not overthink, as can be seen by the flat spans in the curves in Figure 4. In fact, can push these systems to their limits (and the limits of our hardware) and find that our models can solve 74% of the 201 x 201 mazes." See Appendices A.9 and A.10 for examples of 201 201 and 801 x 801 mazes, which our models solve using 2,000 and 20,000 iterations, respectively (the equivalent of 10,004 and 100,004 convolutional layers in depth).
The best models in Figure 4 are DT-Recall mod- els trained with 30 maximum iterations and weighting in the loss of a = 0.01. In order to see how critical recall is in overcoming the overthinking problem, we show in the same fig-
DTT-Rccall Prog.
DT-Recall
Prog-
DT
FF
Train Regine
100
Accuracy (%)
that without recall all models fail to extrap- olate. The benefit of progressive loss is also highlighted by the fact that De-Recaal models with progressive loss achieve on average 97% accuracy while Re-Recall models without pro- gressive loss only reach an average of 83%. In- terestingly, in this domain, 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.4.
60
25
75
160
125
Test-Time Iterations
Figure 5: Chess models trained on the first 600K easiest puzzles extrapolate to 600K-700K. Recall and progressive loss are required to retain accuracy with many iterations.
Chess Puzzles We further find that using recall and progressive loss yields notable improvement on chess puzzles. In Figure 5, we see that our techniques lead to a 4% accuracy improvement compared to networks without them. We show that either removing recall or training without the incremental progress loss will hurt performance. Our best models are DT-Recall networks with 512 convolutional filters in each layer trained with a maximum of 30 iterations and a weight of a = 0.5. Moreover, their accuracy preserved as the number of iterations increases, while that of the other DT networks decays seriously after about 70 iterations the sign of overthinking and the subject of the following section. Results from tests on harder puzzles are in Appendix A.5.
'The seemingly modest performance of 74% accuracy should be understood in context. Of the 100 mazes in this test set, 20 have pixel (out of 166,464) wrong and no single maze has more than seven mistakes.
7

---

The overthinking problem
Deep Thinking networks boast impressive capabilities to solve harder problems by thinking deeper, but they are prone to overthinking. In particular, by testing models data closer in complexity the training data we can better compare our methods to prior work. Figure 6 shows that some recurrent networks that perform algorithmic extrapolation may collapse entirely when performing too many iterations.
DT--Recall Prog.
Prog.
Regime
100
100
Accuracy (%)
75
50
25
100
200
25
50
75
100
Test-Time Iterations
Test-Time leeaations
Figure 6: Left: Prefix sum models trained on 32-bits tested on 48-bit inputs. Right: Maze models trained on 9 x 9 mazes tested on 13 x 13 mazes. In these regimes, we can see that the weaker models can extrapolate slightly harder problems, but importantly only models with recall avoid the overthinking trap.
Our models do not suffer from the sharp decline in accuracy the number of iterations increases. In Figure 6, the decay in each dashed curve and the stability with high iteration counts of the solid curves shows that recall is critical to avoid overthinking, See Figure 5 for an example where progressive loss is needed too. The monotonic increase in accuracy with added iterations in our models is practically useful, 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.
One observation that can made from these results is that the overthinking problem often 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 our models seem to converge to a fixed-point solution as they iterate rather than becoming unstable. This property is desirable and using representative models trained to solve mazes, we explore how robust models are when we manipulate the features during the thinking process. In Appendix A.8, present similar findings for prefix sums.
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 we add Gaussian noise (with mean zero and standard deviation one) to the features after iteration of maze solving. Models with recall can still solve 13 13 mazes even when we perturb the features, but models without recall cannot. See Appendix A.8 for plots, tables, and further discussion.
Next, we ask whether the initial feature maps (after the projection layer p) carry any important information. We test this by replacing the feature maps with zeros after 50 iterations at this point the model has solved the maze and by annihilating the features we remove all information from p. In this case, our DT-recall models naturally regenerate their features and recover to solve the problem again, indicating that the learned algorithm able to find a solution using the input without the initial projection. See Figure 12 in Appendix A.8.
5.2 Manipulating inputs
With the notion that our models may embody convergent process, we turn our attention to investi- gating how the networks determine when to stop manipulating the representation. Perhaps there something in the feature maps output by late iterations that tells the network to stop working. We can test this in two ways. We perturb the input (which we concatenate onto the features) after some number of iterations first subtly, then by swapping it with an entirely different example.
8

---

RT-R
DT-Reeall
Prog.
Avg. Solution
30
100
Features
wagpee
75
20
Accuracy
Soletfess reCHVII
50
10
25
10
20
40
50
100
150 10
50
100
150
Index of Perturbed Bit
Test-Time Iterations
Test-Time Iterations
Figure 7: Left: How long it takes prefix sum models to recover from perturbation recall keeps this quantity small. Center: Test accuracy on 13 x 13 mazes when features are swapped after 50 iterations. Right: The change in the features when solving x 13 mazes.
To start, we flip single bits in the input string for prefix sums. We explore the response of the model to flipping bits different indices after 50 iterations (when the models have already solved the initial 48- bit problem). The left panel of Figure 7 shows recovery time as a function of index of the flipped bit. We observe that our models can indeed recover from single bit flips, and recovery time decays linearly with index. Higher indices are closer to the end and affect fewer bits in the output than lower ones.
Similarly, with mazes, change the input by moving the end of the maze two steps closer to the start. We use this new input concatenated with features generated after 50 iterations of solving the original maze. In this case, we see in Figure 13 (appendix) that models with recall self-correct and solve the new problem in very few iterations.
The last way we 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 obtained 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 a solution for maze A. The center panel of Figure shows the effect of swapping feature maps.
5.3 Converging to a fixed point
We can also study the convergence by measuring the change in the output at each iteration. A decreasing change in the feature maps with each additional iteration suggests that the network manipulates the representation, moving 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 qualitatively categorize each model type (architecture/loss pair) as convergent or non-convergent. To do we measure the change in the solution each iteration with f--norm of the difference. The right-hand panel of Figure shows Aø(n) := I -n-1llea. 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 algorithmic extrapolation power of neural networks. We propose an architecture and loss that lead to gigantic generalization leap from easy training data to much more complex testing examples. Furthermore, we show that our models avoid the overthinking trap. We test algorithmic 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 as 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.
Existing methods deteriorate at such large numbers of iterations creating the need for stopping mechanisms. Even for small numbers of iterations, if the halting decision sub-optimal, then the output may also be sub-optimal. Because our models have the desirable property (for an algorithm) that they converge to fixed points, they run no such risk, require no learned or hand-crafted stopping rules, and can perform well at huge numbers of iterations.

---

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 to testing data can be handled without overthinking. Importantly, our neural networks learn end-to-end how to perform algorithmic extrapolation. Our findings are limited to the domains we consider and further exploration into how our methods will perform in other settings is needed.
The application of machine learning systems to real world data demands effective extrapolation, even though this is often overlooked. Curated research benchmarks designed for measuring generalization (e.g. ImageNet [4]) omit a critical feature of real data in the field: there is often reason to think that data encountered after training could embody more difficult examples than those in the training set. Consider that a practitioner cannot always know a priori the range of complexity that exists in their domain. If we are to employ models that cannot extrapolate, we are to miss out on using our machine learning tools to solve problems harder than those for which we already have answers.
Acknowledgements
This work was supported by the ONR MURI Program, The National Science Foundation DMS- 1912866, and the Office of Naval Research (#4720008163).
References
Andrea Banino, Jan Balaguer, and Charles Blundell. Ponderett: Learning ponder, arXiv preprint arXiv:2107.05007 2021.
[2] Marco Ciccone, Marco Gallieri, Jonathan Masci, Christian Osendorfer, and Faustino Gomez. Nais net: Stable deep networks from non-autonomous differential equations. Advances in Neural Information Processing Systems, 31, 2018.
[3] Mostafa Dehghani, Stephan Gouws. Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. International Conference on Learning Representations, 2018.
[4] Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, Fei-Fei. Imagenet: large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-255. IEEE, 2009.
[5] Cristobal Eyzaguirre and Alvaro Soto. Differentiable adaptive computation time for visual reasoning. In Proceedings the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12817- 12825, 2020.
[6] Karlis Freivalds and Renars Liepins. Improving the neural gpu architecture for algorithm learning. arXiv preprint xxXiv/102.00877,, 2017.
[7] Felix Gers and F. Schmidhuher. Lstm recurrent networks learn simple context-free and context-sensitive languages. IEFF Transactions Neural Networks, 1336):1333 2001.
[8] Alex Graves. Adaptive computation time for recurrent neural networks. arXiv preprint urXiv:1603.08983. 2016.
[9] Alex Graves, Greg Wayne, and Danihelka. Neural turing machines, 2014.
[10] K. Hc, Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 770-778, 2016
[11] Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger. Densely connected convolutional networks. In Proceedings the the IEEE conference computer vision and pattern recognition.
pages 4700 4708, 2017.
[12] Herbert Jaeger. Tatorial training recurrent neural networks, covering BPPT, RTRL, EKF and the" echo state network" approach, volume 5.01. HMMD-Forshunnszzentuum Informationstechmik Bonn, 2002.
[13] Eukasz Kaiser and Ilya Sutskever. Neural gpus learn algorithms. arXiv preprint arXiv:1511.0881 2015.
[14] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint InXiv:412.69900 2014.
10

---

[15] Richard E Ladner and Michael Fischer. Parallel prefix computation. Journal of the ACM (JACM). 27(4): 831-838, 1980.
[16] Kwwabena Nuamah. Deep algorithmic question answering: Towards a compositionally hybrid Al for algorithmic reasoning. arXiv preprint arXi:22109.80006 2021.
[17] Rasmus Berg Palm, Ulrich Paquet, and Ole Winther. Recurrent relational networks. In 32nd Conference Neural Information Processing Systems, pages 3368-2278. Neural Information Processing Systems Foundation, 2018.
[18] Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations. 2021.
[19] Jürgen Schmidhuher. Self-delimiting neural networks. arXiv preprint arXiv:/210.011: 2012.
[20] Jürgen Schmidhuber, Daan Wierstra, Matteo Gagliolo, and Faustino Gomez. Training recurrent networks by evolino. Neural computation, 19(3:757-779, 2007.
[21] Avi Schwarzschild, Fitan Borgnia, Arjun Gupta, Arpit Bansal, Zeyad Emam, Furong Huang. Micah Goldblum, and Tom Goldstein. Datasets for studying generalization from easy hard examples, 2021.
[22] Avi Schwarzschild, Eitan Borgnia, Arjun Gupta, Furong Huang, Uzi Vishkin, Micah Goldblum, and Tom Goldstein. Can you learn an algorithm? generalizing from easy hard problems with recurrent networks. Advances Neural Information Processing Systems, 34:6695-6706, 2021.
[23] Avi Schwarzschild, Arjun Gupta, Amin Ghiasi, Micah Goldhlum, and Tom Goldstein. The uncanny similarity of recurrence and depth. In International Conference on Learning Representations, 2021.
[24] Daniel Selsam, Matthew Lamm, Benedikt Bünz, Percy Liang, Leonardo de Moura, and David I. Dill. Learning a sat solver from single-hit supervision. arXiv preprint arXiv: 1802.0330 2018
[25] Rupesh Kumar Srivastava, Klaus Greff, and Jürgen Schmidhuber. Highway networks. arXiv preprint arXiv:1505.00330 2015.


Reviews:
[{"review_id": "xoPj1WarP2", "reviewer": "Reviewer_fZnJ", "summary": "This paper proposes an approach that allows models to extrapolate to more complex problems than they observed during training, by recurrently applying the same operations many times. The paper offers two key innovations to enable this: 1) a new loss is proposed that encourages the model to converge in as few iterations as possible, rather than learning iteration-specific behaviors, and 2) a recall element is introduced that allows models to continue to have access to the original problem, which appears to prevent a phenomenon they refer to as 'overthinking' in which a model's outputs appear to drift over many iterations, even after initially discovering the correct answer. They evaluate this approach on a range of tasks that each involve generalization to more complex problems than observed during training.", "questions": "My questions are detailed in the 'strengths & weaknesses' section above. ", "limitations": "I don't believe this work has any potential negative societal impact.", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "In general I have a positive view of this work. I think the overall approach is compelling and the results are generally strong. However, I have reviewed this work before, and some of the issues raised during that discussion are addressed here, but some issues are still neglected. Here are the issues that I think still need to be addressed:\n- One of the primary issues that I raised in the previous discussion is the lack of comparison with other baselines. The choice of architecture here, in which only the output from the previous stage is fed as input to the next stage, is interesting, but I still feel that the work would be improved if there was a comparison with other baselines, such as recurrent models that have a hidden state (e.g. lstm), or models with access to an external memory. It would be good to know whether this specific design decision is important to the extrapolation capability of the models.\n- The progressive loss is primarily designed to encourage the models to converge quickly, without relying on having a certain number of iterations to arrive at the correct answer. A natural comparison here is some sort of penalty where loss increases as a function of the number of iterations. In the previous discussion, the authors described some experiments testing this alternative, and reported that it did not perform as well as their approach, but those experiments don't appear in the present manuscript.\n- In the previous version of the paper, there was some mention of the fact that some number models were excluded from analysis because they did not reach a certain training accuracy criterion. In discussion, the authors stated that the overall results were not significantly affected by the exclusion of those models. But now there is no longer any mention of excluding models based on training accuracy, though the results appear to be the same as in the original manuscript.\n- The previous version of the paper also employed something called a 'maximum confidence rule' where the iteration with the maximum confidence value was selected for testing. According to the previous discussion, this also did not affect the results much, but now there is also no mention of this selection procedure in the current paper.\n\nSo, although I'm generally positive about the work, it seems that not much has been added relative to the previous iteration of this work, and in fact some additional results that I think strengthened the paper's conclusions have now been removed. I am eager to hear from the authors about the points above.", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "ddLA7U85YYa", "reviewer": "Reviewer_V9gQ", "summary": "The authors approach the task of 'algorithmic' or 'logical' extrapolation in neural networks, which is to say, the ability of a neural network to generalise patterns learned from small/simple examples in order to solve scaled up/more difficult tasks from the same overall problem, e.g., solve a large maze after seeing many small ones, or a tough chess puzzle after many easy ones.\nThe authors restrict themselves to problem domains in which the outputs to each task are the same size as the inputs, this allows applying convnet-based architectures to arbitrary input size without dimensionality issues.\n\nAn existing approach uses an architecture consisting of 3 convnet-based blocks: one input-processing, one recurrent, and one output-preparing. The blocks feed into each other, though the recurrent block is applied recurrently to its own input some (externally defined) number of times N before feeding to the output-preparing block. The authors revisit this architecture, noting that it a) struggles to extrapolate to especially larger examples from those seen in training, and b) may approach a task solution after some N1 recurrences, but sadly degenerate past that (i.e., giving the wrong output if asked to run for N>N1 steps). They refer to this latter issue as overthinking.\n\nThe authors suggest two modifications to the architecture, motivated by these issues. First, they suggest feeding the original input back into the recurrent net at every step, as part of its input - a \"recall\" architecture; and second, they suggest a careful but straightforward training regime, defining a \"progressive loss\" which encourages the network to output (as best it can) the correct solution at *every* step of the computation and not just the last, this pushes it to a) learn as short a solution as possible and b) stabilize on that solution once reached, avoiding the overthinking issue. Further, to prevent the network from learning iteration-specific algorithms, the authors employ BPTT when using their progressive loss (I am confused on this, see questions).\n\nThe authors evaluate the modified architecture, showing that it scales to far larger/harder problems with the modifications, and perform ablation studies demonstrating the importance of each of their modifications. They even show that their \"recall\" modification is truly being used for the intended purpose by the network, showing how the networks overcome perturbations in the recurrent computation to return the correct solutions for their true inputs.\n", "questions": "Please clarify the following in the paper:\n1. Are all of the internal blocks in the architecture indeed convnets? Could you give (each of) their exact architectures in the paper/appendix?\n2. Figure 2 suggests the original input is fed into the reccurent block (A) at every step, my intuition suggests it should be the processed input (i.e., the arrows should emanate from the 'P' block, not before it), could you clarify?\n3. I am not sure why being iteration agnostic matters? What if the task really can be solved in 3 operations, regardless of input size/difficulty? I understand the value of BPTT to allow training later steps even with limited computational resources! But just don't understand the importance of being iteration agnostic.\n", "limitations": "y", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "strengths:\n- straightforward, intuitive modifications to existing work provide significant improvement in performance\n- well presented, clear work (though please correct/clarify notes in comments and questions!) with convincing evaluation\n- exciting task\n\nweaknesses:\n- while this is a great step for the \"deep thinking\" models, I don't see a lot of comparison to/discussion of other (if any) architectures that could be applied to these tasks. E.g., why not compare to a universal transformer? In this sense the work is not well contextualized.\n\ncomments, including presentation comments and typos:\n\n0. (general) this work uses 'recurrence', but in a different way from RNNs: RNNs recur alongside an input sequence, eating one token at a time and finishing at the end of the input, this architecture simply recurs on its own state (albeit with help of the input recall) for a given number of steps, independently of the input. Given the popularity of RNNs, for clarity, I suggest this distinction is clarified early in the paper, e.g. a sentence of the sort \"this work uses recurrence in a manner different from that commonly referred to in RNN literature, the architecture behaviors should not be confused\"\n1. line 8 \"because behavior degenerates\" I don't know if 'behavior degenerates' so much as iterations continue past correct stop (losing solution), or representation degenerates (also losing solution), or some other issue. Rephrase (unless you prove that indeed \"behavior degenerates\" and is the \"cause\" of the failure to scale).\n2. Fig 1 caption description of 59x59 as \"center\" is confusing, consider just labeling the subfigures\n3. line 20 \"lack the ability to solve complex reasoning tasks in a scalable, algorithmic, way\" again strong concrete statement, not sure I agree, soften/rephrase.\n4. Eq 2 (description of f_recall, r_recall) - bad notation: r_recall takes two inputs, one of which is its own previous output, and the other the original input to the net. By construction such a function cannot be applied recurrently (r^2_recall(phi,x) unrolls to r_recall(r_recall(phi,x)), which makes no sense because the x has to be put in again). Consider instead notation similar to that used in RNN descriptions, e.g.: phi_t+1=r_recall(phi_t,x), phi_1=r(x), f_recall(x,m)=h(phi_m) (something like that - fill it out right according to your architecture. Please also make sure to properly distinguish between the original input x and the processed input P(x), so we can tell which one is being used where!). \n5. line 163: \"re-initialize the network\" - the entire network or just the recurrent part? clarify\n6. line 164: clarify here that gradient is being discarded, easier to follow that way\n7. line 169-170: (too) strong statements! 1) the network may arbitrarily implement counting its iterations, it is not \"prevented\" from doing so, and 2) the network might still learn iteration specific behaviors, whether by luck or even directly through the SGD (given that the standard loss is also still used!). The only claim I would make here is that encouraging random iterations to be correct a) encourages a stabilizing behavior (i.e. not losing the solution once obtained) and b) encourages 'faster' (shallower) solutions (i.e. having the solution ready in the earliest iteration possible) - basically, it pushes the network to have the correct solution available in the maximum number of iterations possible.\n8. Fig 3: would also be nice to see the stability of DT-recall (i.e. no prog loss) *after* reaching the solution, e.g. plot accuracy up to 1000 iterations (given that prog loss is supposed to encourage stability, and here we only just see DT-recall reach the solution but not whether it stabilizes there)\n9. line 254-255: [s/of sampling/of randomly sampling], [s/We modify/To show this, we modify], [s/slighty to always/slightly, to always]\n10. line 259 s/random/randomly\n11. line 265 \"harder\" - harder than what?\n12. line 293-294: explicitly mention that the best models had progressive loss (currently missing from this sentence)\n", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "M2NkcCEZpS", "reviewer": "Reviewer_nMzj", "summary": "The aim of this paper is to extend the generalization of learning approaches on algorithmic challenges.  \n\nThey introduce an architecture for algorithmic extrapolation, which amounts to providing the input to the network as input to each current layer.\n\nThey introduce a training procedure that aims to train a network to improve a partial solution in an incremental manner.  Specifically, the loss term minimizes an error term not relative to the actual input $x$ but instead starting at some intermediate value.", "questions": "If I understand correctly:\n\n1. $\\mathcal{L}_{\\textrm{max iters}} = |\\hat{y}_m - y| = |f_\\text{recall}(x; m) - y|$ and:\n2. $\\mathcal{L}_{\\textrm{progressive}} = |f_\\text{recall}(x_n; k) - y|$ where $x_n = f_\\text{recall}(x; n)$\n\nIs this correct? If so, I suggest adding an expression like this to the paper.  Currently the loss terms are not fully specified.  Section 3.2 and Algorithm 1  give a procedural account of how to compute the gradient step but we don’t need to talk about detaching computation graphs etc to be able to state exactly what loss term we are trying to minimize.\n\nAs described above, are the ablated and non-ablated networks have the same or different number of parameters?  If different, can we control for this?", "limitations": "It would be useful to understand more clearly what if any tradeoffs there are for using progressive loss", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is well written and easy to follow.  I like that the idea is not complicated beyond necessity.\n\nI do not think “overthinking” is an appropriate term.  Deterioration of recurrent networks outside their training regime does not correspond to the common sense notion of overthinking.\n\nGiven how similar this notion of recall is to existing work, it’s not clear to me that one could consider the architecture itself to be a novel contribution.  The insight that this kind of architecture dramatically helps with this kind of extrapolation is a contribution in its own right, though.\n\nThe ablation studies demonstrate the importance of the different components.  However, with the increased input size of $$f_recall$$ over $$f$$, does this change the number of parameters?  A common pitfall for ablation studies in neural networks is to not account for a reduction in the capacity of a network when ablating a component, and it seems like this may be happening here too.\n\nMy main worry is that the results here do not generalize beyond the small number of examples shown.  With such a small and simple change, this kind of paper is most convincing with extensive experimentation.", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "H_6jwdufnUZ", "reviewer": "Reviewer_Z4mS", "summary": "This paper studies recurrent networks, architectures that iterate over shared layers. Previous research by Schwarzchild et al. (ref 22 in the paper) has shown that for three specific tasks, these networks could extrapolate to larger problems than those seen at training, by being given \"ponder time\", i.e. more iterations at inference than at training.\n\nThe authors note that as ponder time gets large, extrapolation accuracy drops, a phenomenon they call \"overthinking\". In practice, overthinking prevents the model to extrapolate to problems much larger than those seen at training.\n\nThe authors propose two improvements to the original architecture: \n- recall: repeating the original problem at every iteration, by concatenating it to the input of the shared layer\n- progressive training: instead of computing the loss over m iterations, compute it over k successive iterations between n and n+k ( n and k randomly sampled). This amounts to randomizing the number of iterations at train time, and limiting the back propagation of gradients to the last k layers (an approach reminiscent of GNN).\n\nOn the three problems considered, they show that recall and progressive training prevent overthinking, and allow the model to generalize to much larger problems than seen at training. ", "questions": "**Sec. 3.1**  The recall model is defined as phi(k) = r(phi(k-1), x), (phi(k) the output of iteration k, x the input sequence). Since the input of the first share layer is p(x), it would make more sense to use r(phi(k), p(x)). Have you tried this? What is the rationale for r(phi, x)? \n\n**Sec. 3.2** Progressive training amounts to preventing the gradients to back-propagate fully. The experiments suggest that whereas it does not improve extrapolation (recall does most of the work), it speed up inference (in the sense that extrapolation is faster). Could such a principle be applied to regular deep networks, i.e. without shared layers?\n\nWhat is the benefit of having the loss a weighted sum of total (m loops) and progressive (k loops) loss? Couldn't the same effect be achieved by sampling (with probability alpha) examples for which the loss would be computed for the whole m steps, and others where it would be computed on the k last? Would there be an advantage in varying alpha as training proceeds (as a form of curriculum learning)?\n\n**Architecture** Both input and output layers (p and h) are convolutional. Also both the prefix sum and maze problems are equivariant with respect to translation, which makes conv-nets a natural choice of embeddings. How dependent on the choice of convnets are your findings? e.g. what happens if you change p (or h) to a linear layer ?\n\nHave you done an ablation on m? in other words, could your extrapolation results allow one to reduce the number of iterations during training (hence saving memory when problems become very large)? \n\n**misc.**\nLines 213-214 : zeroes or ones, or black and white? \nLine 316 : should be chess, not prefix sums?", "limitations": "A discussion about the tasks for which this technique can be applied would reinforce the significance of the paper. What is the importance of translational invariance? Is the technique limited to problems that can be solved to algorithms with little or no memory (to compute the next term of the prefix sum, you only need the previous result, to choose the next move in a maze, you need limited information about your previous moves)?", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is very well written, and the architecture clearly described. The experimental results are compelling: on the binary integration problem, the model extrapolates to sequences 16 times longer, on the maze problem, to arrays 500 times larger. This is a novel idea, and a significant result on a very important subject. It makes no doubt to me that this paper should be accepted. \n\nThe paper could be made **much stronger** by improving the experiments. In particular, \n- the binary integration problem is very simple: it is a first order recurrence y[n] = y[n-1] + x[n] mod 2  (x and y the input and output sequences). Experimenting with larger order (e.g. order two, in Fibonacci fashion) and/or larger moduli would be interesting.\n- as far as extrapolation is concerned, the chess results are inconclusive: both figures 5 and 9 suggest that 30 iterations are enough to achieve best performance. The authors note that recall helps achieve a slightly better accuracy, but this is not a case of Deep Thinking.\n\n", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept