Neurocoder: Learning General-Purpose Computation Using Stored Neural Programs

Artificial Neural Networks are functionally equivalent to special-purpose comput-1 ers . Their inter-neuronal connection weights represent the learnt Neural Program2 that instructs the networks on how to compute the data . However , without stor-3 ing Neural Programs , they are restricted to only one , overwriting learnt programs4 when trained on new data . Here we design Neurocoder , a new class of general-5 purpose neural networks in which the neural network “ codes ” itself in a data-6 responsive way by composing relevant programs from a set of shareable , modular7 programs stored in external memory . For the first time , a Neural Program is ef-8 ficiently treated as a datum in memory . Integrating Neurocoder into current neu-9 ral architectures , we demonstrate new capacity to learn modular programs , reuse10 simple programs to build complex ones , handle pattern shifts and remember old11 programs as new ones are learnt , and show substantial performance improvement12 in solving object recognition , playing video games and continual learning tasks.13 1 Introduction14 From its inception in 1943 until recently , the fundamental architectures of Artificial Neural Net-15 works remained largely unchanged - a program is executed by passing data through a network of16 artificial neurons whose inter-neuronal connection weights are learnt through training with data.17 These inter-neuronal connection weights , or Neural Programs , correspond to a program in modern18 computers [ 32 ] . Memory Augmented Neural Networks ( MANN ) are an innovative solution allow-19 ing networks to access external memory for manipulating data [ 11 , 12 ] . But they were still unable20 to store Neural Programs in such external memory , and this severely limits machine learning . Stor-21 ing inter-neuronal connection weights only in their network does not permit modular separation of22 Neural programs and is analogous to a computer with one fixed program . Recent works introduce23 conditional computation via adjusting or activating parts of a network in an input-dependent manner24 [ 39 , 33 , 4 , 13 , 28 ] , but networks remain monolithic . Current networks forget when retrained , old25 inter-neuronal connection weights are merged with new ones or erased.26 The brain is modular , not a monolithic system [ 8 , 6 ] . Neuroscience research indicates that the brain is27 divided into functional modules [ 19 , 7 , 9 ] . If the neural program for each module is kept in separate28 networks , networks proliferate . Modular neural networks , another form of conditional computation,29 combine the output of multiple expert networks , but as the experts grow , the networks grow drasti-30 cally [ 20 , 14 , 35 , 29 ] . This requires huge computational storage and introduces redundancy as these31 experts do not share common basic programs.32 A pathway out of this bind is to keep such basic programs in memory and combine them as required.33 This brings neural networks towards modern general-purpose computers that use the stored-program34 principle [ 37 , 40 ] to efficiently access reusable programs in external memory . Here we show how35 Neurocoder , a new neural framework , introduces a new class of general-purpose conditional compu-36 Submitted to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Do not distribute . tation machines in which a neural network can be “ coded ” in an input-dependent manner . Efficient37 decomposition of Neural Programs creates shareable modular components that can reconstruct the38 whole program space . These components change their “ shapes ” based on training and are stored39 in an external Program Memory . Then , in a data-responsive way , a Program Controller retrieves40 relevant components to build the Neural Program . This is analogous to shape-shifting Lego bricks41 that can be reused to build unlimited shapes and structures ( See Appendix Fig . 4 ) .42 Using adaptive modular components vastly increases the learning capacity of the neural network43 by allowing re-utilisation of parameters , effectively curbing network growth as programs increase.44 More importantly , unlike pre-defined sub-networks or modules [ 20 , 1 ] that combine at activation45 level , the construction of our modular components is dynamic and performed on the weight space.46 The Neural Program construction is learnt through training via traditional backpropagation [ 30 ] as47 the architecture is end-to-end differentiable.48 2 Methods49 2.1 System overview50 A Neurocoder is a neural network ( Main Network ) coupled to an external Program Memory through51 a Program Controller . The working program of the Main Network processes the input data to pro-52 duce the output . This working program is “ coded ” by the Program Controller by creating an input-53 dependent active program from the Program Memory ( Fig . 1 ) . The following gives a high-level54 description of the Neurocoder framework and then the details.55 Neurocoder stores Singular Value Decomposition of Neural Programs in Program Memory56 The Neural Program needs to be stored efficiently in Program Memory . This is challenging as there57 may be millions of inter-neuronal connection weights , thus storing them directly ( [ 22 ] ) is grossly58 inefficient . Instead , the Neurocoder forms the basis of a subspace spanned by Neural Programs and59 stores the singular values and vectors of this subspace in memory slots of the Program Memory60 ( hereafter referred to as singular programs ) . Based on the input , relevant singular programs are61 retrieved , a new program is reconstructed and then loaded in the Main Network to process the input.62 This representational choice significantly reduces the number of stored elements and allows each63 singular program to effectively represent a unitary function of the active program.64 The active program matrix P can be composed by standard low-rank approximation as65 P = USVT = rm∑ n σnunv > n ( 1 ) where U and V are matrices of the left and right singular vectors , and S the matrix of singular values.66 rm is the total number of components we want to retrieve . { σn } rmn=1 is the attended singular values,67 { un } rmn=1 and { vn } rm n=1 the attended singular vectors of S , U , and V , respectively . The Program68 Memory is crafted as three singular program memories { MU , MV , MS } –each of their memory69 slot stores a singular component or singular program . The process “ codes ” the active program70 using singular programs from the program memories . The coding is conditioned on input xt , yet71 we drop index t for notation simplification and leave the details on the computation of σn , un , vn in72 Sec . 2.2.73 The Program Memory also maintains the status for each singular program in terms of access and74 usage . To access a singular program , program keys ( k ) are used . These keys are low-dimensional75 vectors that represent the singular program function and computed by a neural network that ef-76 fectively compresses the singular program . The program usage ( m ) measures memory utilisation,77 recording how much a memory slot is used in constructing a program . The components of the78 Program Memory are summarised in Fig . 1 ( c ) .79 Recurrent multi-head program attention mechanisms for program storage and retrieval80 Neural networks use the concept of differentiable attention to access memory [ 11 , 2 ] . This de-81 fines a weighting distribution over the memory slots essentially weighting the degree to which each82 memory slot participates in a read or write operation . This is unlike conventional computers that use83 a unique address to access a single memory slot.84 Here we use two kinds of attention . First is content-based attention [ 11 , 12 ] to ensure that the singu-85 lar program is selected based on its functionality and the data input . This is achieved by producing86 a query vector based on the input and comparing it to the program keys ( k ) using cosine similarity.87 Higher cosine similarity scores indicate higher attention weights to the singular programs associated88 with those program keys . Second , to encourage better memory utilisation , higher attention weights89 are assigned to slots with lower program usage ( m ) through usage-based attention [ 12 , 31 ] . The90 attention weights from the two schemas are then combined using interpolating gates to compose the91 final attention weights to the Program Memory.92 We adapt multi-head attention [ 11 , 38 ] that applies multiple attentions in parallel to retrieveH singu-93 lar components . Besides , we introduce a recurrent attention mechanism , in which multi-head access94 is performed recurrently in J steps . The j-th set of H retrieved components is conditioned on the95 previous ones . This recurrent , multi-head attention allows the composition network to incrementally96 search for optimal components for building relevant active programs.97 Neurocoder learns to “ code ” a relevant working program via training98 The structure of the Program Memory and the role of the Program Controller facilitates the au-99 tomatic construction of working programs via training . The Program Controller controls memory100 access through its composition network that creates the attention weight defining how to weight the101 singular programs in the memories . A weighted summation of the singular programs results in the102 attended singular program . Applying the recurrent multi-head attention described earlier , multiple103 attended singular programs are retrieved to construct an active program ( Eq . 1 ) . Then the Program104 Controller generates a residual program using its integration network , adding to the active program105 to produce the working program of the Main Network . This addition enables creation of flexible106 higher-rank working programs , which compensates for the low-rank coding process . The structure107 of the Program Controller is illustrated in Fig . 1 ( b ) .108 The singular programs are trained to represent unitary functions necessary for any computation109 whilst the composition and integration networks are trained to compose the relevant programs for110 the considering task . As such , beside minimising the task loss , we enforce orthogonality of stored111 singular vectors by minimising Lo = MUM > U − I+MVM > V − I . The parameters of the networks,112 and the stored singular programs are adjusted using gradient training via minimising the total loss113 L = Ltask + aLo ( 2 ) where Ltask represents the supervised task loss and Lo represents the orthogonal loss weighted by114 a hyper-parameter a to enforce orthogonality of the singular vectors.115 2.2 Attention mechanisms for Program Memory116 Here we describe program attention mechanisms used in this paper . Given wuin , w v in , w σ in ( jointly117 denoted as wu , v , σin ) –the attention weight to the i-th slot of the singular program memories MU , MV118 and MS , we retrieve the n-th singular vector as follows,119 un = Pu∑ i=1 wuinMU ( i ) ( 3 ) vn = Pv∑ i=1 wvinMV ( i ) ( 4 ) For the singular values , we need to enforce σ1 > σ2 > ... > σrm > 0 , thus we retrieve using120 σn = softplus ( ∑Ps i=1 w σ inMS ( i ) ) n = rm σn+1 + softplus ( ∑Ps i=1 w σ inMS ( i ) ) n < rm ( 5 ) Here , Pu , Pv and Ps are the number of memory slots of MU , MV and MS , respectively . In this121 paper , we set P = Pu = Pv = Ps as the number of memory slots of the Program Memory . We note122 that these notations are specified for some data input xt and the index n later maps to an attention123 head h , and an attention step j , hence the full notation should be wu , v , σtijh . To simplify notations , we124 will drop u , v , σ from now and describe the computation of a representative wtijh for any of the125 three program memories in the following parts.126 Recurrent Access to the Program Memory via the composition network127 To perform program attention , the Program Controller employs a composition network ( denoted128 as fθ ) , which takes the current input xt and produce program composition control signals ( ξ p t ) .129 If fθ performs all attentions concurrently via multi-head attention as in [ 11 , 38 ] , it may lead to130 program collapse [ 22 ] . To have a better control of the component formation and alleviate program131 collapse , we propose to recurrently attend to the program memory . To this end , we implement fθ as132 a recurrent neural network ( LSTM [ 16 ] ) and let it access the program memory J times , resulting in133 ξpt = { ξptj } J j=1 . At access step j , the recurrent network updates its hidden states and generates ξptj134 using recurrent dynamics as135 ξptj , hj = fθ ( xt , hj−1 ) ( 6 ) where h0 is initialized as zeros and ξ p tj is the program composition control signal at step j that136 depends on both on the input data xt and the the previous state hj−1 . Particularly , the control signal137 contains the queries and the interpolation gates for each head to compute the program attention138 weight : ξptj = { qtjh , gtijh } H h=1 . Here , at each attention step , we perform multi-head attention with139 H as the number of attention heads and thus , each ξptj consists of H pairs of queries and gates.140 Hence , the total number of retrieved components rm = J ×H and the index n = j ×H + h.141

Summary: This paper proposes a general-purpose neural architecture with an explicit goal of modularity, called Neurocoder. Neurocoder dynamically produces the weights for a layer in a neural network, based on the input, by combining "programs" from memory. These programs consist of left and right singular vectors ($U$, $V$) and singular values ($S$). For a given input, a recurrent multi-head process is used to retrieve the program from memory. In each step of the recurrent process, we obtain a query $q$ and gate $g$ for each of the $H$ heads and for each of the $U$, $V$ and $S$. The query $q$ is used for content-based attention inside $U$, $V$, and $S$ respectively. The gate $g$ is used to combine the content-attention result with a least-recently-used attention. By putting all of this together we obtain a $u_n$, $v_n$, and $\sigma_n$ for each head in each step of the recurrent process. We then sum $\sum_n \sigma_n u_n v_n^T$ to obtain the weights for the layer.  The authors evaluate this architecture on MNIST, CIFAR, RL on Atari games, and polynomial regression.