Original: End-to-end Algorithm Synthesis with Recurrent Networks: Logical Extrapolation Without Overthinking
Revised: End-to-end Algorithm Synthesis with Recurrent Networks: Logical Extrapolation Without Overthinking

Original: Abstract
Revised: Abstract

Original: Machine learning systems perform well on pattern matching tasks, but their ability to perform algorithmic or logical reasoning is not well understood.
Revised: Machine learning systems perform well on pattern matching tasks, but their ability to perform algorithmic or logical reasoning is not well understood.

Original: 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.
Revised: 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.

Original: Algorithmic extrapolation can be achieved through recurrent systems, which can be iterated many times to solve difficult reasoning problems.
Revised: Algorithmic extrapolation can be achieved through recurrent systems, which can be iterated many times to solve difﬁcult reasoning problems.

Original: 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.
Revised: 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.

Original: " We propose a recall architecture that keeps an explicit copy of the problem instance in memory so that it cannot be forgotten.
Revised: " We propose a recall architecture that keeps an explicit copy of the problem instance in memory so that it cannot be forgotten.

Original: 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.
Revised: We also employ a progressive training routine that prevents the model from learning behaviors that are speciﬁc to iteration number and instead pushes it to learn behaviors that can be repeated indeﬁnitely.

Original: These innovations prevent the overthinking problem, and enable recurrent systems to solve extremely hard extrapolation tasks.
Revised: These innovations prevent the overthinking problem, and enable recurrent systems to solve extremely hard extrapolation tasks.

Revised: 36th Conference on Neural Information Processing Systems (NeurIPS 2022).

Original:  Introduction
Revised:  Introduction

Original: Humans solve complex logical reasoning problems through logical extrapolation – they assemble simple logical primitives into complex strategies.
Revised: Humans solve complex logical reasoning problems through logical extrapolation – they assemble simple logical primitives into complex strategies.

Original: For example, a person taught to prove simple lemmas can in turn prove more complex theorems simply by expending more cognitive effort.
Revised: For example, a person taught to prove simple lemmas can in turn prove more complex theorems simply by expending more cognitive effort.

Original: Neural networks have achieved great success at pattern matching tasks, often exceeding human performance, but they struggle to solve complex reasoning tasks in a scalable, algorithmic way.
Revised: Neural networks have achieved great success at pattern matching tasks, often exceeding human performance, but they struggle to solve complex reasoning tasks in a scalable, algorithmic way.

Original: Recently, Deep Thinking systems have been proposed as a way to represent and learn scalable reasoning processes using recurrent neural networks [22].
Revised: Recently, Deep Thinking systems have been proposed as a way to represent and learn scalable reasoning processes using recurrent neural networks [22].

Original: The word ‘thinking’ in this context refers to sequential processing to solve discrete/logical problems.
Revised: The word ‘thinking’ in this context refers to sequential processing to solve discrete/logical problems.

Original: These systems train recurrent models (networks that recycle parameters between layers) to solve reasoning problems.
Revised: These systems train recurrent models (networks that recycle parameters between layers) to solve reasoning problems.

Original: 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.
Revised: Unlike traditional feed-forward models, which are limited in the complexity of problems they can solve by their ﬁnite depth, the effective depth of recurrent models can be expanded after training simply by iterating the recurrent unit for longer.

Original: When trained properly, thinking systems learn scalable algorithms for solving classes of problems.
Revised: When trained properly, thinking systems learn scalable algorithms for solving classes of problems.

Original: 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.
Revised: 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.

Original: In doing so, the system can achieve algorithmic extrapolation, solving problems of greater difficulty than those in the training set.
Revised: In doing so, the system can achieve algorithmic extrapolation, solving problems of greater difﬁculty than those in the training set.

Original: To date, the level of algorithmic extrapolation observed in thinking systems has been quite modest.
Revised: To date, the level of algorithmic extrapolation observed in thinking systems has been quite modest.

Original: For example it has been demonstrated that a system trained on 9 × 9 mazes can extrapolate to solve a 13 × 13 maze.
Revised: For example it has been demonstrated that a system trained on 9 ⇥ 9 mazes can extrapolate to solve a 13 ⇥ 13 maze.

Original: These systems fail to achieve greater extrapolation because of a problem we call overthinking ; recurrent systems, when extended too far outside their training regime, often deteriorate and fail to produce interpretable outputs.
Revised: These systems fail to achieve greater extrapolation because of a problem we call overthinking ; recurrent systems, when extended too far outside their training regime, often deteriorate and fail to produce interpretable outputs.

Original: In this work, we design purpose-built neural architectures and specialized training loops to make it possible to train systems that do not suffer from overthinking and instead converge to a fixed point when iterated for thousands of iterations.
Revised: In this work, we design purpose-built neural architectures and specialized training loops to make it possible to train systems that do not suffer from overthinking and instead converge to a ﬁxed point when iterated for thousands of iterations.

Original: By doing so, we are able to 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.
Revised: By doing so, we are able to 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.

Original: We experiment on benchmark problems for measuring extrapolation behavior.
Revised: We experiment on benchmark problems for measuring extrapolation behavior.

Original: These tasks include computing prefix sums, finding optimal solutions for two-dimensional mazes, and solving chess puzzles [21].
Revised: These tasks include computing preﬁx sums, ﬁnding optimal solutions for two-dimensional mazes, and solving chess puzzles [21].

Original: Our architectures and training routines significantly outperform existing methods on all tasks in the benchmark suite.
Revised: Our architectures and training routines signiﬁcantly outperform existing methods on all tasks in the benchmark suite.

Original: Additionally, we demonstrate that iterative methods for these reasoning problems are susceptible to overthinking, a problem that is overcome by our new architectures and training loops.
Revised: Additionally, we demonstrate that iterative methods for these reasoning problems are susceptible to overthinking, a problem that is overcome by our new architectures and training loops.

Original: Our contributions can be summarized as follows:
Revised: Our contributions can be summarized as follows:

Original: • We provide a recurrent architecture for algorithmic extrapolation, in which the problem input is concatenated directly to the feature stack of certain layers in the recurrent thinking module.
Revised: • We provide a recurrent architecture for algorithmic extrapolation, in which the problem input is concatenated directly to the feature stack of certain layers in the recurrent thinking module.

Original: This prevents the problem instance from being forgotten if deep features become noisy, corrupted, or lossy.
Revised: This prevents the problem instance from being forgotten if deep features become noisy, corrupted, or lossy.

Original: • We develop a new training routine that incentivizes recurrent networks to make incremental improvements towards a solution, improving the feature representation after each iteration.
Revised: • We develop a new training routine that incentivizes recurrent networks to make incremental improvements towards a solution, improving the feature representation after each iteration.

Original: 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.
Revised: This training process removes information about how many times the recurrent module has been applied, preventing the network from learning iteration-speciﬁc behaviors and instead allowing models to learn scalable behaviors that can be iterated indeﬁnitely for extrapolation.

Original: • We analyze the overthinking problem and show that our models overcome this phenomenon.
Revised: • We analyze the overthinking problem and show that our models overcome this phenomenon.

Original: 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.
Revised: 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.

Original: Our improvements in performance on the easy-to-hard benchmark datasets can be categorized in several ways.
Revised: Our improvements in performance on the easy-to-hard benchmark datasets can be categorized in several ways.

Original: First, our models yield uniformly higher accuracy across the most difficult tasks used for testing in previous work.
Revised: First, our models yield uniformly higher accuracy across the most difﬁcult tasks used for testing in previous work.

Original: 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.
Revised: 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.

Original: Lastly, we show that our models do not forget solutions in settings where previous models overthink.
Revised: Lastly, we show that our models do not forget solutions in settings where previous models overthink.

Original: Preliminary work [22] shows that simple recurrent architectures, when trained to solve various reasoning problems, can exhibit algorithmic extrapolation while their feed-forward counterparts cannot.
Revised: Preliminary work [22] shows that simple recurrent architectures, when trained to solve various reasoning problems, can exhibit algorithmic extrapolation while their feed-forward counterparts cannot.

Original: In this section, we contextualize this approach amongst prior work on algorithm learning, adaptive neural models, and logical extrapolation.
Revised: In this section, we contextualize this approach amongst prior work on algorithm learning, adaptive neural models, and logical extrapolation.

Original: Algorithm learning describes models that learn scalable processes from data.
Revised: Algorithm learning describes models that learn scalable processes from data.

Original: Early works on this topic study the ability of recurrent neural networks (RNNs) to process input strings of arbitrary lengths [7, 20].
Revised: Early works on this topic study the ability of recurrent neural networks (RNNs) to process input strings of arbitrary lengths [7, 20].

Original: More recent work by Graves et al.
Revised: More recent work by Graves et al.

Original: [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.
Revised: [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.

Original: These methods, and various improvements to 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].
Revised: These methods, and various improvements to 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].

Original: Since they are based on classical RNNs, however, the amount of computation they perform is 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).
Revised: Since they are based on classical RNNs, however, the amount of computation they perform is 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 difﬁculty of the problem).

Original: 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.
Revised: 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.

Original: This makes it difficult to apply them to problems where the solution cannot be decomposed into incremental parts (e.g., chess).
Revised: This makes it difﬁcult to apply them to problems where the solution cannot be decomposed into incremental parts (e.g., chess).

Original: Constraint satisfiability problem (CSP) solving networks disentangle the amount of work from the input size [24].
Revised: Constraint satisﬁability problem (CSP) solving networks disentangle the amount of work from the input size [24].

Original: Specifically, message passing neural networks can execute more passes to solve harder CSPs [24].
Revised: Speciﬁcally, message passing neural networks can execute more passes to solve harder CSPs [24].

Original: These systems are specific to the problem of constraint satisfaction for boolean expressions, but they are an early demonstration of scalable algorithmic behavior.
Revised: These systems are speciﬁc to the problem of constraint satisfaction for boolean expressions, but they are an early demonstration of scalable algorithmic behavior.

Original: Adaptive neural networks are designed to expend varying amounts of computation on different inputs, thus overcoming the limitation of classical RNNs.
Revised: Adaptive neural networks are designed to expend varying amounts of computation on different inputs, thus overcoming the limitation of classical RNNs.

Original: 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].
Revised: 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].

Original: Adaptive compute time (ACT) is an algorithm that provides RNNs with a halting unit, which estimates the probability that computation should continue.
Revised: Adaptive compute time (ACT) is an algorithm that provides RNNs with a halting unit, which estimates the probability that computation should continue.

Original: This algorithm penalizes ‘ponder time’ during training to encourage the network to solve problems quickly [8].
Revised: This algorithm penalizes ‘ponder time’ during training to encourage the network to solve problems quickly [8].

Original: Eyzaguirre and Soto [5] exhibit strong performance on visual question answering by introducing a differentiable version of ACT.
Revised: Eyzaguirre and Soto [5] exhibit strong performance on visual question answering by introducing a differentiable version of ACT.

Original: Adaptive transformer-based language models also exist, notably Universal Transformers and Depth-Adaptive Transformers, which utilize
Revised: Adaptive transformer-based language models also exist, notably Universal Transformers and Depth-Adaptive Transformers, which utilize

Original: ACT to determine the work required for each input [3, 18].
Revised: ACT to determine the work required for each input [3, 18].

Original: Similarly, iterative residual networks like
Revised: Similarly, iterative residual networks like

Original: NAIS-Nets repeat blocks in stages and perform well on image classification [2].
Revised: NAIS-Nets repeat blocks in stages and perform well on image classiﬁcation [2].

Original: 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.
Revised: 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.

Original: Logical extrapolation describes the task of generalizing to test sets which comprise more computationally complex samples than the training data.
Revised: Logical extrapolation describes the task of generalizing to test sets which comprise more computationally complex samples than the training data.

Original: Nuamah [16] 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.
Revised: Nuamah [16] 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.

Original: Similarly, Palm et al.
Revised: Similarly, Palm et al.

Original: [17] propose recurrent relational networks, which operate on graphs by iteratively passing messages.
Revised: [17] propose recurrent relational networks, which operate on graphs by iteratively passing messages.

Original: They also claim that classical architectures lack the inductive bias to reason about relationships between objects.
Revised: They also claim that classical architectures lack the inductive bias to reason about relationships between objects.

Original: Several recent works call these claims into question.
Revised: Several recent works call these claims into question.

Original: Schwarzschild et al.
Revised: Schwarzschild et al.

Original: [22] employ recurrent networks based on weight-sharing architectures, which can be made deeper at test time independent of the input size.
Revised: [22] employ recurrent networks based on weight-sharing architectures, which can be made deeper at test time independent of the input size.

Original: These systems exhibit logical extrapolation behavior in several domains.
Revised: These systems exhibit logical extrapolation behavior in several domains.

Original: More details on these methods are discussed in Section 3 as our algorithms build on these directly.
Revised: More details on these methods are discussed in Section 3 as our algorithms build on these directly.

Original: Banino et al.
Revised: Banino et al.

Original: [1] 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.
Revised: [1] 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 preﬁx sums.

Original: The ‘thinking’ systems proposed by Schwarzschild et al.
Revised: The ‘thinking’ systems proposed by Schwarzschild et al.

Original: [22] depart from classical recurrent networks for text, which learn from step-by-step supervision to produce output tokens one at a time.
Revised: [22] depart from classical recurrent networks for text, which learn from step-by-step supervision to produce output tokens one at a time.

Original: In contrast, ‘thinking’ systems autonomously synthesize a scalable algorithm end-to-end with no supervision over what each algorithmic step should do.
Revised: In contrast, ‘thinking’ systems autonomously synthesize a scalable algorithm end-to-end with no supervision over what each algorithmic step should do.

Original: Even more importantly, they can be applied to solve complex problems that are difficult or impossible to decompose by hand.
Revised: Even more importantly, they can be applied to solve complex problems that are difﬁcult or impossible to decompose by hand.

Original: In other words, sequence-to-sequence models, even adaptive RNNs, have a severe limitation in that the problem needs to be represented as a sequential input and that each token in the output is generated by the same function.
Revised: In other words, sequence-to-sequence models, even adaptive RNNs, have a severe limitation in that the problem needs to be represented as a sequential input and that each token in the output is generated by the same function.

Original: ‘Thinking’ systems, on the other hand, provide a mechanism for domains where this type of decomposition is unnatural, difficult, or even impossible.
Revised: ‘Thinking’ systems, on the other hand, provide a mechanism for domains where this type of decomposition is unnatural, difﬁcult, or even impossible.

Original:  Methods
Revised:  Methods

Original: We begin with some terms and definitions.
Revised: We begin with some terms and deﬁnitions.

Original: We study networks that share weights across blocks of layers during training.
Revised: We study networks that share weights across blocks of layers during training.

Original: For example, instead of three distinct residual blocks, a single residual block is repeated three times (see Figure 2 for a graphical depiction).
Revised: For example, instead of three distinct residual blocks, a single residual block is repeated three times (see Figure 2 for a graphical depiction).

Original: At test time, networks trained this way can be made “deeper” to extend their compute budget simply by repeating the block more times.
Revised: At test time, networks trained this way can be made “deeper” to extend their compute budget simply by repeating the block more times.

Original: We refer to the number of layers applied in a recurrent network as its “depth,” and this quantity grows as the number of recurrent iterations increases.
Revised: We refer to the number of layers applied in a recurrent network as its “depth,” and this quantity grows as the number of recurrent iterations increases.

Original: The number of feature maps produced by each layer (or the number of filters in a convolutional layer) is referred to as its “width.”
Revised: The number of feature maps produced by each layer (or the number of ﬁlters in a convolutional layer) is referred to as its “width.”

Original: More formally, let r be a function representing a recurrent block, e.g. a ResNet block [10], and let r n denote n recurrences of that function, e.g. r 2 ( x ) =
Revised: More formally, let r be a function representing a recurrent block, e.g. a ResNet block [10], and let r n denote n recurrences of that function, e.g. r 2 ( x ) =

Original: r ( r ( x )) .
Revised: r ( r ( x )) .

Original: Let ϕ denote a feature map, or an output of r , and let ϕ n = r n ( x ) .
Revised: Let  denote a feature map, or an output of r , and let  n = r n ( x ) .

Original: 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 to outputs.
Revised: We also consider an initial “embedding function” denoted by p , which projects an input instance into feature space, and also a ﬁnal “output head” denoted by h , which maps features to outputs.

Original: A Deep Thinking (DT) network with m iterations of the recurrent block can then be expressed as follows.
Revised: A Deep Thinking (DT) network with m iterations of the recurrent block can then be expressed as follows.

Original: [Equation]
Revised: [Equation]

Original: In our systems p comprises a convolutional layer followed by a ReLU, r is a single four-layer residual block, and h is a set of three convolutional layers with ReLUs after the first and second.
Revised: In our systems p comprises a convolutional layer followed by a ReLU, r is a single four-layer residual block, and h is a set of three convolutional layers with ReLUs after the ﬁrst and second.

Original: We fix m during training and compute gradients for optimization by backpropagating through the unrolled network.
Revised: We ﬁx m during training and compute gradients for optimization by backpropagating through the unrolled network.

Original: Then, m can be increased for testing, allowing these networks to increase their processing power and solve larger and harder problems.
Revised: Then, m can be increased for testing, allowing these networks to increase their processing power and solve larger and harder problems.

Original: Below, we consider two new approaches for achieving improved extrapolation: recall architectures and a modified optimization process.
Revised: Below, we consider two new approaches for achieving improved extrapolation: recall architectures and a modiﬁed optimization process.

Original: 3.1 Recall architectures
Revised: 3.1 Recall architectures

Original: When humans think for a long time to solve a problem, we often stop to reread the question or review the task at hand.
Revised: When humans think for a long time to solve a problem, we often stop to reread the question or review the task at hand.

Original: We improve DT architectures to periodically recall the input exactly.
Revised: We improve DT architectures to periodically recall the input exactly.

Original: We incorporate this capability into architectures by concatenating the input problem to the features output from each instance of the recurrent block.
Revised: We incorporate this capability into architectures by concatenating the input problem to the features output from each instance of the recurrent block.

Original: Popular architectures in computer vision typically incorporate skip connections that similarly pass information from earlier layers forward.
Revised: Popular architectures in computer vision typically incorporate skip connections that similarly pass information from earlier layers forward.

Original: In fact, empirical evidence suggests that skip connections, for example in highway networks, ResNets, and DenseNets, stabilize training [25, 10, 11].
Revised: In fact, empirical evidence suggests that skip connections, for example in highway networks, ResNets, and DenseNets, stabilize training [25, 10, 11].

Original: Our architectural modification is 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.
Revised: Our architectural modiﬁcation is 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.

Original: Recall allows the system to reproduce any missing or damaged features, and makes it impossible to “forget” the problem being solved.
Revised: Recall allows the system to reproduce any missing or damaged features, and makes it impossible to “forget” the problem being solved.

Original: To formalize this architectural change, adding recall to the network can be expressed using the notation defined above as follows.
Revised: To formalize this architectural change, adding recall to the network can be expressed using the notation deﬁned above as follows.

Original: [Equation]
Revised: [Equation]

Original: Whereas the input to r at iteration k is usually ϕ k − 1 , with recall, the input to r at iteration k is [ ϕ k − 1 , x ] , or the concatenation of the input with the feature map output by the previous recurrence.
Revised: Whereas the input to r at iteration k is usually  k  1 , with recall, the input to r at iteration k is [  k  1 , x ] , or the concatenation of the input with the feature map output by the previous recurrence.

Original: We add a single convolutional layer to map the input [ ϕ k − 1 , x ] to a feature map of the same shape as ϕ k − 1 .
Revised: We add a single convolutional layer to map the input [  k  1 , x ] to a feature map of the same shape as  k  1 .

Original: We refer to DT networks with concatenating skip connections as DT-Recall models.
Revised: We refer to DT networks with concatenating skip connections as DT-Recall models.

Original: 3.2 Promoting forward progress through optimization
Revised: 3.2 Promoting forward progress through optimization

Original: We propose a training objective to encourage the system to incrementally make progress from any starting point.
Revised: We propose a training objective to encourage the system to incrementally make progress from any starting point.

Original: We do this by inputting a problem instance and running the recurrent module for some random number of iterations.
Revised: We do this by inputting a problem instance and running the recurrent module for some random number of iterations.

Original: 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.
Revised: 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.

Original: This incremental training process has two benefits.
Revised: This incremental training process has two beneﬁts.

Original: 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.
Revised: 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.

Original: 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-specific behaviors, such as behaviors that get executed only on iteration five, for example.
Revised: 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-speciﬁc behaviors, such as behaviors that get executed only on iteration ﬁve, for example.

Original: Rather, the network is encouraged to learn iteration-agnostic behaviors that are effective at any stage of the problem solving process.
Revised: Rather, the network is encouraged to learn iteration-agnostic behaviors that are effective at any stage of the problem solving process.

Original: Input: parameter vector θ , integer m , weight α for batch_idx = 1, 2, ... do Choose n ∼ U { 0 , m − 1 } and k ∼ U { 1 , m − n } Compute ϕ n with n iterations w/o tracking gradientsCompute
Revised: Input: parameter vector ✓ , integer m , weight ↵ for batch_idx = 1, 2, ... do Choose n ⇠ U { 0 , m  1 } and k ⇠ U { 1 , m  n } Compute  n with n iterations w/o tracking gradients

Original:  ˆ y prog with additional k iterations Compute ˆ y m with new forward pass of m iterations Compute L max_iters with ˆ y m and L progressive with ˆ y prog .
Revised: Compute ˆ y prog with additional k iterations Compute ˆ y m with new forward pass of m iterations Compute L max_iters with ˆ y m and L progressive with ˆ y prog .

Original: Compute L = (1 − α ) · L max_iters + α · L progressive Compute ∇ θ L and update θ end for
Revised: Compute L = (1  ↵ ) · L max_iters + ↵ · L progressive Compute r ✓ L and update ✓ end for

Original: 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 to improve this partial solution.
Revised: 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 to improve this partial solution.

Original: We then update the network’s parameters to minimize loss after n + k total iterations when it starts with the partial solution.
Revised: We then update the network’s parameters to minimize loss after n + k total iterations when it starts with the partial solution.

Original: This is done by detaching the recurrent module’s output after n iterations from the computation graph before computing the gradient of the loss at iteration n + k .
Revised: This is done by detaching the recurrent module’s output after n iterations from the computation graph before computing the gradient of the loss at iteration n + k .

Original: The process above is an analog of truncated backpropagation through time [12], with a random starting and end point.
Revised: The process above is an analog of truncated backpropagation through time [12], with a random starting and end point.

Original: 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 .
Revised: During training, we ensure that the sum of n and k is less than a ﬁxed maximum number of iteration m , which we call the training regime .

Original: 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.
Revised: 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.

Original: The training loop and computation of the loss are given in Algorithm 1.
Revised: The training loop and computation of the loss are given in Algorithm 1.

Original: The incremental progress term is referred to as L progressive (or progressive loss ), and the contribution to the loss from the fully unrolled network is denoted by L max_iters .
Revised: The incremental progress term is referred to as L progressive (or progressive loss ), and the contribution to the loss from the fully unrolled network is denoted by L max_iters .

Original: See Appendix A.3 for more details.
Revised: See Appendix A.3 for more details.

Original: 3.3 Datasets
Revised: 3.3 Datasets

Original: We evaluate our methods on the benchmark problems available in the Python package easy-to-hard-data , which generates reasoning problems of various difficulties.
Revised: We evaluate our methods on the benchmark problems available in the Python package easy-to-hard-data , which generates reasoning problems of various difﬁculties.

Original: The three problems considered are computing prefix sums, finding the optimal path in a two-dimensional maze, and solving chess puzzles.
Revised: The three problems considered are computing preﬁx sums, ﬁnding the optimal path in a two-dimensional maze, and solving chess puzzles.

Original: We briefly review the input and output structures for each problem and refer the reader to Schwarzschild et al.
Revised: We brieﬂy review the input and output structures for each problem and refer the reader to Schwarzschild et al.

Original: [21] for more detail, including the data generation process.
Revised: [21] for more detail, including the data generation process.

Original: Note that the architectures we consider are fully convolutional, and produce outputs of the same dimension as their inputs.
Revised: Note that the architectures we consider are fully convolutional, and produce outputs of the same dimension as their inputs.

Original: Furthermore, the training and testing datasets we consider have labels of the same dimension as their inputs.
Revised: Furthermore, the training and testing datasets we consider have labels of the same dimension as their inputs.

Original: Therefore, a network trained on inputs of one size can then trivially be applied to inputs of a different size.
Revised: Therefore, a network trained on inputs of one size can then trivially be applied to inputs of a different size.

Original: We begin with the toy problem of computing prefix sums modulo two.
Revised: We begin with the toy problem of computing preﬁx sums modulo two.

Original: The inputs and targets for this problem consist of bit strings.
Revised: The inputs and targets for this problem consist of bit strings.

Original: The i th bit of the target is the mod two sum of all bits prior to and including the i th bit of the input.
Revised: The i th bit of the target is the mod two sum of all bits prior to and including the i th bit of the input.

Original: We control the “difficulty” of the problem by changing the length of the bit string.
Revised: We control the “difﬁculty” of the problem by changing the length of the bit string.

Original: Note that computing prefix sums of greater length is known to require a greater number of sequential operations [15].
Revised: Note that computing preﬁx sums of greater length is known to require a greater number of sequential operations [15].

Original: 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.
Revised: 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.

Original: The mazes we 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.
Revised: The mazes we 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.

Original: 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.
Revised: 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.

Original: We make more challenging datasets by increasing the size of the mazes.
Revised: We make more challenging datasets by increasing the size of the mazes.

Original: All training in our work is done on 9 × 9 mazes.
Revised: All training in our work is done on 9 ⇥ 9 mazes.

Original: Despite this small training size, we benchmark extrapolation performance on mazes of size 59 × 59 , and observe models solving mazes as large as 201 × 201 with high success rates.
Revised: Despite this small training size, we benchmark extrapolation performance on mazes of size 59 ⇥ 59 , and observe models solving mazes as large as 201 ⇥ 201 with high success rates.

Original: The chess puzzles we consider are mid-game boards represented by twelve 8 × 8 planes indicating the positions of all 12 types of pieces (there are six distinct piece classes and two colors on a chess board).
Revised: The chess puzzles we consider are mid-game boards represented by twelve 8 ⇥ 8 planes indicating the positions of all 12 types of pieces (there are six distinct piece classes and two colors on a chess board).

Original: The goal of this task is to find the best next move, and each target encodes this information in an 8 × 8 array with zeros everywhere except the origin and destination of the piece to be moved, which are populated with ones.
Revised: The goal of this task is to ﬁnd the best next move, and each target encodes this information in an 8 ⇥ 8 array with zeros everywhere except the origin and destination of the piece to be moved, which are populated with ones.

Original: The chess dataset is sorted by difficulty rating, as determined by Elo scores computed via human trials on Lichess.org.
Revised: The chess dataset is sorted by difﬁculty rating, as determined by Elo scores computed via human trials on Lichess.org.

Original: 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.
Revised: Problems are sorted by difﬁculty and we train on the ﬁrst 600K easiest puzzles and test our models on indices ranging from 700K to 1.1M.

Original: Our experiments have two categories: (I) Testing extrapolation at extreme leaps in problem difficulty/size and (II) Testing scenarios that emphasize the overthinking trap and how we avoid it.
Revised: Our experiments have two categories: (I) Testing extrapolation at extreme leaps in problem difﬁculty/size and (II) Testing scenarios that emphasize the overthinking trap and how we avoid it.

Original: The following two sections describe each of these experiments, respectively.
Revised: The following two sections describe each of these experiments, respectively.

Original:  Extreme extrapolation
Revised:  Extreme extrapolation

Original: In each problem domain, we show that our models with recall and trained with progressive loss perform best.
Revised: In each problem domain, we show that our models with recall and trained with progressive loss perform best.

Original: We compare our models to a baseline of DT models trained with m = 30 and problemspecific widths (detailed below), which are both wider and deeper than the models previously studied.
Revised: We compare our models to a baseline of DT models trained with m = 30 and problemspeciﬁc widths (detailed below), which are both wider and deeper than the models previously studied.

Original: 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.
Revised: However, to provide a fair comparison, we use consistent values for these parameters across our experiments since we ﬁnd that all architectures beneﬁt from additional width and depth.

Original: We also compare to feed-forward models without weight sharing of the same effective depth as a 30-iteration
Revised: We also compare to feed-forward models without weight sharing of the same effective depth as a 30-iteration

Original: DT network, which are ResNets [10].
Revised: DT network, which are ResNets [10].

Original: We include ablation studies lending credence to our claim that both the architecture and the loss modifications are instrumental in boosting extrapolation power.
Revised: We include ablation studies lending credence to our claim that both the architecture and the loss modiﬁcations are instrumental in boosting extrapolation power.

Original: For a discussion on hyperparameter choice and training set-ups for all of our experiments, see Appendix
Revised: For a discussion on hyperparameter choice and training set-ups for all of our experiments, see Appendix

Original: A.4.
Revised: A.4.

Original: Also, code to reproduce all of our experiments is publicly available.
Revised: Also, code to reproduce all of our experiments is publicly available.

Original: 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.
Revised: Preﬁx 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.

Original: 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 α = 1 when progressive loss is used.[Figure]
Revised: We train our models on 32-bit strings with a maximum of only m = 30 recurrent iterations (indicated in the ﬁgures as the shaded train regime section), 400 convolutional ﬁlters per layer, and incremental loss weight ↵ = 1 when progressive loss is used.

Original: inputs extrapolate to 512-bit data.
Revised: inputs extrapolate to 512-bit data.

Original: The value of our recall and progressive loss is clear by how quickly and accurately our models solve this very large problem.
Revised: The value of our recall and progressive loss is clear by how quickly and accurately our models solve this very large problem.

Original: [Figure]

Original: Mazes this large cannot be solved without recall, and furthermore progressive loss leads to more accurate models.
Revised: and furthermore progressive loss leads to more accurate models.

Original: 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-Recall models with progressive loss can solve more than 97% after approximately 200 iterations.
Revised: 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-Recall models with progressive loss can solve more than 97% after approximately 200 iterations.

Original: To better understand the individual effects of our proposed approach, we perform an ablation study.
Revised: To better understand the individual effects of our proposed approach, we perform an ablation study.

Original: Figure 3 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.
Revised: Figure 3 also makes clear that DT-Recall networks trained without progressive loss achieve the same ﬁnal accuracy as models trained with it; but they require approximately twice the number of iterations to get there.

Original: 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.
Revised: 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 speciﬁcally.

Original: 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.
Revised: 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.

Original: One other study in this domain reveals the importance of randomly setting n in Algorithm 1.
Revised: One other study in this domain reveals the importance of randomly setting n in Algorithm 1.

Original: To show this, we modify the algorithm slightly, to always use n = 0 , rather than randomly choosing that value.
Revised: To show this, we modify the algorithm slightly, to always use n = 0 , rather than randomly choosing that value.

Original: While this change seems minor, and indeed makes a small impact on the results, it does lower the accuracy.
Revised: While this change seems minor, and indeed makes a small impact on the results, it does lower the accuracy.

Original: Specifically, DT networks with recall on average solve 90.27% of 512-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.12% accuracy.
Revised: Speciﬁcally, DT networks with recall on average solve 90.27% of 512-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.12% accuracy.

Original: From this ablation experiment, we conclude that randomly sampling these values is, in fact, beneficial.
Revised: From this ablation experiment, we conclude that randomly sampling these values is, in fact, beneﬁcial.

Original: Mazes As with the prefix sum problem, we can improve performance on hard mazes by combining incremental training and recall architectures.
Revised: Mazes As with the preﬁx sum problem, we can improve performance on hard mazes by combining incremental training and recall architectures.

Original: In particular, while Figure 3 may not convince the reader that our proposed loss is critical since DT-Recall models without it perform very well, with more complex data a drastic difference emerges.
Revised: In particular, while Figure 3 may not convince the reader that our proposed loss is critical since DT-Recall models without it perform very well, with more complex data a drastic difference emerges.

Original: We show in Figure 4 that on the significantly harder test set of 59 × 59 mazes, our models exhibit strong algorithmic extrapolation, while previous methods, both feed-forward and DT systems, completely fail.
Revised: We show in Figure 4 that on the signiﬁcantly harder test set of 59 ⇥ 59 mazes, our models exhibit strong algorithmic extrapolation, while previous methods, both feed-forward and DT systems, completely fail.

Original: 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.
Revised: Not only do our models achieve a higher peak accuracy, but they do not overthink, as can be seen by the ﬂat spans in the curves in Figure 4.

Original: In fact, we can push these systems to their limits (and the limits of our hardware) and find that our models can solve 74% of the 201 × 201 mazes.
Revised: In fact, we can push these systems to their limits (and the limits of our hardware) and ﬁnd that our models can solve 74% of the 201 ⇥ 201 mazes.

Original: 4 See
Revised: 4 See

Original: Appendices A.9 and A.10 for examples of 201 × 201 and 801 × 801 mazes, which our models solve using 2,000 and 20,000 iterations, respectively (the equivalent of  layers  depth).[Figure]
Revised: Appendices A.9 and A.10 for examples of 201 ⇥ 201 and 801 ⇥ 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).

Original: The best models in Figure 4 are DT-Recall models trained with 30 maximum iterations and a weighting in the loss of α = 0 . 
Revised: The best models in Figure 4 are DT-Recall models trained with 30 maximum iterations and a weighting in the loss of ↵ = 0 . 01 .

Original: In order to see how critical recall is in overcoming the overthinking problem, we show in the same figure that without recall all models fail to extrapolate.
Revised: In order to see how critical recall is in overcoming the overthinking problem, we show in the same ﬁgure that without recall all models fail to extrapolate.

Original: The benefit of progressive loss is also highlighted by the fact that DT-Recall models with progressive loss achieve on average 97% accuracy while DT-Recall models without progressive loss only reach an average of 83%.
Revised: The beneﬁt of progressive loss is also highlighted by the fact that DT-Recall models with progressive loss achieve on average 97% accuracy while DT-Recall models without progressive loss only reach an average of 83%.

Original: Interestingly, in this domain, we find that the best model has α = 0 .
Revised: Interestingly, in this domain, we ﬁnd that the best model has ↵ = 0 .

Original: 01 in the loss, a much smaller weight than we use for prefix sums.
Revised: 01 in the loss, a much smaller weight than we use for preﬁx sums.

Original: More on finding the optimal value for α can be found in
Revised: More on ﬁnding the optimal value for ↵ can be found in

Original: Appendix A.4.
Revised: Appendix A.4.

Original: Chess Puzzles We further find that using recall and progressive loss yields notable improvement on chess puzzles.
Revised: Chess Puzzles We further ﬁnd that using recall and progressive loss yields notable improvement on chess puzzles.

Original: In Figure 5, we see that our techniques lead to a 4% accuracy improvement compared to networks without them.
Revised: In Figure 5, we see that our techniques lead to a 4% accuracy improvement compared to networks without them.

Original: We show that either removing recall or training without the incremental progress loss will hurt performance.
Revised: We show that either removing recall or training without the incremental progress loss will hurt performance.

Original: 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 α = 0 .
Revised: Our best models are DT-Recall networks with 512 convolutional ﬁlters in each layer trained with a maximum of 30 iterations and a weight of ↵ = 0 .

Original: Moreover, their accuracy is preserved as the number of iterations increases, while that of the other DT networks decays seriously after aboutiterations – the sign of overthinking and the subject of the following section.
Revised: Moreover, their accuracy is preserved as the number of iterations increases, while that of the other DT networks decays seriously after aboutiterations – the sign of overthinking and the subject of the following section.

Original: Results from tests on harder puzzles are in Appendix A.5.
Revised: Results from tests on harder puzzles are in Appendix A.5.

Original:  The overthinking problem
Revised:  The overthinking problem

Original: Deep Thinking networks boast impressive capabilities to solve harder problems by thinking deeper, but they are prone to overthinking.
Revised: Deep Thinking networks boast impressive capabilities to solve harder problems by thinking deeper, but they are prone to overthinking.

Original: In particular, by testing models on data closer in complexity to the training data we can better compare our methods to prior work.
Revised: In particular, by testing models on data closer in complexity to the training data we can better compare our methods to prior work.

Original: Figure 6 shows that some recurrent networks that can perform algorithmic extrapolation may collapse entirely when performing too many iterations.[Figure]
Revised: Figure 6 shows that some recurrent networks that can perform algorithmic extrapolation may collapse entirely when performing too many iterations.

Original: Our models do not suffer from the sharp decline in accuracy as the number of iterations increases.
Revised: Our models do not suffer from the sharp decline in accuracy as the number of iterations increases.

Original: 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.
Revised: 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.

Original: See Figure 5 for an example where progressive loss is needed too.
Revised: See Figure 5 for an example where progressive loss is needed too.

Original: The monotonic increase in accuracy with added iterations in our models is practically useful, as it allows for pre-defining a large iteration number at which to terminate the recurrence rather than having to carefully choose a stopping iteration to maximize performance while avoiding degradation.
Revised: The monotonic increase in accuracy with added iterations in our models is practically useful, as it allows for pre-deﬁning a large iteration number at which to terminate the recurrence rather than having to carefully choose a stopping iteration to maximize performance while avoiding degradation.

Original: One observation that can be 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.
Revised: One observation that can be 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.

Original: Another way to describe this is that our models seem to converge to a fixed-point solution as they iterate rather than becoming unstable.
Revised: Another way to describe this is that our models seem to converge to a ﬁxed-point solution as they iterate rather than becoming unstable.

Original: 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.
Revised: 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.

Original: In Appendix A.8, we present similar findings for prefix sums.
Revised: In Appendix A.8, we present similar ﬁndings for preﬁx sums.

Original: 5.1 Manipulating feature maps
Revised: 5.1 Manipulating feature maps

Original: First, we investigate sensitivity to adding noise in the feature maps before concatenating the inputs.
Revised: First, we investigate sensitivity to adding noise in the feature maps before concatenating the inputs.

Original: We examine model behavior when we add Gaussian noise (with mean zero and standard deviation one) to the features after one iteration of maze solving.
Revised: We examine model behavior when we add Gaussian noise (with mean zero and standard deviation one) to the features after one iteration of maze solving.

Original: Models with recall can still solve 13 ×mazes even when we perturb the features, but models without recall cannot.
Revised: Models with recall can still solve 13 ⇥mazes even when we perturb the features, but models without recall cannot.

Original: See Appendix A.8 for plots, tables, and further discussion.
Revised: See Appendix A.8 for plots, tables, and further discussion.

Original: Next, we ask whether the initial feature maps (after the projection layer p ) carry any important information.
Revised: Next, we ask whether the initial feature maps (after the projection layer p ) carry any important information.

Original: 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 .
Revised: 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 .

Original: In this case, our DT-recall models naturally regenerate their features and recover to solve the problem again, indicating that the learned algorithm is able to find a solution using the input without the initial projection.
Revised: In this case, our DT-recall models naturally regenerate their features and recover to solve the problem again, indicating that the learned algorithm is able to ﬁnd a solution using the input without the initial projection.

Original: See Figure 12 in Appendix A.8.
Revised: See Figure 12 in Appendix A.8.

Original: 5.2 Manipulating inputs
Revised: 5.2 Manipulating inputs

Original: With the notion that our models may embody a convergent process, we turn our attention to investigating how the networks determine when to stop manipulating the representation.
Revised: With the notion that our models may embody a convergent process, we turn our attention to investigating how the networks determine when to stop manipulating the representation.

Original: Perhaps there is something in the feature maps output by late iterations that tells the network to stop working.
Revised: Perhaps there is something in the feature maps output by late iterations that tells the network to stop working.

Original: We can test this in two ways.
Revised: We can test this in two ways.

Original: 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.[Figure]
Revised: We perturb the input (which we concatenate onto the features) after some number of iterations – ﬁrst subtly, then by swapping it with an entirely different example.

Original: To start, we flip single bits in the input string for prefix sums.
Revised: To start, we ﬂip single bits in the input string for preﬁx sums.

Original: We explore the response of the model to flipping bits at different indices after 50 iterations (when the models have already solved the initial 48bit problem).
Revised: We explore the response of the model to ﬂipping bits at different indices after 50 iterations (when the models have already solved the initial 48bit problem).

Original: The left panel of Figure 7 shows recovery time as a function of index of the flipped bit.
Revised: The left panel of Figure 7 shows recovery time as a function of index of the ﬂipped bit.

Original: We observe that our models can indeed recover from single bit flips, and recovery time decays linearly with index.
Revised: We observe that our models can indeed recover from single bit ﬂips, and recovery time decays linearly with index.

Original: Higher indices are closer to the end and affect fewer bits in the output than lower ones.
Revised: Higher indices are closer to the end and affect fewer bits in the output than lower ones.

Original: Similarly, with mazes, we change the input by moving the end of the maze two steps closer to the start.
Revised: Similarly, with mazes, we change the input by moving the end of the maze two steps closer to the start.

Original: We use this new input concatenated with features generated after 50 iterations of solving the original maze.
Revised: We use this new input concatenated with features generated after 50 iterations of solving the original maze.

Original: In this case, we see in Figure 13 (appendix) that models with recall self-correct and solve the new problem in very few iterations.
Revised: In this case, we see in Figure 13 (appendix) that models with recall self-correct and solve the new problem in very few iterations.

Original: 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?
Revised: 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?

Original: 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.
Revised: 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.

Original: The center panel of
Revised: The center panel of

Original: Figure 7 shows the effect of swapping feature maps.
Revised: Figure 7 shows the effect of swapping feature maps.

Original: 5.3 Converging to a fixed point
Revised: 5.3 Converging to a ﬁxed point

Original: We can also study the convergence by measuring the change in the output at each iteration.
Revised: We can also study the convergence by measuring the change in the output at each iteration.

Original: 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 it will hold onto this representation (or stay nearby) once it is reached.
Revised: 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 it will hold onto this representation (or stay nearby) once it is reached.

Original: Moreover, we seek to qualitatively categorize each model type (architecture/loss pair) as convergent or non-convergent.
Revised: Moreover, we seek to qualitatively categorize each model type (architecture/loss pair) as convergent or non-convergent.

Original: To do we measure the change in the solution at each iteration with ℓ 2 -norm of the difference.
Revised: To do we measure the change in the solution at each iteration with ` 2 -norm of the difference.

Original: The right-hand panel of Figure 7 shows ∆ ϕ ( n ) :
Revised: The right-hand panel of Figure 7 shows   ( n ) :

Original: = ∥ ϕ n − ϕ n − 1 ∥ ℓ 2 .
Revised: = k  n   n  1 k ` 2 .

Original: We see that our models appear to converge, while DT networks without recall explode, providing another view into overthinking.
Revised: We see that our models appear to converge, while DT networks without recall explode, providing another view into overthinking.

Original:  Conclusion
Revised:  Conclusion

Original: In this work, we improve the algorithmic extrapolation power of neural networks.
Revised: In this work, we improve the algorithmic extrapolation power of neural networks.

Original: We propose an architecture and a loss that lead to a gigantic generalization leap from easy training data to much more complex testing examples.
Revised: We propose an architecture and a loss that lead to a gigantic generalization leap from easy training data to much more complex testing examples.

Original: Furthermore, we show that our models avoid the overthinking trap.
Revised: Furthermore, we show that our models avoid the overthinking trap.

Original: 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.
Revised: We test algorithmic extrapolation with chess puzzles where the spacial dimension is consistent across difﬁculties and with maze solving and preﬁx sum computation where our models can extrapolate to larger problems as well.

Original: 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.
Revised: 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.

Original: Existing methods deteriorate at such large numbers of iterations creating the need for stopping mechanisms.
Revised: Existing methods deteriorate at such large numbers of iterations creating the need for stopping mechanisms.

Original: Even for small numbers of iterations, if the halting decision is sub-optimal, then the output may also be sub-optimal.
Revised: Even for small numbers of iterations, if the halting decision is sub-optimal, then the output may also be sub-optimal.

Original: 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.
Revised: Because our models have the desirable property (for an algorithm) that they converge to ﬁxed points, they run no such risk, require no learned or hand-crafted stopping rules, and can perform well at huge numbers of iterations.

Original: 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.
Revised: Learning scalable processes and performing logical extrapolation are difﬁcult 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.

Original: Importantly, our neural networks learn end-to-end how to perform algorithmic extrapolation.
Revised: Importantly, our neural networks learn end-to-end how to perform algorithmic extrapolation.

Original: Our findings are limited to the domains we consider and further exploration into how our methods will perform in other settings is needed.
Revised: Our ﬁndings are limited to the domains we consider and further exploration into how our methods will perform in other settings is needed.

Original: The application of machine learning systems to real world data demands effective extrapolation, even though this is often overlooked.
Revised: The application of machine learning systems to real world data demands effective extrapolation, even though this is often overlooked.

Original: 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.
Revised: Curated research benchmarks designed for measuring generalization (e g. ImageNet [4]) omit a critical feature of real data in the ﬁeld: there is often reason to think that data encountered after training could embody more difﬁcult examples than those in the training set.

Original: Consider that a practitioner cannot always know a priori the range of complexity that exists in their domain.
Revised: Consider that a practitioner cannot always know a priori the range of complexity that exists in their domain.

Original: 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.
Revised: 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.

Original: Acknowledgements
Revised: Acknowledgements

Original: This work was supported by the ONR MURI Program, The National Science Foundation DMS- and the Office of Naval Research (#4720008163)A Appendix A.1 Architectures
Revised: This work was supported by the ONR MURI Program, The National Science Foundation DMS- and the Ofﬁce of Naval Research (#4720008163)

Original: The architectures we use in our experiments are very similar to those described by Schwarzschild et al.

Original: [22], but we provide the details here for convenience.

Original: In all the models, the first layer p is a convolutional layer with 3 × 3 filters (or filters of length three in the 1D case) followed by a ReLU non-linearity, which projects the input into feature space.

Original: Each filter strides by one, and inputs are padded with one unit in each direction.

Original: The number of filters is specified per dataset in Table 1 as the “width” of the network.

Original: The internal blocks are standard residual blocks with four convolutional layers (followed by ReLUs) and skip connections every two layers [10].

Original: These blocks share weights in recurrent networks and are simply repeated with distinct parameters in feed-forward models.

Original: Models with recall have additional convolutional layers that map the concatenated inputs and features [ ϕ, x ] to the shape of the feature maps ϕ .

Original: The feature maps have the same spacial dimension as the input and w channels, where w denotes the width of the model.

Original: The final block h is composed of three convolutional layers with decreasing widths (specific numbers are in the Table 1), and ReLUs after the first two.

Original: The third and final convolutional layer in h has two channel outputs used for binary pixel

Original: [Table]

Original: A.2 Computational resources

Original: We use Nvidia GeForce RTX 2080Ti GPUs for all of our experiments.

Original: The models that solve prefix sums train in under an hour on a single GPU.

Original: The maze models take up to eight hours to train on one

Original: GPU.

Original: Finally, the chess puzzle networks were trained using four GPUs at once and can take between 24 and 48 hours.

Original: A.3 The loss

Original: Prefix sum problems with n bits are input to the model as a vector x ∈ { 0 , 1 } n , and the output is denoted by ˆ y ∈ R n × 2 .

Original: The target output is y ∈ { 0 , 1 } n .

Original: We compute the loss as follows.

Original: [Equation]

Original: where [ · , · ] indexes the output array.

Original: Therefore, for a batch of B instances, the total loss is

Original: [Equation]

Original: This loss applies to all three problem types we consider.

Original: Note that a maze represented by an n × n ×input has n 2 pixels and the loss can be averaged over those pixels.

Original: The same applies to chess, where there are always 64 pixels.

Original: When computing the progressive loss in Algorithm 1, we compute L max_iters and L progressive as follows.

Original: [Equation]

Original: Note, in Algorithm 1, when the weight α is equal to 1 , there is no contribution from the full forward pass and also that α = 0 corresponds to full backpropagation through all iterations with no incremental objective.

Original: [Figure]

Original: A reviewer has brought up the interesting comparison of our method to a loss that directly penalizes the run time by weighting the loss with increasing coefficients as a function of iteration.

Original: We have run some preliminary experiments, where we find models trained this way exhibit worse generalization to test sets which are harder than their training data.

Original: Specifically, on prefix sums, these models solve only 72 .

Original: 45 ± 8 .

Original: 91% of 512-bit inputs.

Original: On maze data, these models were not able to solve any 59x mazes in our test set, and we note that the training accuracy is >99% and the accuracy on 33x mazes is 93.8%, indicating that training was successful.

Original: We believe a more thorough exploration of

Original: Deep Thinking networks with variants of such run-time penalties would be interesting future work.

Original: A.4 Hyperparameters

Original: We describe the training set-up for each experiment and discuss hyperparameter choices.

Original: All prefix sum models are optimized using Adam [14] and gradient clipping.

Original: Maze models are also trained with Adam, but no clipping is used.

Original: Chess models train stably with SGD without gradient clipping.

Original: All training is done with a weight decay coefficient of 0.0002 and training with SGD uses a momentum coefficient of 0.9.

Original: In each training run, we hold out 20% of the training data to use as a validation set, and we save for testing the checkpoint with the highest validation accuracy.

Original: Coarse experimentation with learning rates and decay schedules is used to determine optimal choices of these hyperparameters for reproducibility and speed of training.

Original: Models were trained with exponential warm-up for a small number of epochs before the main training routine.

Original: In Table 2, we present the training hyperparameters shared among models presented in Section 4 for each problem instance.

Original: [Table]

Original: We perform width and depth experiments to demonstrate that our best prefix sum models are roughly on a plateau in hyperparameter space with respect to performance on 512-bit data.

Original: We see from

Original: Figure A.4 that performance increases steadily with depth, but no additional performance is achieved by increasing m beyond 30.

Original: From Figure 8, we also see performance steadily increases with network width.

Original: We choose to use a width of 400 filters per layer and m = 30 because this adequately solves 512-bit data.

Original: These experiments are performed on prefix sums (since they are realtively cheap to run) to exemplify the general correlation between depth, width, and performance observed on all tasks.

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

Original: We determine the optimal weight to be used it the combination of the loss terms using a coarse grid search.

Original: For prefix sums we test models with α values in [0 , 0 . 25 , 0 . 5 , 0 .

Original: 75 , 1] , for mazes we test values [0 , 0 .

Original: 01 , 0 . 1 , 0 . 5 , 1] , and for chess we test values [0 , 0 .

Original: A reviewer has pointed out the possibility of exploring a modified version of progressive loss, where n in Algorithm 1 is always equal to zero.

Original: Without any further modification, this would mean k is sampled from the larger interval { 1 , m } and would require more computation on average due to the increased expectation for number of backward passes.

Original: However, one could then sample k from a smaller interval to match the expectation in the case where n is sampled from { 1 , m } , which would save computation on forward passes.

Original: We believe such optimizations could be an interesting avenue for future work, but we note that for the case of prefix sums a preliminary evaluation of the expensive version of the n = 0 method produces a less effective learned algorithm.

Original: Specifically, DT nets with recall achieve 90 .

Original: 27 ± 5 .

Original: 95% on 512-bit data with the n = 0 loss, and with the sampling of n used in Algorithm 1, they achieve 97 .

Original: 12 ± 1 . 88% .

Original: A.5 Extended results

Original: Tables 3, 4, and 5 show the peak accuracy of each of the curves presented in the plots in Section 4.

Original: Since we aim to learn an algorithm for each of the three tasks ( i.e. the model should work for all training examples), we use a strict training accuracy convergence criteria.

Original: Models that do not reach a training accuracy of at least 99% are considered unconverged and are removed from the averages in the tables below to maintain fair comparisons across the various experimental groups.

Original: For prefix sums, 6 out of 20 of the models in the progressive loss with no recall category were filtered out this way, and no models were filtered for the other categories.

Original: For mazes, 1 out every 5 models trained with recall were filtered out this way, and no models were filtered the other three categories.

Original: For all categories, about half of the chess models did not converge.

Original: We would like to stress that this filtering was done solely based on training accuracy, which is standard practice.

Original: We would like to highlight that the additional performance gained from the recall architecture is not due to the marginally higher parameter count.

Original: Schwarzschild et al.

Original: [23] explore the effect of additional depth on the performance of vanilla DT nets, and the increased performance of DT nets with recall far outweighs the benefit from additional capacity of deeper or wider models.

Original: [Table]

Original: [Table]

Original: Figure 9 shows the test performance of chess models on even harder test sets.

Original: It is interesting to note that recall and our loss both still help dramatically, but there is an apparent ceiling on

Original: [Table]

Original: [Figure]

Original: generalizing from the easy puzzles to much harder ones.

Original: We leave investigation into this phenomenon for future work.

Original: A.6 In-distribution results

Original: Information on how models perform in distribution reveals that each class of model fits the training data and generalizes in the classical sense.

Original: Specifically, we report the accuracy on a held-out validation set from the same problem difficulty as the training data.

Original: For recurrent models, we evaluate the performance using the maximum number of iterations used during training, irrespective of the training objective.

Original: See Table 6.

Original: A.7 Hard to easy

Original: One natural query about our models is how well they perform on test sets that are easier/smaller than the data used for training.

Original: Figure 10 shows that recurrent models can generally solve easier/smaller problems in fewer iterations.

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

Original: [Figure]

Original: [Table]

Original: A.8 More on the overthinking problem

Original: We present additional plots and tables to give a more complete picture of how some models overcome the overthinking problem presented in Section 5.

Original: Table 7 shows the results from several experiments.

Original: The first column indicates the average number of iterations to solve a maze and the best accuracy (over all the test iterations considered).

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

Original: Next, we see the same trend when we replace the features at the first iteration with arrays of zeros (‘Zeros’ column).

Original: Finally, we swap the features after 50 iterations with those obtained from solving an entirely different maze, and again the models without recall cannot recover from this, but those with recall achieve near perfect accuracy.

Original: Figures 11-13 shows accuracy curves for these experiments.

Original: Figure 14 shows the norm of the difference in features when we input random noise instead of actual mazes.

Original: This plot shows that the stable behavior of our models is not limited to the portion of input space occupied by real mazes, rather the process they learn exhibits convergent behavior even on random inputs.

Original: Similar experiments on prefix sum computation reveal that recurrent models with recall seem to converge as well.

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Table]

Original: A.9 A fun maze