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 solution proposed is the combination of an RNN (LSTM) and Fast Weighted Memory (FWM). The LSTM produces a query to the memory used to retrieve information from the memory and be presented at the model output. It also controls the memory through fast weights that are updated through a Hebbian mechanism. The FWM is based on Tensor Product Representations (TPR). The FWM is differentiable and builds upon the work of TPR-RNN from Schlag and Schmidhuber and Metalearned Neural Memory (MNM) by Munkhdalai et al. In the experimental section, the authors propose a concatenated version of the bAbI dataset to test their model with language modeling and question answering. Further the model is trained on a meta-learning task over POMDPs on graphs, and on language modeling on the PennTree Bank dataset. They show that the LSTM-FWM model generalizes better than without memory and similar models and with smaller capacity.