Learning Associative Inference Using Fast Weight Memory

1 INTRODUCTION . Humans continually adapt in order to understand new situations in changing environments . One important adaptive ability is associative inference for composing features extracted from distinct experiences and relating them to each other ( Schlichting & Preston , 2015 ; Gershman et al. , 2015 ) . Suppose Alice has shared with you pictures of her toddler . Later , at the office party , you see a man carrying the depicted toddler . Since the toddler yields a shared feature in two different contexts , it may be plausible to infer that the man is Alice ’ s partner , without ever seeing him and Alice together . The ability to rapidly associate and bind together novel stimuli can help to derive knowledge systematically , in addition to the knowledge gained directly from observation . Virtually all modern cognitive architectures applied to challenging artificial intelligence problems are based on deep artificial neural networks ( NNs ) . Despite their empirical successes and theoretical generality , NNs tend to struggle to generalise in situations similar to the given example ( Lake et al. , 2017 ; Phillips , 1995 ; Lake & Baroni , 2017 ) . This weakness becomes even more severe if the training and test data exhibit systematic differences ( Atzmon et al. , 2016 ; Agrawal et al. , 2017 ) . For example , during training , the man ’ s representation might never be associated with the toddler ’ s , but during testing , this association might be necessary to make a useful prediction . In problems where humans excel , this sort of inference is likely ubiquitous since data is often combinatorially complex in a way that observations used during training will likely cover just a small fraction of all possible compositions . Such a lack of productivity and systematicity is a long-standing argument against the use of NNs as a substrate of an artificial cognitive architecture ( Fodor & Pylyshyn , 1988 ; Hadley , 1994 ; McLaughlin , 2009 ) . The hidden state of a neural model is a learned representation of the task-relevant information extracted from the input . To generalise to never-seen-before compositions of stimuli , the function which produces the state representation must be able to systematically construct all possible states . This requires a general and preferrably differentiable method , such as the Tensor Product Representation ( TPR ; Smolensky ( 1990 ) ) . TPRs provide a general and differentiable method for embed- 1Source code and data used in this paper is available at github.com/ischlag/Fast-Weight-Memory-public ding symbolic structures in vector spaces . A TPR state representation is constructed via the tensor product ( i.e . the generalised outer-product ) of learned component representations . Under certain constraints , such a mechanism guarantees a unique representation for every possible combination of components ( Smolensky , 1990 ; 2012 ) . In this work , we augment a recurrent NN ( RNN ) with an additional TPR-like memory representation . To facilitate the learning of multi-step associative inference , the TPR memory can be queried multiple times in a row , allowing the model to chain together various independent associations . In contrast to previous work on fast weights , we apply our memory-augmented RNN to much longer sequences . This requires the model to update its associative memory . Furthermore , we demonstrate the generality of our method by applying it to meta-reinforcement learning and small scale language modelling problems . In the next section , we cover related memory-augmented NNs . Section 3 describes the FWM in detail . Section 4 demonstrates the generality of our method through experiments in the supervised , self-supervised , and meta-reinforcement learning setting . The supervised-learning experiments in subsection 4.1 consist of a more challenging version of the bAbI dataset dubbed concatenated-bAbI or catbAbI . The meta-reinforcement learning experiment in section 4.2 demonstrates the FWM ’ s ability to learn to explore a partially observable environment through its ability to perform associative inference . Finally , the self-supervised experiments in subsection 4.3 demonstrate that the FWM can compete with the state-of-the-art word-level language models on small benchmark datasets . 2 RELATED WORK . RNNs such as the Long Short-Term Memory ( LSTM ; Hochreiter & Schmidhuber ( 1997 ) ; Gers et al . ( 2000 ) ) are in theory capable of implementing any algorithm ( Siegelmann & Sontag , 1991 ) . However , the linear growth of the hidden state of a fully connected RNN leads to quadratic growth in the number of trainable weights . Early work addressed this issue through the use of additional memory ( Das et al. , 1992 ; Mozer & Das , 1993 ) and differentiable fast weights ( Schmidhuber , 1992 ; 1993 ) . Recently , memory-augmented NNs have solved algorithmic toy problems ( Graves et al. , 2014 ; 2016 ) as well as reasoning and inference problems in synthetic and natural language ( Weston et al. , 2015b ; Xiong et al. , 2016 ) . Inspired by the random-access memory of computer architectures , a common approach is to incorporate a soft and differentiable lookup table into the NN model . Such slot-based memory matrices have shown to be difficult to train ( Munkhdalai & Yu , 2017b ) and require sophisticated mechanisms for the allocation and deallocation of memory ( Csordas & Schmidhuber , 2019 ) . The TransformerXL ( TXL ; Dai et al . ( 2019 ) ) , an autoregressive language model variant of the Transformer ( Vaswani et al. , 2017 ) , can be understood as a slot-based memory-augmented RNN where every new state is pushed into an immutable queue of finite size . Although it is recurrent , the layers of a transformer architecture are strictly forced to use inputs from a lower layer which limits its generality . Nevertheless , a sufficiently deep and well regularised TXL model has achieved state-of-the-art performance in large scale language modelling tasks . A biologically more plausible alternative of increasing the memory capacity of NNs are fastchanging weights , i.e . stateful weights that can adapt as a function of its input . Non-differentiable fast weights or “ dynamic links ” have been published since 1981 ( von der Malsburg , 1981 ; Feldman , 1982 ; Hinton & Plaut , 1987 ) . Subsequent work showed that a regular network can be trained by gradient descent to control the fast weights of a separate network ( Schmidhuber , 1992 ) or of itself ( Schmidhuber , 1993 ) in an end-to-end differentiable fashion . Recently , fast weights have made a comeback and achieved good results in small toy problems where regular NNs fall short ( Ba et al. , 2016a ; Schlag & Schmidhuber , 2017 ; Munkhdalai & Yu , 2017a ; Pritzel et al. , 2017 ; Ha et al. , 2017 ; Zhang & Zhou , 2017 ; Miconi et al. , 2018 ; 2019 ; Schlag & Schmidhuber , 2018 ; Munkhdalai et al. , 2019 ; Bartunov et al. , 2020 ) . Most memory-augmented NNs are based on content-based or key-based lookup mechanisms . An alternative to the storage of patterns in a lookup table is the idea that patterns are reconstructed through the implicit iterative minimisation of an energy function , such as in the classical Hopfield network ( Steinbuch , 1961 ; Willshaw et al. , 1969 ; Hopfield , 1982 ; Kanerva , 1988 ) or the modern Hopfield network ( Krotov & Hopfield , 2016 ; Demircigil et al. , 2017 ; Ramsauer et al. , 2020 ) . This is often described as an auto-associative type of memory as it reconstructs a previously stored pattern that mostly resembles the current pattern . A much less studied variation is the hetero-associative memory ( see e.g . Kosko ( 1988 ) ) where the retrieved pattern is different from the input pattern . This is more relevant for our use case . We aim to train an LSTM to construct , maintain , and edit its associative memory . The ability to edit Hopfield networks partially is not very well studied . For this reason , we employ a simple ( multi- ) linear hetero-associative memory as it is more closely related to the theory of TPRs ( whose manipulation is well understood ) and because the association is retrieved in a single step . Our work directly builds on two examples of differentiable fast weight memories : the TPR-RNN by Schlag & Schmidhuber ( 2018 ) and the Metalearned Neural Memory ( MNM ) by Munkhdalai et al . ( 2019 ) . The TPR-RNN is a sentence-level model for reasoning on text . It achieves excellent results on the regular bAbI tasks but it underperforms on word-level bAbI ( Schlag et al. , 2019 ) or algorithmic toy problems ( Le et al. , 2020 ) . In contrast , the MNM is a word-level model which augments the LSTM with a fully-connected multi-layer feed-forward network as its memory and trains it using a meta-learning objective . Both , MNM and TPR-RNN were developed on the regular bAbI dataset which only contains short sequences and does not require the model to remove deprecated associations from its memory . In this work , we train on an infinite sequence of bAbI stories where our FWM achieves excellent performance and improves over MNM . We further demonstrate strong performance in small-scale language modelling and meta reinforcement-learning which demonstrates the generality of our contribution . 3 PROPOSED METHOD . Our FWM is a fast-changing , multi-linear map which is controlled by a slowly-changing , non-linear LSTM . The slow weights of the LSTM are regular NN weights which are updated during training by gradient descent . In contrast , the fast weights of the FWM are updated by the LSTM at every step of the input sequence through a Hebb-like differentiable mechanism . This allows the FWM function to change rapidly even during testing—hence the name fast weights . Along with updating the fast weights , the LSTM also generates a memory query which is used to retrieve information that was previously stored . The retrieved information then becomes part of the model ’ s output . 3.1 THE FAST WEIGHT MEMORY Given a sequence of tokens x = ( x1 , ... , xT ) from a vocabulary V , the task of language modelling is to train a model which maximizes the joint probability p ( x ) which we factorize autoregressively p ( x1 : T ) =∏T t=1 p ( xt|x0 : t−1 ) where x0 is an artificial start token.2 In this work , we train an RNN model to encode the input sequence x1 : t into ht , the hidden state of the LSTM , and Ft , the fast weight tensor of the FWM , to maximize the probability of the next token xt+1 . At step t of the input sequence , the input token xt is embedded in a dE-dimensional vector space using a lookup table et = embedding ( xt ) . An LSTM with dLSTM hidden units encodes the sequence of embedded tokens into a fixed size vector representation ht = LSTM ( et , ht−1 ) . The probability distribution over the next token x̂t+1 = softmax ( W ( s ) ( ht + FWM ( ht , Ft ) ) where Ft ∈ RdFWM×d 2 FWM are the fast weights of the FWM at step t and W ( s ) ∈ R|V|×dLSTM . Note that the fast weight matrix Ft is a reshaped third-order tensor Ft ∈ RdFWM×dFWM×dFWM . This allows us to describe third-order tensor operations using matrix multiplications . We ’ ll now describe in detail the FWM function and how its fast weights are updated . 2We use the notation x1 : t to refer to the sequence ( x1 , x2 , ... , xt ) .

Summary: The authors present a working memory model composed of a recurrent neural network trained via gradient descent and an associative memory based on the approach taken by Ba et al. (2016) in "Using Fast Weights to Attend to the Recent Past".  The model consists of an LSTM to which takes the input and its own state from the previous step to produce an output (or new state) which is then passed to a fast weight memory (FWM) module.