Thinking Deeper With Recurrent Networks: Logical Extrapolation Without Overthinking

1 INTRODUCTION . Humans solve complex logical reasoning problems through the process of logical extrapolation – they assemble simple logical primitives into complex strategies . For example , a person taught to prove simple lemmas 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 finite depth , the effective depth of recurrent models can be expanded after training simply by iterating the recurrent unit for longer . When trained properly , thinking systems learn scalable algorithmic processes for solving problems . After training to solve small/easy problem instances with few recurrent iterations , the 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 difficulty than those in the training set . To date , the level of logical extrapolation observed in thinking systems has been quite modest . For example it has been demonstrated that a system trained on 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 fixed point when iterated for thousands of iterations . By doing so , we are able to build thinking systems that exhibit extreme logical extrapolation behaviors , and leap from solving small/simple training problems with tens of iterations to solving large and complex problem instances at test time using hundreds or even thousands of iterations . We experiment on the benchmark problems made available by Schwarzschild et al . ( 2021a ) for measuring extrapolation behavior . These tasks include computing prefix sums , finding optimal solutions for two-dimensional mazes , and solving chess puzzles . We modify model architectures and propose a novel training routine for recurrent networks that significantly boosts performance on all tasks in the benchmark suite . Additionally , we demonstrate that existing models that have shown positive logical extrapolation on these datasets specifically are susceptible to overthinking , a problem that is largely overcome by 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-specific behaviors ( e.g. , a special behavior for iteration five and another for iteration six ) rather than scalable behaviors than can be iterated indefinitely to yield extrapolation . III . We analyze the overthinking problem and show that our models overcome this phenomenon . Our improvements in performance on the easy-to-hard benchmark datasets ( Schwarzschild et al. , 2021a ) can be categorized in several ways . First , our models yield uniformly higher accuracy across the most difficult tasks used for testing in previous work . Second , we show that our models can extrapolate to much harder/larger examples than are considered in previous work , where the prior methods generalize poorly , if at all . Lastly , we show that our models do not forget solutions in settings where previous models overthink . 2 RELATED WORK . The work in this paper conceptually builds upon the systems recently studied by Schwarzschild et al . ( 2021b ) , which showed that simple recurrent architectures , when trained to solve various reasoning problems , can then perform logical extrapolation while their feed-forward counterparts can not . 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 difficulty of the problem ) . Moreover , classical RNNs are often trained incrementally to produce one bit at a time , rather than synthesizing an algorithm for solving an entire problem end-to-end . This makes it difficult to apply them in situations where the solution to a problem can not be decomposed into incremental parts ( e.g. , chess ) . Constraint satisfiability problem ( CSP ) solving networks disentangle the amount of work from the input size ( Selsam et al. , 2018 ) . Specifically , message passing neural networks can execute more passes to solve harder CSPs ( Selsam et al. , 2018 ) . These systems are specific to the problem of constraint satisfaction for boolean expressions , but they are an early demonstration of scalable algorithmic behavior . Adaptive neural networks are designed to expend varying amounts of computation on different inputs , thus overcoming the limitation of classical RNNs . Self delimiting neural networks use one neuron to determine when to stop updating the hidden state in RNNs , and in doing so , they perform more or less computation for each token in an input sequence ( 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 classification ( 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 ... can not , 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 prefix sums . The thinking systems proposed by Schwarzschild et al . ( 2021b ) depart from classical recurrent networks for ( e.g . ) text , which learn from step-by-step supervision to produce output tokens one at a time . In contrast , thinking systems autonomously synthesize 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 difficult or impossible to decompose by hand . 3 METHODS . We begin with some terms and definitions . We study networks that share weights across blocks of layers during training . For example , instead of four distinct residual blocks , a single residual block is repeated four times ( see Figure 2 for 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 rn denote n recurrences of that function , e.g . r2 ( x ) = r ( r ( x ) ) . Let φ denote a feature map , or an output of r , and let φn = rn ( x ) . Finally , we 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 . A Deep Thinking ( DT ) network with m iterations of the recurrent block can then be expressed as follows . fm ( x ) : = h ( r m ( p ( x ) ) ) ( 1 ) 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 first and second . We fix m during training and compute gradients for optimization by backpropagating through the unrolled network . Then , m can be increased for testing , allowing these networks to increase their processing power and solve larger and harder problems . Below , we consider two modifications to standard approaches for achieving improved extrapolation : modified architectures and modified optimization .

Summary: This paper is focused on training networks to solve problems via extrapolation, e.g. solving large mazes by learning to solve small mazes. In order to achieve this, two modifications to recurrent convolutional networks are proposed: a) adding a concatenating skip connection from the input to the recurrent layers that stabilizes extrapolation (Fig. 2), and b) using a modified training method and loss that encourages the network to learn computations independent of current iteration (Alg. 1). Together, these techniques lead to substantial improvements in results on the problems of computing prefix sums, solving 2D mazes and solving chess puzzles.