Can You Learn an Algorithm?  Generalizing from Easy to Hard Problems with Recurrent Networks

1 Introduction . In computational theories of mind , an analytical problem is tackled by embedding it in “ working memory ” and then iteratively applying transformations to the representation until the problem is solved [ Baddeley , 2012 , Baddeley and Hitch , 1974 ] . Iterative processes underlie the human ability to solve sequential reasoning problems , such as complex question answering , proof writing , and even object classification [ Liao and Poggio , 2016 , Kar et al. , 2019 ] . They also enable humans to extrapolate their knowledge to solve problems of potentially unbounded complexity , including harder problems than they have seen before , by thinking for longer . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . This work examines whether recurrent neural networks trained on easy problems can extrapolate their knowledge to solve hard problems . We find that recurrent networks can indeed generalize to harder problems simply by increasing their test time iteration budget ( i.e. , thinking for longer than they did at train time ) . Moreover , we find that the performance of recurrent models improves as they recur for more iterations , even without adding parameters or re-training in the new , more challenging problem domain . This ability is specific to recurrent networks , as standard feed-forward networks rely on layer-specific behaviors that can not be repeated to extend their reasoning power . The behavior we observe in recurrent networks falls outside the classical notions of generalization in which models are trained and tested on the same distribution . Because we train and test on problems of different sizes/difficulties , our training and test distributions are disjoint , and systems must extrapolate to solve problems from the test distribution . Outside the field of machine learning , computers achieve a functionally similar extrapolation ability through the use of algorithms , which encode the process required to solve a class of problems , and can therefore scale to problems of arbitrary size , albeit with longer runtime . By training networks to solve problems iteratively , we hope to find models that encode a scalable method for solving problems rather than memorizing a mapping between input features and outputs . In short , the goal is to create recurrent architectures that are capable of learning an algorithm . Our focus is on three reasoning problems that are classically solved using hand-crafted algorithms : computing prefix sums , solving mazes , and playing chess . Sequential reasoning tasks like these are ideal for our study because one can directly quantify the difficulty of a problem instance . In the case of mazes , for example , we can easily swap to a more challenging domain by increasing the size of the search space . For each class of problems , recurrent networks are trained on a set of “ easy ” problems using a fixed number of iterations of the recurrent module on the forward pass . After training is complete , we assess whether our models exhibit logical extrapolation behaviors by testing them on “ hard ” problems , with varying numbers of additional iterations . Remarkably , models trained on easy examples exhibit little extrapolative behavior until their iteration budget is increased — generalizing to harder problems requires thinking deeper . Moreover , we find recurrent models tested with a sufficient number of extra iterations outperform the inflexible feed-forward models of comparable depth , often by a wide margin . Finally , we visualize the iterative behavior of the recurrent module to gain insights into the problem solving process they discover . 1.1 Related works . Our investigation into generalizing from easy to hard examples builds on several bodies of work . Logical extrapolation encompasses a special kind of distributional shift . A number of existing works on domain generalization instead explore shifts , such as re-stylization and image corruptions , which do not represent an increase in scale or computational complexity [ Arjovsky et al. , 2019 , Shu et al. , 2020 ] . Also , the basic neural architectures we use are not new and build upon prior studies of weight sharing and recurrence [ Pinheiro and Collobert , 2014 , Liang and Hu , 2015 , Alom et al. , 2018 , Bai et al. , 2018 , 2019 , Lan et al. , 2020 , Jaegle et al. , 2021 ] . Networks with variable numbers of test time iterations/layers have also been studied , including variable depth networks [ Graves , 2016 , Huang et al. , 2016 , Kaya et al. , 2019 , Eyzaguirre and Soto , 2020 ] . Existing work on algorithm learning involves recurrent neural network ( RNN ) based approaches . For example , neural turing machines and neural GPUs can learn simple algorithms for tasks such as binary addition and multiplication [ Graves et al. , 2014 , Kaiser and Sutskever , 2015 ] . Like most RNNs , the compute budget for these methods is inextricably tied to input length . Motivated by the fact that input sequence length is not necessarily correlated with the computational burden required to solve a problem , Graves [ 2016 ] develops a method for RNNs to adaptively select a compute time limit . This work considers only sequence inputs and shows the benefits of decoupling compute budget from input length . A differentiable extension of the technique can also be applied to visual question answering [ Eyzaguirre and Soto , 2020 ] . The above works leverage neural networks with adaptive computation budgets to speed up and strengthen inference when learning on stationary distributions . In contrast , our work studies the logical extrapolation behaviors that recurrent networks possess when both computation budgets and problem difficulties are extended beyond the train-time regime . In the domain of constraint satisfiability problems ( CSPs ) , Selsam et al . [ 2018 ] show message passing neural networks trained on small CSPs can generalize to larger problems if more messages are passed at test-time – a similar type of extrapolation to our methods , but for a very specific problem formulation . The particular problems we use to study this type of extrapolation include prefix sum computation and maze solving , two problems analyzed in the classical algorithms literature . For example , there are many ways to solve mazes , both classical ( e.g . breadth first search ) and learned ( e.g . value iteration networks ) , but our goal is not to develop the best solver , rather to use mazes as a test bed for logical extrapolation [ Tamar et al. , 2016 ] . Our third case study is the game of chess , which has also been the focus of much artificial intelligence work [ Romstad et al. , Biswas and Regan , 2015 , Silver et al. , 2017 , McIlroy-Young et al. , 2020 ] . However , those efforts to play chess rely heavily on hand-crafted search algorithms , often paired with neural networks or opening books , and aim to play games from start to finish , both methods and goals that diverge from ours . As opposed to hard-coded algorithms , which scale by design , we are interested in studying whether learned processes can generalize from the data on which they are trained to even harder problems . 2 Dataset descriptions . We conduct experiments in three problem domains that are classically solved using hand-crafted algorithms . For each , we define datasets that have quantifiable notions of difficulty . This makes it possible to train models on easy/small examples and test them on harder/larger ones . We consider the task of computing prefix sums modulo two of binary bit strings , solving two-dimensional mazes , and finding the best move in chess puzzles . Prefix sums This problem is inspired by a similar dataset used by Graves [ 2016 ] . Each training sample is a binary string . The goal is to output a binary string of equal length , where each bit represents the cumulative sum of input bits modulo two . Our models accept input strings of any size , and we consider longer strings to be more difficult to process than shorter ones . Each dataset contains 10,000 uniform random binary strings without duplicates . We use datasets with input lengths of 32 , 44 , and 48 bits . See Figure 1 for an example input and the corresponding target output . Input : [ 1 , 0 , 0 , 1 , 0 , 1 , 0 , 1 , 0 , 1 , 0 , 1 , 1 , 0 , 1 , 1 , 0 , 1 ] Target : [ 1 , 1 , 1 , 0 , 0 , 1 , 1 , 0 , 0 , 1 , 1 , 0 , 1 , 1 , 0 , 1 , 1 , 0 ] Figure 1 : Prefix sum input and target example . Mazes We consider mazes generated using a depth-first search algorithm [ Hill , 2017 ] . We train on 50,000 small ( 9 ⇥ 9 ) mazes , and we test on 10,000 larger ( 13 ⇥ 13 ) mazes . Our models are convolutional , and receive a maze as a N ⇥N three-channel image , where the maze walls are black , and the start and goal locations are red and green , respectively . The label for each maze is a binary two-dimensional mask containing the locations of positions along the shortest path solution . Our models output a two logits per pixel which are then thresholded to a binary mask of the whole input , and we consider a candidate solution correct only if it exactly predicts the labeled path . See Figure 10 for an example . More examples are available in Appendix A.1.2 . Chess puzzles The third dataset we use is a corpus of chess puzzles . The data is furnished by Lichess , an online open-source chess server [ Lichess , 2021 ] . From billions of games , Lichess compiles “ puzzles ” – mid-game boards for which a sequence of unique best moves is determinable ( e.g . sequences of moves leading to a forced checkmate ) . From this database , we compiled labeled data where the inputs are 8⇥ 8⇥ 12 arrays indicating the position of each piece on the board ( one channel per piece type and color ) and the outputs are 8 ⇥ 8 binary masks showing the origin and destination positions for the optimal move . See the example in Figure 3 , where it is white to move and the solution to the puzzle is to move the queen from F3 to F7 . The puzzles each have a difficulty rating determined by an Elo-like system ( a standard system for rating players using tournament play ) [ Elo , 1978 ] . When Lichess users , each with an Elo rating , attempt to solve the puzzle , the rating of the player and puzzle is updated as if they “ played ” against each other . After enough players encounter the puzzle , the rating of the puzzle reaches an equilibrium which gets recorded in the database . We use these ratings to distinguish between easy and hard ; 600,000 puzzles of difficulty rating less than 1,385 are used for training , and testing is performed on 100,000 examples with ratings greater than 1,385 . Our models output a confidence score between zero and one for each square on the board . We take the two highest scores and compare their locations to the target to measure correctness – only exact matches are considered correct . 3 Model architectures & training . In all of the experiments in this work , we employ network architectures based on ResNets [ He et al. , 2016 ] . Our feed-forward networks are slight deviations from the most commonly used ResNets , in that the width does not change except at the first layer and after the last residual block , and we do not use batch normalization . This is done so that the recurrent models , whose internal recurrent module has the same input and output dimensions , can be as similar as possible to the feed-forward ResNets . In fact , the only difference during training between a feed-forward and recurrent model of the same effective depth is that the weights are shared between the residual blocks in the recurrent models . We refer to the recurrent portion of the network as the recurrent module . Also , our models are fully convolutional with no fully connected heads . For solving prefix sums , which involves onedimensional strings , we further deviate from classical ResNets by using one-dimensional convolutions . For complete architectural details , see Appendix A.2 . We measure recurrent models in terms of iterations and effective depth . An iteration is a repetition of the recurrent residual block , which contains four layers in all of our models . Therefore , the effective depth is equal to four times the number of iterations , plus non-recurrent encoder and head layers that sandwich the recurrent module . For example , the models used for computing prefix sums have one convolutional encoder layer , followed by the recurrent block and then by a three layer convolutional head . In this case , a 10-iteration model has effective depth 1 + 10⇥ 4 + 3 = 44 layers . For each training sample we consider , the label is an array of binary classification variables , one per input pixel . The training loss is simply the mean cross-entropy loss across output values . In general , hyperparameters were determined with the goal of finding convergent models . The specific batch sizes , learning rate and decay schedules , and other hyperparameter values are all available in Appendix A.3 . 1

Summary: This paper describes a method that allows models to extrapolate to larger problems than those they are trained on, a known limitation of many neural network architectures. It proposes a network with an encoder, a decoder and a recurrent unit (a 4 layer residual convolution network) that is iterated through several times during training, the output of an iteration being used as the input to the next. This amounts to a deep ResNet with shared middle layers.   The authors show that models trained with this architecture can generalize to larger problems by increasing the number of iterations of the recurrent unit during inference. This achieves better accuracy than feed-forward ResNets of equivalent depth.  Three problems are considered: parity integration of a binary string, path finding in random mazes (encoded as binary images), and solving chess problems (generated from Lichess). The models are shown to generalize to longer strings, larger mazes, and, to a lesser extent, harder problems (measured by their Elo ranking).