Towards Modular Algorithm Induction

1 INTRODUCTION . Many applications in artificial intelligence require the ability to learn and perform tasks that have algorithmic structure , such as learning a sequence of precise actions together with conditional decisions and branching . Therefore , an important research problem is learning to induce algorithms from input-output examples . A particular focus of research attention has been on neural algorithm induction , which are neural architectures for representing algorithms by including unbounded intermediate state and mechanisms for learning control flow . The neural Turing machine ( Graves et al. , 2014 ) and its successor the Differentiable Neural Computer ( Graves et al. , 2016 ) augment a recurrent network with an external memory using differentiable read/write mechanisms , and are able to learn simple algorithmic tasks such as array copying and sorting . Neural Random-Access Machines ( NRAM ) ( Kurach et al. , 2016 ) learn algorithms by generating a fuzzy circuit comprising of pre-defined modules together with modules for dereferencing and storing values in a fixed memory . Existing architectures do not fully leverage a key fact that algorithmic tasks can be solved by flexibly combining the results of smaller , reusable procedures , which we call modules . A modular architecture aims to represent in a learning system these aspects of programming languages intended for human developers : ( i ) procedural abstractions that perform computations to produce outputs given some inputs and that can be reused in the algorithm multiple times ; ( ii ) control flow constructs such as branching and loops to compose and combine intermediate results . Although there has been a recent movement toward modular neural architectures for supervised tasks ( Andreas et al. , 2016 ; Kirsch et al. , 2018 ; Rosenbaum et al. , 2018 ; Valkov et al. , 2018 ; Chang et al. , 2018 ) , the design space of modular networks has not been fully explored for algorithm induction . Towards this end , we present a new architecture MAIN ( short for Modular Algorithm Induction Network ) for neural algorithm induction . Our architecture consists of a neural controller that interacts with a variable-length read/write tape where inputs , outputs , and intermediate values are stored . Each module is a small computational procedure that reads from and writes to a small fixed number of tape cells ( a given fixed set of modules are specified in advance ) . At each time step , the controller selects a module to use together with the tape location of the module ’ s input arguments and the write location of the module output . This architecture is trained end-to-end using reinforcement learning . A comparison of architectural design choices in MAIN with those in recent neural program induction approaches NTM ( Graves et al. , 2014 ) , LSA ( Zaremba et al. , 2016 ) , NRAM ( Kurach et al. , 2016 ) , NGPU ( Kaiser & Sutskever , 2016 ) , CRL ( Chang et al. , 2018 ) , and DNC ( Graves et al. , 2016 ) is presented in Table 1 . Unlike previous architectures , MAIN allows for learning to compose modules together with the corresponding argument values for the chosen modules using a general domainagnostic mechanism for module and argument choices . It uses a generic linear tape layout together with a parallel history tape of landmark symbols that indicate the most recently read and written cells . Finally , the architecture allows for random access of the memory cells and uses a memoryless controller with a length-invariant self-attention based encoding of input tape contents . There are two key design choices in MAIN that we found crucial for good performance . The first is in the representation of the previous history of the computation . Like previous work ( Graves et al. , 2016 ) , we find that providing history to the model is an important source of information . We introduce a simple but effective discrete representation , introducing a parallel history tape of landmark symbols that indicate the most recently read and written cells . This helps the model to learn common patterns of control flow , such as maps and reduce operations over the tape . The second choice is in the encoder architecture of the abstracted tape view with the controller . We find that using an attention-based encoder performs much better than a recurrent network based encoding within the controller , which has been employed by previous architectures for neural algorithm induction . We evaluate our architecture MAIN on five algorithmic tasks including array copying , reversing , increment , filter-even , and multi-digit add . MAIN can learn these tasks in an end-to-end manner using input-output examples and is also able to generalize to longer input lengths . We observe that both parallel tape history and length-invariant input tape encoding based on self-attention are necessary for the architecture to learn these tasks . Moreover , for tasks such as filter-even and multidigit add that require the controller to perform some computations , we observe that abstracting the tape contents to only certain landmark positions can help it learn the corresponding algorithms . 2 RELATED WORK . Learning to Compose and Modular Networks : Compositional Recursive Learner ( CRL ) ( Chang et al. , 2018 ) is a framework for learning algorithmic procedures for composing representation transformations , where both the transformations and their compositions are simultaneously learnt from a sparse supervision . The controller that learns to compose transformation is trained using reinforcement learning , whereas the transformations themselves are trained using supervised learning . The CRL architecture supports two forms of transformations : reducers and translators . Because of specialized transformations , CRL imposes a restriction on always selecting either full tape or three consecutive input tokens as arguments ( that result in 1 resulting token ) for the transformations . In contrast , our architecture allows the controller to select arbitrary locations on the input tape to select as module arguments and also the write locations for their outputs . There has also been some recent related work on learning Modular Networks ( Kirsch et al. , 2018 ) and Routing Networks ( Rosenbaum et al. , 2018 ) . In routing networks , a router learns to select a sequence of function blocks to compose given an input , and is trained using reinforcement learning . Kirsch et al . ( 2018 ) use a probabilistic model to represent the module choice as a latent variable and use Expectation Maximization ( EM ) to learn the module parameters and the decomposition choice in an end-to-end manner by maximizing a variational lower bound of the likelihood . Both of these approaches learn modular composition of functions with the aim of reusability and better generalization across multiple tasks . Our architecture , in contrast , is designed for learning algorithms where the modules are pre-specified but the task of the controller agent is to learn to compose those modules to achieve desired input-output behavior . Neural Program Induction and Synthesis : Several neural architectures have been proposed recently to learn algorithmic tasks . Neural Turing Machine ( NTM ) ( Graves et al. , 2014 ) extends an LSTM controller with an external memory together with differentiable read-write mechanism . Differentiable Neural Computer ( DNC ) ( Graves et al. , 2016 ) , a successor to NTM , added capabilities of freeing up unused memory for processing long sequences and a temporal link matrix for better tracking the write order . Zaremba et al . ( 2016 ) propose an architecture with an RNN controller that learns to navigate a structured input grid by selecting move actions and outputs tokens on an output grid . Unlike our architecture , which separates control from compute , these architectures require the controller to learn both the control structure as well as the desired computations . Neural RAM ( NRAM ) ( Kurach et al. , 2016 ) uses a neural controller that learns to produce a fuzzy circuit consisting of a fixed number of modules . The controller learns to wire the modules with inputs coming from a fixed set of registers or intermediate outputs , and also learns to write output back to registers and memory tape . There are three key differences between Neural RAM and our architecture . First , NRAM uses a fixed sequence of 14 modules ( with an option to repeat the whole sequence multiple times ) . In contrast , our controller learns to select appropriate module at each time step . Second , since registers and memory tape in NRAM architecture store distributions over integers , the modules also need to be differentiable to compute output over such inputs . Our architecture , on the other hand , does not require the modules to be differentiable and also writes discrete values on the input tape . Finally , our architecture uses a parallel history tape of recent read and write positions for accessing the input tape unlike pointer dereferencing in NRAM . We summarize the characteristics of these architectures , comparing them to our work in Table 1 . To understand the rows of this table , “ Modules ” means whether computation is performed by composable modules . “ General module selection ” means the modules can be selected in any order at any point in the computation ( NRAM enforces a fixed ordering ) . “ General argument selection ” means module arguments can come from anywhere in memory without restriction ( CRL enforces inputs be locally adjacent ) . “ Read/write history ” means that history of previous read and write locations is provided as context to the controller . “ Random access ” means the architecture can look up arbitrary memory contents and agregate information across memory regions , rather than be restricted by relative head movements ( LSA ) . “ Memoryless controller ” means the architecture does not require the controller to have hidden memory over time , i.e . the controller is not recurrent , and only uses external memory for context . Finally , the “ Encoder type ” indicates how memory is consumed by the controller , feed-forward ( FF ) , convolutional ( CNN ) , recurrent ( RNN ) , or attention-based ( Attn ) . There is also a growing interest in using neural models to generate symbolic programs as output . RobustFill ( Devlin et al. , 2017 ) trains an attention based encoder-decoder model that generates a program given a set of input-output examples . Deepcoder ( Balog et al. , 2017 ) learns a probability distribution over a set of functions in a DSL to guide an enumerative search . Chen et al . ( 2019 ) also use an encoder-decoder based approach to synthesize Karel programs ( Bunel et al. , 2018 ) , but in addition also execute the partially decoded program to guide the decoder . In contrast , our architecture uses a more general mechanism to read and write on the variable length input tape , and uses reinforcement learning to learn to compose the desired modules for each individual task . 3 ARCHITECTURE We describe MAIN , an architecture which can express arbitrary modular Turing machines . Our architecture has three parts : the memory which stores the state of the computation , the modules which update the tape , and the controller which chooses a module Mi to execute at each stage of the computation as illustrated in Figures 3 . Once Mi is selected out of the module set ( 4 modules depicted ) , and the read/write locations on the tape are chosen , the module sees only the read inputs , and modifies the tape only at the write position . Each module is a user provided function ( which could in principle also be learned ) . In our experiments we picked modules of two arguments and one output , but our architecture is agnostic to the number of read/write heads as shown in Figure 1 . The memory is a finite tape of discrete tokens/symbols . Formally , let s0 , s1 , . . . , sT be memory states at each computation step t = 0 , . . . , T . Each st is itself an array of tokens , indexed as st [ 0 ] , . . . , st [ L − 1 ] . At the start of the computation , s0 is initialized with the program input , and necessary empty space to perform intermediate computation and write the output . The memory length L is dynamically set based on the size of the given input , so every neural architecture component that depends on L needs to be length invariant . The memory contains landmark tokens , which are are task-specific tokens , e.g . ‘ $ ’ , ‘ + ’ , ‘ & ’ , that provide positional information . For example , these tokens indicate where the input starts , where the output should be written , etc . In our experiments we found that the controller would often overwrite the landmark tokens during training . As a workaround , we provide the positions of landmark tokens as immutable metadata . These are the lambdas λ ( 1 ) , λ ( 2 ) , . . . inside σt shown above . Additionally , we place Start-Of-Tape and End-Of-Tape landmarks at the first and last tape positions . The initial tape s0 contains the input and additional space for scratch work . At the end , the output of computation will be read from designated regions of the tape , called target locations , which are initialized with the empty token ‘ . ’ . Note that the architecture is free to overwrite any position during the course of computation , including important input tokens . When scoring the final tape sT , we only look at designated target positions , designated by the ‘ . ’ token . The initial tape configurations for Copy and Multi-digit Add Tasks is shown in Figure 2 . The tape configuration in addition to the input tape also consists of landmark positions as well as the read/write heads history . The modules M = { m1 , . . . , mk } are functions that read a narrow portion of the memory , and compute new values that are stored in memory . Each module mi is a function of R arguments , and outputs a vector of size W , i.e . number of read/write heads corresponds to the number of inputs and outputs of the modules . During the computation , each input and output of the modules will be a single cell on the memory tape . Some example modules are the maximum module which returns the maximum of two inputs , and the sum module which returns the sum of two inputs mod base B . The controller is a policy over actions that specifies a module-tape interaction . The controller can not directly modify memory . Instead it selects a module and the locations on the tape the module will read from and write to . Specifically , the controller defines a distribution π ( H ( r,1 ) , . . . , H ( r , R ) , H ( w,1 ) , . . . , H ( w , W ) , M | ct ) , where H ( ... ) are random variables over tape positions ( support is { 0 , . . . , L − 1 } ) ; one for each head . The variables H ( r,1 ) , . . . , H ( r , R ) select read-head locations , and H ( w,1 ) , . . . , H ( w , W ) select write-head locations . The random variable M is one of the modules . Finally ct is the context , which contains the current state of memory and information about the computation history . We will define this in more detail shortly . The controller defines an end-to-end computation as follows . At each step t of the computation , let the current state of the memory by st. To choose the next module and the locations of the read and write heads , we sample from the controller . This produces a module choice mt , and locations h ( r,1 ) t , . . . , h ( r , R ) t for the read heads , and h ( w,1 ) t , . . . , h ( w , W ) t for the write heads . Because the controller is free to choose any tape location for each read and write head independently , the memory is random access . The tape is updated by calling the module mt , with the tape contents under the read heads as input , and then writing to the position specified by the write heads . More formally , st+1 [ h ( w , j ) t ] : = m ( st [ h ( r,1 ) t ] , . . . , st [ h ( r , R ) t ] ) [ j ] ∀ j ∈ { 1 , . . . , W } st+1 [ i ] : = st [ i ] ∀ i /∈ { h ( w,1 ) t , . . . , h ( w , W ) t } . The selected write indices are updated with module ’ s output , and other tape elements are unchanged . The context ct is the input to the controller , which represents both the current tape contents st and an action history that represents information about the previous computation . The action history consists of two parts : • Fixed-size part : Contains the module choice and head locations chosen at the previous time step . – Module choice : mt−1 – Tape values underneath heads : st [ h ( r,1 ) t−1 ] , . . . , st [ h ( r , R ) t−1 ] , st [ h ( w,1 ) t−1 ] , . . . , st [ h ( w , W ) t−1 ] . • Variable-size part : One-hot encoding of { h ( ... ) t−1 } , the read and write head locations at the previous step of the computation . The tape values underneath the heads are technically redundant information , because the controller can use st combined with the head locations from the variable-sized part to lookup the corresponding tape values . However , we found that the controller had a hard time doing this in our experiments , and providing the head values as auxiliary input to the controller proved helpful . Formally , the context is a tuple ct = ( ξt , σt ) , where ξt is a fixed sized encoding ( does not depend on the tape length L ) , and σt is a variable sized encoding ( depends on the tape length L ) . First , ξt describes the current tape values underneath the tape heads chosen at the previous time step : ξt = 〈st [ h ( r,1 ) t−1 ] , . . . , st [ h ( r , R ) t−1 ] , st [ h ( w,1 ) t−1 ] , . . . , st [ h ( w , W ) t−1 ] , mt−1〉 . where mt−1 is the previous module choice . Second , the variable sized encoding σ is a matrix concatenation that provides complete information about st : σt = 1V { st [ 0 ] } 1V { st [ 1 ] } . . . 1V { st [ L− 1 ] } 1L { λ ( 1 ) } 1L { λ ( 2 ) } ... 1L { h ( r,1 ) t−1 } ... 1L { h ( w,1 ) t−1 } ... . We can think of σt as a stack of binary channels , each a row-vector of lengthL . The one-hot function 1D { d } produces a vector of length D which is filled with 0s , and 1 at position d. Each 1V { st [ ` ] } is a one-hot vector of token st [ ` ] . Their horizontal concatenation produces a channel for each token ( indicates its presence or absense ) . The 1L { λ ( i ) } and 1L { h ( ... ) t−1 } channels indicate the positions of landmarks and heads respectively ( λ ( i ) are fixed throughout the episode and not indexed by t ) . Note that at t = 0 no actions have been taken , so the previous-action encodings are all 0s . Now we describe the controller architecture . Given the context ct , the controller begins with a sequence encoder for the variable length encoding σt . We pass σt through two 1D convolutional layers ( along the L dimension ) with filter width 3 and stride 1 . This is the only aspect of our encoder that provides awareness of the local ordering of cells on the tape . After that , we explore two different seq2fixed encoders : • RNN encoder : BiLSTM produces fixed length embedding ( concat of embedding for each direction ) which is fed into the controller . • Attention encoder : A context-independent query set is learned ( fixed during evaluation ) . Queries are fed into attention over the tape , and the resulting weighted sum over tape values is a fixed length embedding that is fed into the controller . Next , we pass the resulting fixed-sized embedding for σt along with ξt into a feed forward network which outputs attention queries for the read/write head actions , and logits for the module selection action . Logits for the read/write head actions are produced with dot-product attention over σt ( independent attention head for each read/write head ) . Learning the Controller . We frame our setup in the reinforcement learning paradigm , where the controller is the agent , and other components ( memory and modules ) are part of the environment . From the perspective of algorithm induction , only the program input and target outputs are external , while everything in MAIN is part of the black box that performs the computation . The controller is trained end-to-end on all its actions with Impala ( Espeholt et al. , 2018 ) , a distributed variant of REINFORCE . Simultaneously learning to halt with RL , while learning what computation to perform , proved to be unstable in our experiments . We removed the additional complication of learning when to halt by providing a halting oracle . The oracle is given the correct output , and immediately ends the episode if the answer is in memory . The oracle is used in evaluation as well .

Summary: This paper proposes Modular Algorithm Induction Network (MAIN) that learns algorithms given input-output examples. MAIN is equipped with several components that make it perform better than baselines, but probably the most important part is its use of modules to break-down algorithmic tasks into simpler problems. MAIN is learned end-to-end using reinforcement learning and demonstrated to perform well in several tasks.