T HINKING D EEPER W ITH R ECURRENT N ETWORKS : L OGICAL E XTRAPOLATION W ITHOUT O VERTHINKING
Anonymous authors
Paper under double-blind review
A BSTRACT
Classical machine learning systems perform best when they are trained and tested on the same distribution, and they lack a mechanism to increase model power after training is complete.
In contrast, recent work has observed that recurrent networks can exhibit logical extrapolation; models trained only on small/simple problem instances can extend their abilities to solve large/complex instances at test time simply by performing more recurrent iterations.
While preliminary results on these “thinking systems” are promising, existing recurrent systems, when iterated many times, often collapse rather than improve their performance.
This “overthinking” phenomenon has prevented thinking systems from scaling to particularly large and complex problems.
In this paper, we design a recall architecture that keeps an explicit copy of the problem instance in memory so that it cannot be forgotten.
We also propose an incremental training routine that prevents the model from learning behaviors that are speciﬁc to iteration number and instead pushes it to learn behaviors that can be repeated indeﬁnitely.
Together, these design choices encourage models to converge to a steady state solution rather than deteriorate when many iterations are used.
These innovations help to tackle the overthinking problem and boost deep thinking behavior on each of the benchmark tasks proposed by Schwarzschild et al.
(2021a).
I NTRODUCTION
Humans solve complex logical reasoning problems through the process of logical extrapolation – they assemble simple logical primitives into complex strategies.
For example, a person taught to prove simple lemmas 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 lack the ability to solve complex reasoning tasks in a scalable way.
Recently, deep thinking systems (Schwarzschild et al., 2021b) have been proposed as a way to represent and learn scalable reasoning processes using recurrent neural networks.
The word ‘thinking’ in this context refers to sequential processing.
These systems train recurrent models (networks that recycle parameters between layers) to solve reasoning problems.
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.
When trained properly, thinking systems learn scalable algorithmic processes for solving problems.
After training to solve small/easy problem instances with few recurrent iterations, the system is then extended to run for more iterations at test time.
In doing so, the system can achieve logical extrapolation, and solve problems of greater difﬁculty than those in the training set.
To date, the level of logical extrapolation observed in thinking systems has been quite modest.
For example it has been demonstrated that a system trained on 9 × 9 mazes can extrapolate to solve a 13 × 13 maze.
These systems fail to achieve greater extrapolation because of the overthinking problem; recurrent systems, when extended too far outside their training regime, often deteriorate and fail to produce interpretable outputs.
In this work, we design purpose-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.
By doing so, we are able to build thinking systems that exhibit extreme logical extrapolation behaviors, and leap from solving small/simple training problems with tens of iterations to solving large and complex problem instances at test time using hundreds or even thousands of iterations.
We experiment on the benchmark problems made available by Schwarzschild et al.
(2021a)
for measuring extrapolation behavior.
These tasks include computing preﬁx sums, ﬁnding optimal solutions for two-dimensional mazes, and solving chess puzzles.
We modify model architectures and propose a novel training routine for recurrent networks that signiﬁcantly boosts performance on all tasks in the benchmark suite.
Additionally, we demonstrate that existing models that have shown positive logical extrapolation on these datasets speciﬁcally are susceptible to overthinking, a problem that is largely overcome by our new architectures and training loops (Schwarzschild et al., 2021b).
Our contributions can be summarized as follows.
I. We provide a recurrent architecture for logical extrapolation in which the problem input is concatenated directly to the feature stack of certain layers in the recurrent thinking module.
This prevents the problem instance from being forgotten if deep features become noisy, corrupted, or lossy. II.
We develop a new training routine that incentivizes recurrent networks to make incremental improvements towards a solution, improving the feature representation on each iteration. This training process removes information about how many times the
recurrent module
has been applied , which in turns prevents the network from learning iteration-speciﬁc behaviors (e g., a special behavior for iteration ﬁve and another for iteration six) rather than scalable behaviors than can be iterated indeﬁnitely to yield extrapolation.
III. We analyze the overthinking problem and show that our models overcome this phenomenon.
Our improvements in performance on the easy-to-hard benchmark datasets (Schwarzschild et al., 2021a) can be categorized in several ways.
First, our models yield uniformly higher accuracy across the most difﬁcult tasks used for testing in previous work.
Second, we show that our models can extrapolate to much harder/larger examples than are considered in previous work, where the prior methods generalize poorly, if at all.
Lastly, we show that our models do not forget solutions in settings where previous models overthink.
R ELATED WORK
The work in this paper conceptually builds upon the systems recently studied by Schwarzschild et al.
(2021b), which showed that simple recurrent architectures, when trained to solve various reasoning problems, can then perform logical extrapolation while their feed-forward counterparts cannot.
In this section, we contextualize this approach amongst prior work on algorithm learning, adaptive neural models, and logical extrapolation.
Algorithm learning describes models that learn recurrent processes from data.
Early works on this topic study the ability of recurrent neural networks (RNNs) to process input strings of arbitrary lengths (Gers & Schmidhuber, 2001; Schmidhuber et al., 2007).
More recently, Graves et al.
(2014) introduced neural Turing machines designed to mimic programmable computers, and Kaiser &
Sutskever (2015) propose a parallel version inspired by massively parallel graphical processing units.
These methods, and improvements to them, show promising results 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 (Graves et al., 2014; Kaiser & Sutskever, 2015; Freivalds & Liepins, 2017).
These methods are based on classical RNNs so 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 the difﬁculty 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 difﬁcult to apply them in situations where the solution to a problem cannot be decomposed into incremental parts (e.g., chess).
Constraint satisﬁability problem (CSP) solving networks disentangle the amount of work from the input size (Selsam et al., 2018).
Speciﬁcally, message passing neural networks can execute more passes to solve harder CSPs (Selsam et al., 2018).
These systems are speciﬁc 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 (Schmidhuber, 2012).
Adaptive compute time (ACT) is an algorithm that provides RNNs with a halting unit which estimates the probability that computation should continue.
This algorithm penalizes ‘ponder time’ during training to encourage the network to solve problems quickly (Graves, 2016).
Eyzaguirre & Soto (2020) exhibit strong performance on visual question answering by introducing a differentiable version of
ACT.
Adaptive transformer-based language models also exist, notably Universal Transformers and
Depth-Adaptive Transformers, which utilize ACT to determine the work required for each input (Dehghani et al., 2018; Press et al., 2021).
Similarly, iterative residual networks like NAIS-Nets repeat blocks in stages and perform well on image classiﬁcation (Ciccone et al., 2018).
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 denotes the task of generalizing to test sets which comprise more computationally complex samples than the training data.
Nuamah (2021) claims that “neural network models with end-to-end training pipelines ... cannot, on their own, successfully perform algorithmic reasoning,” and instead proposes a hybrid hand-crafted and learned approach.
Similarly, Palm et al.
(2018) propose recurrent relational networks, which operate on graphs by iteratively passing messages.
They also claim that classical architectures lack the inductive bias to reason about relationships between objects.
Several recent works call these claims into question.
Schwarzschild et al.
(2021b) employ recurrent networks based on weight-sharing architectures, which can be made deeper at test time independent of the input size (Schwarzschild et al., 2021c).
These systems were shown to exhibit logical extrapolation behavior in several domains.
More details on these methods are discussed in
Section 3 as our algorithms build on these directly.
Banino et al.
(2021) reformulate the halting unit in
ACT leading to a probabilistic model with improved performance called PonderNet, and importantly their method outperforms ACT on logical extrapolation tasks for preﬁx sums.
The thinking systems proposed by Schwarzschild et al.
(2021b) depart from classical recurrent networks for (e.g.) text, which learn from step-by-step supervision to produce output tokens one at a time.
In contrast, thinking systems autonomously synthesize a scalable algorithm end-to-end with no supervision over what each algorithmic step should produce.
For this reason, they can potentially be applied to solve complex problems that are difﬁcult or impossible to decompose by hand.
M ETHODS
We begin with some terms and deﬁnitions.
We study networks that share weights across blocks of layers during training.
For example, instead of four distinct residual blocks, a single residual block is repeated four times (see Figure 2 for a graphical depiction).
At test time, networks trained this way can be made “deeper” to extend their compute budget simply by repeating the block more times.
We refer to the number of convolutional layers applied in a recurrent network as its “depth”, and this quantity grows as the number of recurrent iterations increases.
The number of feature maps produced by each convolution is referred to as its “width.”
More formally, let r be a function representing a recurrent block, e.g. a ResNet block (He et al., 2016), and let r n denote n recurrences of that function, e.g. r 2 ( x ) =
r ( r ( x )) .
Let φ denote a feature map, or an output of r , and let φ n = r n ( x ) .
Finally, we 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.
A Deep Thinking (DT) network with m iterations of the recurrent block can then be expressed as follows.
[Equation]
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 a ReLU after the ﬁrst and second.
We ﬁx m during training and compute gradients for optimization by backpropagating through the unrolled network.
Then, m can be increased for testing, allowing these networks to increase their processing power and solve larger and harder problems.
Below, we consider two modiﬁcations to standard approaches for achieving improved extrapolation: modiﬁed architectures and modiﬁed optimization.
3. A DDING RECALL TO THE ARCHITECTURE
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 modify thinking 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 in highway networks, ResNets, and DenseNets, stabilize training (Srivastava et al., 2015;
He et al., 2016; Huang et al., 2017).
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.
Recall allows the system to look back at the problem at any time and reproduce any missing or damaged features.
To formalize this architectural change, adding recall to the network can be expressed using the notation deﬁned above as follows.
[Equation]
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.
We add a single convolutional layer to map the input [ φ k − 1 , x ] to a feature map of the same shape as φ k − 1 .
We refer to DT networks with concatenating skip connections as DT-Recall models.
3. P ROMOTING FORWARD PROGRESS THROUGH OPTIMIZATION
We modify the training objective to encourage the system to incrementally make progress from any starting point.
We do this by inputting a problem instance and running the recurrent module for some random number of iterations.
We then take the output of this process, re-initialize the network with these features as if iteration had just begun, and train the model to produce the correct solution after a random number of iterations.
This incremental training process has two beneﬁts.
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 prevent the network from internally counting iterations and learning iteration-speciﬁc behaviors, such as behaviors that get executed only on iteration ﬁve, for example.
Rather, the network is forced to learn iteration-agnostic 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 to improve this partial solution.
We then update the network’s parameters to minimize loss after n + k total iterations when it starts with the initial partial solution.
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 .
The process above is an analog of truncated backpropagation through time (Jaeger, 2002), with a random starting and end point.
During training, we ensure that the sum of n and k is less than a ﬁxed maximum number of iteration m , which we call the training regime .
The incremental loss described above is added to the standard loss computed with a full forward and backward pass through the unrolled m -iteration network.
The training loop and computation of the loss are given in Algorithm 1.
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 .
Note that 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.
Input: parameter vector θ , integer m , weight α for batch_idx = 1, 2, .. do
Sample inputs x and targets y .
Choose n ∼ U { 0 , m − 1 } and k ∼ U { 1 , m − n } Compute
φ n with forward pass through n iterations without tracking gradients Compute ˆ y prog with
forward
pass of additional k iterations Compute ˆ y m with
E XPERIMENTS
We evaluate our improvements on the benchmark problems available in the Python package easy-to-hard-data .
The three problems considered are computing preﬁx sums, ﬁnding the optimal path in a two-dimensional maze, and solving chess puzzles.
We brieﬂy review the input and output structures for each problem and refer the reader to Schwarzschild et al.
(2021a)
for more detail, including the data generation process.
Note that the architectures we consider are fully convolutional, and produce outputs of the same dimension as their inputs.
Furthermore, the training and testing datasets we conside have labels of the same dimension as their inputs.
Therefore, a 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 preﬁx sums modulo two.
The inputs and targets for this problem consist of bit strings.
The i th bit of the target is the mod 2 sum of all bits prior to and including the i th bit of the input.
We control the “difﬁculty” of the problem by changing the length of the bit string.
Note that computing preﬁx sums of greater length is known to require a greater number of sequential operations (Ladner & Fischer, 1980).
All of our training is done on 32-bit strings and we explore the behavior of our models on longer strings, even showing strong performance on 512-bits.
The mazes 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.
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 × 9 mazes.
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.
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.)
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.
The chess dataset is sorted by difﬁculty rating, as determined by Elo scores computed via human trials on Lichess.org.
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.
In each problem domain, we show that architectures with recall trained with the progressive loss perform best.
We include ablation studies lending weight to our claim that both the architecture and the loss modiﬁcations are instrumental in boosting logical extrapolation power.
For a discussion on hyperparameter choice and training set-ups for all of our experiments, see Appendix A.3.
Even with tuned hyperparameters, training is sometimes unstable – e.g. one out of every ﬁve maze-solving models fails to train.
Code to reproduce all of our experiments is publicly available.
Preﬁx Sums
Though Schwarzschild et al.
(2021b) only perform tests on strings with up to 48 bits, we use 48-bit and 512-bit testing sets to illustrate the extent of our performance improvements.
We train our models on 32-bit strings with a maximum of only m = 30 recurrent iterations (indicated in the ﬁgures as the shaded ‘training regime’ section), 400 convolutional ﬁlters per layer, and incremental loss weight α = 1 when the new loss is used.
[Figure]
In Figure 3, we show that DT networks without our proposed modiﬁcations are unable to solve very long binary strings, achieving 0% accuracy on 512-bit data, while DT-Recall models reach 100% after approximately 200 iterations.
Notably, our models with recall and progressive loss do not suffer from the sharp decline in accuracy as the number of iterations increases beyond 20.
The monotonic increase in accuracy with added iterations in our model is practically useful, as it allows for pre-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.
We compare our models to a baseline of DT models trained with m = 30 and 400 ﬁlters per layer, which is both wider and deeper than the models studied by Schwarzschild et al.
(2021b).
However, to provide a fair comparison, we use the same values for these parameters since we ﬁnd that all architectures beneﬁt from additional width and depth.
We also compare to feed-forwards model without weight sharing of the same effective depth as a 30-iteration DT network.
To better understand the individual effects of our proposed approach, we perform an ablation study.
On both test sets, DT-Recall networks trained without the progressive loss achieve the same ﬁnal accuracy as models trained with it, but require approximately twice the number of iterations to get there, see Figure 4.
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 initerations speciﬁcally.
DT models trained with the progressive loss reach approximately 3% higher peak accuracy than DT networks on 48-bit data, which is also sustained for more iterations.
However, both these models (without recall) fail on 512-bit data and they suffer from overthinking. [Figure]
Mazes As above with the preﬁx sum problem, we can improve performance on hard mazes by combining incremental training and recall architectures.
In particular, while Figure 4 may not convince the reader that our proposed loss is critical since DT-Recall models without it perform very well, with more complex data a drastic difference emerges.
We show in Figure 5 that our models trivially solve the test set of 13 × 13 mazes considered to be ‘hard’ in prior work (Schwarzschild et al., 2021b).
Moreover, on the signiﬁcantly harder test set of 59 × 59 examples, our models exhibit strong logical extrapolation, while previous methods, both feed-forward and thinking systems, completely fail.
Not only do our models achieve a higher peak accuracy, but they do not overthink, as can be seen by the ﬂat spans in the red curves in Figure 5.
In fact [Figure] we can push these systems to their limits (and the limits of our hardware) and ﬁnd that our models can solve nine of ten 201 × 201 mazes using more than 2,000 iterations.
See Appendix A. for an example of an input this large
[Figure]
The
[Figure]
We seek to understand the importance of recall and training with the improved optimization procedure.
Since the best models in Figure 5 are DT-Recall models trained with 30 maximum iterations and a weighting in the loss of α = 0 .
01 , we compare to models with other combinations of design elements.
In order to see how critical recall is in overcoming the overthinking problem, we show in the left-hand pane of Figure 6 that without recall all models suffer from overthinking.
The decay in each dashed curve and the stability with high iteration counts of the solid curves makes this clear.
Notice, however, that testing on 13 × 13 mazes does not reveal large differences between DT-Recall models (with/without progressive loss).
For this, we shift our attention to tests on 59 × 59 mazes, where DT-Recall models with progressive loss achieve on average 97% accuracy while DT-Recall models without progressive loss only reach an average of 83%.
Interestingly, on mazes, we ﬁnd that the best model has α = 0 .
01 in the loss, a much smaller weight than we use for preﬁx sums.
More on ﬁnding the optimal value for α can be found in Appendix A.3.
Chess Puzzles We further ﬁnd that using recall and the progressive loss yields notable improvement on chess puzzles.
In Figure 7, these modiﬁcations generate a 4% accuracy improvement compared to vanilla DT networks.
Moreover, the accuracy of the DT-Recall networks is preserved as the number of iterations increases, while that of the DT networks drops to nearly 0% after 100 iterations.
As with the above benchmark datasets, we justify our hyperparameter choices with ablation studies.
In the right-hand plot of Figure 7, we show that either removing recall or training without the incremental progress loss will damage performance.
Our best models are DT-Recall networks with 512 convolutional ﬁlters in each layer trained with a maximum of 30 iterations and a weight for the combination of loss terms of α = 0 .
Results from tests on harder puzzles are in Appendix A.4.
[Figure]
T HE OVERTHINKING PROBLEM
Deep Thinking networks boast impressive capabilities to solve harder problems by thinking deeper, but they are prone to overthinking.
In particular, we see in Section 4 that some recurrent networks that can perform logical extrapolation may collapse entirely when performing too many iterations.
One observation that can be made from the results above is that the overthinking problem seems to disappear when skip connections are added to provide networks with an uninterrupted view of the input.
Another way to describe this is that our models seem to converge to a ﬁxed-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, we present similar ﬁndings for preﬁx sum models.
5. M ANIPULATING 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 one 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, results, and further discussion.
Next, we ask whether or not the initial feature maps (after the projection layer p ) carry any important information.
We test this by replacing the initial feature maps with zeros after 50 iterations.
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.
See the top-left panel of Figure 20 in Appendix A.8.
5. M ANIPULATING INPUTS
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.
Perhaps there is 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 – ﬁrst subtly, then by swapping the maze with an entirely different example.
[Figure]
To start, we change the input by moving the end of the maze two steps closer to the start of the maze.
We use this new input concatenated with features generated after 50 iterations of solving the original maze.
In this case, we see in the bottom-left pane of Figure 20 (appendix) that models with recall correct to solve the new problem in very few iterations.
A ﬁnal way to 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. Figureshows the effect of swapping feature maps.
5. C ONVERGING TO A FIXED POINT
[Figure]
Recall keeps
Finally, we can 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 it will hold onto this representation (or stay nearby) once it is reached.
Moreover, we seek to qualitatively categorize each model type (architecture/loss pair) as convergent or nonconvergent.
To do we measure the change in the solution at each iteration with (cid:96) 2 -norm of the difference.
Figure 9 shows ∆ φ ( n ) :
= (cid:107) φ n − φ n − 1 (cid:107) (cid:96) 2 .
We see that our models appear to converge, while DT networks without recall explode, providing another view into overthinking.
C ONCLUSION
In this work, we improve the logical extrapolation power of neural networks.
We propose an architecture modiﬁcation and a new loss that lead to a 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 logical 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.
In fact, our models use more than 2,000 iterations, the equivalent of more than 10,000 convolutional layers, to solve the largest mazes we consider.
Learning scalable processes and performing logical extrapolation are 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.
Importantly, our neural networks learn end-to-end how to perform logical extrapolation.
Observing neural models that handle this important type of domain shift prompts future inquiry into online stopping conditions, through which models will determine on their own when to stop thinking and into applications to even more complex real world problems.
E THICS AND REPRODUCIBILITY
The experiments in our work, on academic logical reasoning problems, indicate that our models can extrapolate their reasoning, but real-world domain shift is not limited to logical extrapolation, and our paradigm might fail for numerous unforseen reasons.
Thus, we caution against the assumption that our methods will translate to critical applications, which were not tested in this work.
In order to reproduce our ﬁndings refer to the code and the supplementary material.
Both publicly available they provide reference for the exact implementations we use along with speciﬁc hyperparameters and training set-ups.
The data we use is freely accessible as part of an open-source Python package called easy-to-hard-data , and we use version 1.0.0, released in August 2021.
A A PPENDIX A. A RCHITECTURES
The architectures we use in our experiments are very similar to those described by Schwarzschild et al.
(2021b)
, but we provide the details here for convenience.
In all the models, the ﬁrst layer p is a convolutional layer with 3 × 3 ﬁlters (or ﬁlters of length three in the 1D case) followed by a ReLU non-linearity, which projects the input into feature space.
Each ﬁlter strides by one, and inputs are padded with one unit in each direction.
The number of ﬁlters is speciﬁed per dataset in Table 1 as the “width” of the network.
The internal blocks are standard residual blocks with four convolutional layers (followed by ReLUs) and skip connections every two layers (He et al., 2016).
These blocks share weights in recurrent networks and are simply repeated with distinct parameters in feed-forward models.
Models with recall have additional convolutional layers that map the concatenated inputs and features [ φ, x ] to the shape of the feature maps φ .
The feature maps have the same spacial dimension as the input and w channels, where w denotes the width of the model.
The ﬁnal block h is composed of three convolutional layers with decreasing widths (speciﬁc numbers are in the Table 1), and ReLUs after the ﬁrst two.
The third and ﬁnal convolutional layer in h has two channel outputs used for binary pixel classiﬁcation.
[Table]
. T HE LOSS
Preﬁx 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 .
The target output is y ∈ { 0 , 1 } n .
We compute the loss as follows.
[Equation]
where [ · , · ] indexes the output array.
Therefore, for a batch of B instances, the total loss is
[Equation]
This loss applies to all three problem types we consider.
Note that a maze represented by an m × m input has m 2 pixels and the loss can be averaged over those pixels.
The same applies to chess, where there are always 64 pixels.
A. H YPERPARAMETERS
We describe the training set-up for each experiment and discuss hyperparameter choices.
Due to instability of training for preﬁx sum networks, all models are optimized using Adam (Kingma & Ba, 2014) and gradient clipping.
Maze models are also trained with Adam, but no clipping is used.
Chess models train stably with SGD without gradient clipping.
All training is done with a weight decay coefﬁcient of 0.0002 and training with SGD uses a momentum coefﬁcient of 0.9.
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.
Additionally, when we compute averages, we only include models that trained beyond a threshold training accuracy.
Our threshold choices for 90% preﬁx sums and for chess and 99% for mazes.
Coarse experimentation with learning rates and decay schedules is used to determine optimal choices of these hyperparameters for reproducibility and speed of training.
Models were trained with exponential warm-up for a small number of epochs before the main training routine.
All DT models presented in 4 are trained using early stopping to combat overﬁtting – network parameters are saved for model with highest accuracy on hold-out set during training.
Feed-forward networks are tested with both early stopping and full epoch training, and the best performing class of model is taken as the baseline.
In Table 2, we present the training hyperparameters shared among all models presented in Sectionfor each problem instance.
[Table]
[Figure]
[Table]
We perform width and depth experiments to demonstrate that our best preﬁx sum models are, roughly speaking, on a plateau with respect to performance on 512-bit data in hyperparameter space.
We see from the right panel of Figure 10 that performance increases steadily with depth, but no additional performance is achieved by increasing m past 30.
From the left panel of Figure 10, we also see performance steadily increases with network width.
We choose to use a width of 400 ﬁlters per layer and m = 30 because this adequately solves 512-bit data, and it is memory inefﬁcient to use any more for the task at hand.
[Figure]
These experiments are performed on preﬁx sums, which are relatively memory and computationally inexpensive, to exemplify the general correlation between depth, width, and performance observed on all tasks.
The choice of depth and width of maze and chess models is also made to be only as large as necessary to avoid excessive memory and computational costs.
We determine the optimal weight to be used it the combination of the loss terms using a coarse grid search.
For preﬁx sums we test models with α values in [0 , 0 . 25 , 0 . 5 , 0 .
75 , 1] , for mazes we test values [0 , 0 .
01 , 0 . 1 , 0 . 5 , 1] , and for chess we test values [0 , 0 .
We choose the best weight for each task to present in the main experiments.
For evaluating recurrent models, Schwarzschild et al.
(2021b) propose a maximum conﬁdence rule.
This algorithm takes the output from the iteration at which the outputs have the highest conﬁdence.
This is measured by averaging the per pixel conﬁdence, where conﬁdence is computed as the maximum of the softmax over the output channels.
In order to fairly compare our recurrent models to theirs, we use the same evaluation technique.
For more detail see Appendix A.6.
A. E XTENDED RESULTS
Tables 3, 4, and 5 show the peak accuracy of each of the curves presented in the plots in Section 4.
[Figure] [Figure]
[Figure] [Figure]
r
[Figure]
[Table] [Table]
Figures 13 and 14 show the test performance of chess models on even harder test sets.
It is interesting to note that recall and our loss both still help dramatically, but there is an apparent ceiling on generalizing from the easy puzzles to much harder ones.
We leave further investigation into this phenomenon for future work.
A. I N - DISTRIBUTION RESULTS
Information on how models perform in distribution reveals that each class of model ﬁts the training data and generalizes in the classical sense.
Speciﬁcally, we report the accuracy on a held-out validation set from the same problem difﬁculty as the training data.
For recurrent models, we evaluate the performance using the maximum number of iterations used during training, irrespective of the training objective.
See Tables 6, 7, and 8.
[Table]
[Table] [Table]
[Table] [Table]
A. E XIT RULES
When a DT network processes a single input, it can produce output at each iteration.
When we choose a number of test-time iterations M , we are specifying the maximum number of iterations, but we may choose from any of those to extract a single output.
A naive approach is to take that output at iteration M ; we call this We call this the default exit rule .
A slightly more complicated process, and what Schwarzschild et al.
(2021b) do, is to take the output from iteration m ∗ which solves the following maximization problem.
Let ˆ y ( m ) be the output from iteration m , and let ˆ y ( m ) [ k, j ] be the ( k, j ) entry in the output that corresponds to the conﬁdence that the k th pixel is in class j , ﬁnally let k ∈ { 1 , 2 , ..., K } and let j ∈ { 0 , 1 } .
[Equation]
Choosing the output using this maximization technique is called the maximum conﬁdence exit rule .
In order to stay faithful to the prior work, all plots in the main body of this paper use the maximum conﬁdence exit rule (Schwarzschild et al., 2021b).
In this section, we compute accuracy at iteration n by using the output after n iterations instead of taking the highest conﬁdence output at all iterations up to n .
We call this the default exit rule .
[Figure] [Figure]
To better understand the evolution of conﬁdence as the iteration count increases, we plot this in Figure 18.
Note the conﬁdences are high in the early iterations even before the models are correct.
A. H ARD TO EASY
One natural query about our models is how well they perform on test sets that are easier/smaller than the data used for training.
Figure 19 shows that recurrent models can generally solve easier/smaller problems in fewer iterations.
Even though they solve the easier problems in fewer than the maximum number of iterations used in training, we still see that models without recall suffer from overthinking.
[Figure]
A. M ORE ON THE OVERTHINKING PROBLEM
The experiments we present in Section 5 shed light on how some models overcome the overthinking problem.
Here, we present additional plots and tables to give a more complete picture of our ﬁndings.
Table 9 shows the results from several experiments.
The ﬁrst column indicates the average number of iterations to solve a maze and the best accuracy (over all the test iterations considered).
Then, we examine the behavior when the features are perturbed after the ﬁrst iteration by adding mean zero and variance one Gaussian noise (‘Noise’ column) and we ﬁnd that this completely destroys the models without recall while only slightly slowing down those with recall.
Next, we see the same trend when we replace the features at the ﬁrst iteration with arrays of zeros (‘Zeros’ column).
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.
Figure 20 shows accuracy curves for these experiments.
The bottom right plot in Figure 20 shows the norm of the difference in features when we input random noise instead of actual mazes.
This plot shows that the stable behavior of our models 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.
[Table]
Similar experiments on preﬁx sum computation reveal that recurrent models with recall seem to converge nicely as well.
Note, the top two plots in Figure 21 come from identical experiments as those deﬁned above.
Whereas with mazes we shortened the optimal path, in this setting we ﬂip single bits in the input string.
We explore the response of the model to ﬂipping bits at different indices after 50 iterations (when the models have already solved the initial 48-bit problem).
In the lower left plot in Figure 21, we see that the recovery time as a function of which bit is ﬂipped reveals two things: (i) [Figure]
[Figure]
our models can recover from single bit ﬂips
[Figure]
and (ii) the recovery time decays linearly with the index of the bit ﬂipped.
Higher indices are closer to the end and affect fewer bits in the output than lower ones.
[Figure] [Figure] [Figure]