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 proposes the idea of learning an iterative refinement function within the deep learning framework using recurrence where training is done over a fixed number of iterations, but testing can be performed over an arbitrary number of iterations to let the learned model find a solution in case it needs to “think longer”. The authors have experimentally shown that the learned models can indeed solve more out-of-distribution problems if they're let run longer for three Prefix-sum, Maze and Chess Puzzles problems. They have also shown for the same out-of-distribution problems, a feed-forward model of the same effective depth performs poorly compared to their proposed recurrent model. Based on these results, the authors have claimed that the recurrent model has superior out-of-distribution generalization compared to its feed-forward counterpart. These out-of-distribution datapoints are specifically characterized as “larger” and “harder” throughout the paper, which in turn has paved the way to make the main claim of paper: RNNs can be used to train on easy problems and then generalize over the hard ones at the test time by letting them run longer.