Evolving Neural Update Rules for Sequence Learning

1 INTRODUCTION . The field of neural sequence processing has become dominated by neural networks trained by backpropagation , with the best models being transformers currently , and LSTM recurrent networks in the not so recent past . These networks can be used for a range of problems such as label prediction , modelling and reinforcement learning . The basic computational nature of these networks can be characterized as follows : One can design an arbitrary ( differentiable ) forward computation that updates part of the network ( activations ) and then one executes a specific algorithm to do the updates of another set of parameters ( weights - using back-propagation and variants of stochastic gradient descent ) . This has an advantage of freedom to flexibly design a forward computation . However , the hand-coded update algorithm is fixed , and may not be the most effective means of achieving a given training objective . There has been interest , as we review below , in directly finding the full computation ( both activations and weights updates ) . In this paper , we aim to push this approach further by finding update rules that work better and that scale , by comparing different functional forms and methods for optimizing the search for these rules . We pursue what we call end-to-end search for update rules . A good way to explain this is to consider the analogous process in nature . Roughly speaking the brain of an organism is produced from the information of the parents , ” run ” for the lifetime of the organism , and selected based on how it , as well its descendants perform ( for example keep making descendants ) . Over time , this process developed neural networks that can learn , adapt quickly and that possess whichever other abilities were needed for the success of the organisms . Similarly we consider a space of neural networks and search for updates to their parameters , the activations and the weights in our case , that work well on problems we are interested in . This can allow networks to discover whatever updates are needed to solve these problems . We consider parameterization of neural networks that is somewhat similar to classic artificial recurrent neural networks , with a state consisting of activations ht and weights wt at time t. There are , however , two primary differences . The first is that the entire computation ( including learning or whatever computation the network does ) operates online , that is , the network receives input xt and updates according some function ht , wt , pt = fθ ( xt , ht−1 , wt−1 ) ( 1 ) where fθ is the update function parameterized by hyper-parameters θ that are fixed for the lifetime of the network . The pt is the output - in our case the probability over the next character in the sequence ( we will discuss the problems we solve later ) . The second difference is that the activation ht , i of a given neuron i as well as the weight wt , ij between neurons j and i are ( potentially ) small vectors , as in ( Bertens & Lee , 2019 ; Gregor , 2020 ) . Biological neurons were an inspiration for artificial neural networks , and so is the case here . The former are complex objects , and representing the cell body state as well as the synapse state by vectors ( rather then scalars ) should allow for a closer representation of the neuron ’ s computation . Biological neurons furthermore provide a proof of principle that online neural computation is capable of powerful learning . There are several additional motivations for pursuing this approach . Learning general update rules end to end might allow networks to use weights in a more interesting or powerful fashion than that prescribed by back-propagation . Finding online rules might allow for easier hardware implementation , especially in neural hardware architectures ( Modha , 2017 ; Davies et al. , 2018 ) . Developing the ability to search for update rules might help us to discover how brains compute , for example if we can implement what we know about the computational structure and search for what we don ’ t . Finally in artificial life , entities propagate into the future often not based on an objective but simply if they find means of making descendants ( and if those descendants make more descendants and so on ) ( Ray , 1991 ; Aguilar et al. , 2014 ; Soros & Stanley , 2014 ; Gregor & Besse , 2021 ) . Being able to evolve learning algorithms without objective might provide a path to intelligent artificial life . 1.1 RELATION TO PREVIOUS WORKS .. Evolving network weights There are many works that directly evolve network weights , for example ( Wierstra et al. , 2014 ; Salimans et al. , 2017 ) . In ( Stanley & Miikkulainen , 2002 ) they parameterize the weights by compositional pattern producing networks ( CPPN ) and in ( Risi & Stanley , 2010 ) they consider an adaptive version that updates the weights based on pre- and post-synaptic activations . While in one sense this is similar to our approach in that the weights are updated , it is also very different - with a very different parameterization ( CPPN ’ s and no hidden network units ) , and problems ( evolving the system to solve a T-maze task ) , and without a focus on pure end-to-end learning as the parameters are also evolved for the task . End-to-end learning and meta-learning A range of related works explore end-to-end learning and meta-learning , but don ’ t strictly operate in the local update rules setting . Furthermore , many of these works learn the full weights of the network instead of a smaller number of update rule hyper-parameters . In ( Schmidhuber , 1992 ) they consider an end-to-end setting , parametrizing updates of a ( fast ) network , by another ( slow ) network , that outputs vectors used to compute Hebbian updates to parameters . In AutoML zero ( Real et al. , 2020 ) they evolve a sequence of operations such as multiplication and non-linearities that implement a machine learning algorithm ( that includes learning ) . A series of works such as ( Ravi & Larochelle , 2016 ; Metz et al. , 2019 ; Wichrowska et al. , 2017 ; Bello et al. , 2017 ; Li & Malik , 2017 ; Lv et al. , 2017 ) focus on learning generalisations of gradientdescent optimisers that utilise back-propagated gradient information to update weights . Although these works optimise the parameters of the update rule over multiple update steps , the tasks themselves are typically feed-forward in nature , such as image classification . In some works ( Ravi & Larochelle , 2016 ) , they explicitly consider few-shot learning , which is related to online learning , although task relevant information is preserved across tasks in the full network weights . In ( Miconi , 2016 ; Miconi et al. , 2018 ) they consider recurrent neural networks , but introduce perweight state variables that are updated in a Hebbian fashion with parameters of the updates trained by back-propagation , along with a standard set of weights . This is capable of learning fast weight adaptations . Alternative network parametrizations In ( Ha et al. , 2016 ) they consider a recurrent network with weights parameterized by a smaller set of hyper-weights . This is trained by standard back-propagation over a short periods . Because the number of hyper-weights is still relatively large , the system is still relying on back-propagation learning to encode the structure of the current data-sequence , rather than a learning algorithm primarily . Learning local update rules end-to-end There are a number of works more closely related to our objective of end-to-end learning of local update rules ( Bengio et al. , 1992 ; Orchard & Wang , 2016 ; Gu et al. , 2019 ; Munkhdalai et al. , 2019 ; Bertens & Lee , 2019 ; Gregor , 2020 ; Kirsch & Schmidhuber , 2020 ) . For example Bengio et al . ( 1992 ) parameterizes scalar weight updates as a linear combination of terms derived from local scalar activations and previous weights and searches for these parameters using genetic algorithms , stochastic gradient descent or simulated annealing applied to toy problems . In Orchard & Wang ( 2016 ) a population of agents is evolved for the task of foraging in a 2d world . However both the initial weights of the controller network as well as the synaptic update network are evolved , in principle putting good policies in initial weights ( but showing that updates improve performance ) . Bertens & Lee ( 2019 ) ; Gregor ( 2020 ) introduce the idea of using vectors to represent both the activation and weight states , and updates using LSTM ’ s on T-maze tasks in the former and MLP ’ s on sequence memorization respectively in latter . A key drawback in all these works is that the networks used are tiny , of usually less than ten neurons . In Gregor ( 2020 ) , they attempt to use larger hidden layers , but find that the network does all its computation in the input layer . In Kirsch & Schmidhuber ( 2020 ) they consider more general network parameterizations , using LSTM-based local updates and scaling to a larger number of neurons . However they only consider problems such as MNIST classification that do not require an algorithm to learn from long-range time dependencies . That is , because the same classes often appear close to one another in a random sequence , there is a signal over just a few steps ( longer range dependencies might be needed for good performance , however the reported one is low ) . The main contribution of this paper is to find a parameterization of update rules and a method of search that scales to large number of neurons ( we tested up to a thousand hidden units ) and weights ( million ) and that can learn recurrent network training over long time spans ( thousands ) . 2 TASKS . We train the networks on two tasks both defined on a sequence of characters of a text . We use the pg19 data-set ( Rae et al. , 2019 ) containing a large collections of books as our text source . In the first task we consider sequences of length N taken at random points from the text . The N is of the order of ten thousand in our experiments , but in principle it can be arbitrarily long or we should ideally use the entire data-set , this is just the scale we have managed to achieve so far . At each iteration , we initialize a neural network with random weights and zero activations , run it through this sequence online as in eq ( 1 ) , and measure the total log-likelihood : L = − N∑ t=1 log p ( xt|ht−1 ) ( 2 ) We are looking for a set of hyper-parameters θ parameterising the update rule f that minimizes the loss L. The resulting θ should encode activation and weight dynamics that learn a model of the sequence while running through it online . In the second task we play such randomly sampled sequence twice and measure the likelihood on the second part , testing network ’ s ability to memorize . 3 OPTIMIZATION . There are two algorithms we compare for optimizing θ . The first algorithm is ( meta ) -gradient , where we repeat the following process . We sample a mini-batch of randomly selected sequences from the text . We run the network forward and back-propagate , both through the whole batched sequence , and take stochastic gradient optimizer step - we use Adam ( Kingma & Ba , 2014 ) in the experiments . The second technique we employ is natural evolution strategies ( NES ) ( Wierstra et al. , 2014 ; Salimans et al. , 2017 ) . Instead of optimising only the parameters θ , this method instead maintains a search distribution over the parameters . We use separable NES ( SNES ) , in which this distribution is a diagonal normal distribution parameterised by the parameter means and variances . At each iteration , a population of update rules ( population of θ ’ s ) is sampled from this distribution as well as a single sequence from the dataset . The networks are run forward to obtain fitnesses , which are then used to estimate improved values of the parameter means and variances . In contrast to ( Salimans et al. , 2017 ) , we find that updating the variances at each iteration is beneficial . Further details for our implementation are given in App . A.3 .

Summary: The authors design and apply recurrent neural network models with "fast weights" to sequence compression problems. The "slow" parameters that determine how the "fast" weights change are determined by gradient descent or evolution strategies. Experimentally, the latter method appears to perform better for the architectures designed by the authors. The paper focuses on a long-sequence working memory problem and additionally presents some results on language modeling.