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: This paper presents an approach that trains local Hebbian learning rules to allow a neural network to perform reasonably well in two types of problems: sequence memorization and prediction. Two approaches to train these models are compared, which are based on meta-gradients and evolutionary strategies. The evolved model is able to perform sequence predictions of length 1000.