Augmenting Transformers with KNN-Based Composite Memory

1 INTRODUCTION . Machine learning solutions to various tasks , such as game-playing or dialogue , are often dependent on external information . This information can take multi-modal forms , including structured knowledge bases , free text , and images , and also comes in overwhelmingly large quantities . A pressing challenge is to create models that can identify which specific elements of multiple information sources are relevant , and incorporate them into standard architectures on each task . Previous work has explored incorporating large external memories into neural network layers ( Weston et al. , 2014 ; Sukhbaatar et al. , 2015 ; 2019 ; Lample et al. , 2019 ) . Many existing approaches focus on using attention over the memory slots , which is computationally intensive and becomes less effective as the the size of the memory grows . In this work , we propose representing multiple sources of external information as fixed encodings and using K Nearest Neighbors search to fetch relevant information . KNN search is computationally efficient and scalable , and libraries like faiss ( Johnson et al. , 2019 ) allow KNN to be easily used on GPUs and integrated into neural networks . As the external memories are kept fixed , they do not require any training to learn the memories along with the model . We can thus scale more easily to larger memories by learning only the KNN-based read operation to identify relevant information from the memory . Our core contribution proposes an efficient , KNN-based Information Fetching ( KIF ) module that can access relevant external knowledge , combine knowledge from different sources , and integrate this information into standard sequence to sequence architectures . We apply these flexible modules to two dialogue datasets , challenging tasks where generative models can leverage external information to write coherent , on-topic responses . We show that relevant information can be identified from hundreds of thousands of candidates in a multi-modal , multi-knowledge-source setting to improve the performance of generative dialogue models . On both datasets , we achieve state of the art results compared to generative models and match the quality of retrieval models . 2 RELATED WORK . Incorporating External Knowledge into Neural Networks . Augmenting neural networks with memory , or longer term components that can be accessed with read and write operations , has been explored in various proposed architectures . For example , Memory Networks ( Weston et al. , 2014 ; Sukhbaatar et al. , 2015 ; 2019 ) introduce attention mechanisms over large external memories . Neural cache models ( Grave et al. , 2016 ) simplify these to access previous memories with a dot product . Previous work has studied how to read and write into these memory architectures ( Rae et al. , 2016 ; Graves et al. , 2014 ; Joulin & Mikolov , 2015 ) . Another line of research has focused on computational scalability for larger external memories . For example , Chandar et al . ( 2016 ) propose a hierarchical memory network rather than a flat one and Rae et al . ( 2016 ) learn sparse operations to read and write . Lample et al . ( 2019 ) focus on learning memories of up to one million slots and how to efficiently access the slots using product keys . Beyond explicit memory representations , it may be possible to store information implicitly during training time by memorizing common patterns present in text ( Petroni et al. , 2019 ) . We focus on learning to fetch relevant information from multiple explicit external multi-modal knowledge sources and integrate them into one network . Work has also focused on computationally efficient softmax operations ( Mnih & Hinton , 2009 ; Grave et al. , 2017 ; Chen et al. , 2015 ) . Many approximate softmaxes use KNN-like operations to form clusters , and the softmax operation is constrained by the slow calculation of the exponential . Our use KNN benefits from efficient and scalable libraries such as faiss and nmslib . Generative Dialogue . We develop a general architecture for incorporating external information and apply it to the case of generative dialogue models . Previous work in dialogue has leveraged knowledge as necessary information to accomplish the task . For example , airline and restaurant booking tasks often use API calls to access information about reservation times and availability ( Bordes et al. , 2016 ) . In contrast , our work focuses on how to incorporate unstructured knowledge , such as free text found on the web . Previous work has employed architectures that attend over the available knowledge and identify relevant pieces of information , which scales poorly with large quantities of information ( Dinan et al. , 2018 ; Qin et al. , 2019 ; Lian et al. , 2019 ) . In this work , we replace the use of attention over external information with the output of our KNN module . On the modeling side , work has explored both generative ( Serban et al. , 2016a ; b ) and retrieval based models ( Zhang et al. , 2018 ) , which identify the best utterance from the training set to return as the dialogue response . This often leverages self-attention or cross-attention mechanisms ( Humeau et al. , 2019 ) . Further work has explored hybrid models , for example using the output of a retrieval model as input for a generative model ( Dinan et al. , 2018 ; Weston et al. , 2018 ) . We extend these approaches by augmenting generative models with retrieval-like operations based on KNN search , allowing dialogue models to flexibly incorporate various sources of external knowledge . 3 KNN-BASED INFORMATION FETCHING MODULES . Broadly , the KNN-based Information Fetching ( KIF ) module assumes a model M can access inputs X = x1 , x2 , . . . , xn to produce outputs O = o1 , o2 , . . . , on . In a setting without additional supporting information , the model will process inputs to make output predictions : M ( xi ) = ôi . However , in many tasks , additional information is present , represented as E = { e1 , e2 , . . . , em } . To incorporate E into M , we encode each element of X and E to a fixed-size vector representation . This can be accomplished in a variety of ways , for example with an encoder neural network . Then , to make predictions , the model encodes xi and uses K Nearest Neighbors to find the closest related information in E. The representations of the identified nearest neighbors are combined in a weighted sum , where each of the k retrieved neighbors is weighted by its similarity to xi . These operations are differentiable , so they can be incorporated into neural networks in a straightforward way . All elements of the knowledge source E are pre-computed and kept fixed — we do not backpropagate to affect the embeddings of the pre-encoded knowledge . However , this lack of backpropagation can introduce a mismatch between the encoding of E and the model that is training , as the training model has constantly changing representations because the weights are being learned . The model must learn a function to align its representations to the external memory . To circumvent this misalignment , we learn a mapping operator fE ( x ) that maps elements of X into the information representation space E. Concretely , fE ( x ) is a multi-layer perceptron with ReLU nonlinearities . From the input elements of X , fE ( x ) learns a representation of an output close to the corresponding projection of X into E. This can be interpreted as learning a read operation on a fixed external memory . If there was no change to the encoding of the model compared to the pre-computed knowledge , then the ideal mapping operator would be the identity function . However , as the model changes significantly during the training process , the nonlinear mapping capability of fE ( x ) is essential to be able to identify the correct knowledge E from the input X . Thus , a model augmented with KIF will incorporate external knowledge in the following manner . First , we find the k nearest elements to the projection of xi in E based on KNN search using inner product , and then the relevant elements are encoded by M . KIFi = { M ( e ) | e ∈ KNearest ( E , k , fE ( xi ) ) } ( 1 ) These elements are weighted by their nearest neighbor scores and summed . This is then concatenated to the representation of xi and used by M to form the prediction : M ( [ xi , WeightedSum ( KIFi ) ] ) = ôi ( 2 ) This is easily extended to using multiple modules simultaneously . For instance , two sources of information , E1 and E2 , can be combined by identifying the top candidates of each information source . The weighted sum of the KIF output on each information source is concatenated with xi . Finally , different sources of information may not be required for every prediction and some information sources can be more important than others . To allow the model to make more fine-grained decisions about what information to use from what source , and how much of it , we add a gating mechanism using a sigmoid function around each weighted sum of KNN representations . KIF1i and KIF2i denote the KIF module from Equation 1 applied to E1 and E2 respectively . WS1i = WeightedSum ( KIF1i ) ( 3 ) WS2i = WeightedSum ( KIF2i ) ( 4 ) M ( [ xi , σ ( WS1i ) · WS1i , σ ( WS2i ) · WS2i ] ) = ôi ( 5 ) 4 APPLYING KIF TO DIALOGUE TASKS . We describe how to apply our method to the task of generative dialogue , a challenging setting where models must autoregressively generate engaging and on-topic responses . We investigate dialogue for two main reasons : first , dialogue agents must be able to consult relevant information to maintain the topic of the conversation . Second , retrieval-based agents have strong performance compared to generative ones , due to their ability to copy dialogue utterances from the training set . Using KIF , we can incorporate the benefits of retrieval architectures into generative , knowledge-based models . KIF for Generative Dialogue In a dialogue setting , xi represents the text of the conversation i . A conversation consists of multiple back-and-forth utterances ( or turns ) . For example , a conversation could consist of 4 turns : xi = [ xi,1 , xi,2 , xi,3 , xi,4 ] where xi,4 is the direct utterance the model should respond to , and the earlier utterances are the conversation context . Standard generative dialog models use a Transformer neural network as M and want to produce an output oi that is an appropriate response to the conversation . However , in many cases , the conversation history alone does not include all of the information required to produce an appropriate response . To incorporate knowledge , models often concatenate a knowledge source E such as Wikipedia to xi , such that M ( [ xi , e1 , e2 , . . . , en ] ) = ôi , and use attention modules to identify the most relevant knowledge . However , this approach is computationally intensive when handling large quantities of information . Further , attention mechanisms have been found to operate poorly over long sequences , as the mechanism is blurry and struggles to make fine-grained decisions ( Fan et al. , 2018 ) . The same is true for hierarchical approaches , which lack scalability . We augment Transformer sequence-to-sequence ( seq2seq ) networks with KIF . We experiment on two dialogue tasks , Wizard of Wikipedia ( Dinan et al. , 2018 ) and Engaging Imagechat ( Shuster et al. , 2018 ) . We use faiss ( Johnson et al. , 2019 ) to perform KNN search efficiently and at scale . Wizard of Wikipedia The goal of the Wizard of Wikipedia dataset is to train knowledgeable agents that can chat in any domain . The dataset contains 1,365 various topics discussed in 18,430 dialogues in the training set , totalling 166,787 training utterances . The topic is included as the first utterance of the conversation . The dataset includes relevant Wikipedia sentences for each turn of the chat , identified by an information retrieval system . Our model for Wizard of Wikipedia has access to two sources of external information , E1 and E2 : • E1 is Wikipedia Knowledge provided by the dataset as evidence to support knowledgeable chitchat . The scale of this KNN search is to filter through an average of 34 sentences . The KIF module uses dialogue features to fetch relevant knowledge to condition upon to generate the subsequent utterance . • E2 is Training Utterances . To incorporate the benefits of retrieval-based dialogue models to the generative setting , we use KIF to identify relevant utterances from the training set and take their responses as input . If many conversations about dogs have already occurred , models should be able to take advantage of these human-written examples to improve their generations . There are around 170K dialogue utterances as inputs to KNN search . This can be interpreted as incorporating the benefits of retrieval models by identifying an utterance with similar structure as the text the model would like to generate . Access to these two sources of knowledge can be seen as as learning a template and a topic separately . Sample templates can be identified from the training utterances , and topic-specific information learned by accessing the Wikipedia knowledge . To better identify relevant training utterances from the large quantity available , we break down xi into conversation sub-features for a more fine-grained match in the KNN search step . We concatenate the encoding of the most recent dialogue utterance ( e.g . xi , last ) with the encoding of the dialogue context and the turn number . These are known to be salient conversation features . The most recent dialogue utterance is the direct turn the model is responding to , and the dialogue context may provide additional clues . The turn number is important , as earlier turns are often more generic ( e.g . hi , how are you doing today ) and later turns are more specific . Engaging ImageChat The goal of Engaging ImageChat is to create agents capable of chitchatting about images , selected from the YFFC100M dataset ( Thomee et al. , 2015 ) . The dataset contains 186,782 dialogues in the training set , each about a unique image , totalling 355,862 utterances . Agents are assigned one of 215 personalities ( e.g . sweet ) to increase engagingness . We use a MultiModal neural network designed to handle both image input and text input . Following Shuster et al . ( 2018 ) , the images are encoded using a ResNeXt ( Xie et al. , 2017 ) . To extract the final image representation , we project the 2048-dimensional output of the image encoder to 512-dimensions using a deep multi-layer perceptron with ReLU activation units . The conversation history , which includes the personality , is encoded with a Transformer encoder network . The image and conversation are combined using the Multimodal Sum Combiner proposed in Shuster et al . ( 2018 ) . Our model for Engaging Imagechat has access to two sources of external information , E1 and E2 : • E1 is Chat on Similar Images . While there are over 180K different images used in this dataset , many of the images are similar . For example , conversations associated with two pictures of dogs could be relevant to each other . The model is able to fetch from around 160K different images and returns 6 turns of related chat for each image . Fetching from E1 Model F1 F1 ( Seen ) ( Unseen ) Retrieval Trans . MemNet * 15.4 12.4 2-Stage Generative MemNet * 18.9 17.4 Generative Trans . MemNet * 16.9 14.4 + Reddit Pre-Train 17.6 16.3 Retrieve and Refine 18.2 17.9 KIF-Augmented Transformer 26.9 23.3 Table 1 : Results on the Wizard of Wikipedia dataset . * denotes results from Dinan et al . ( 2018 ) Model F1 Retrieval Trans . * 9.81 Generative Trans . MemNet 7.1 + Reddit Pre-Train 12.8 Retrieve and Refine 13.6 KIF-Augmented Transformer 14.4 Table 2 : Results on the Engaging Imagechat dataset . * denotes results from Shuster et al . ( 2018 ) consists of identifying Related Image Chats , or conversations on similar topics ( as similar images are likely to have similar conversations ) . • E2 is Training Utterances . Similar to the motivation for the previous dataset , we allow the model to identify training utterances that could be useful for responding in the current conversation . The scale of this fetching task is large : 350K dialogue utterances . This could be interpreted as identifying utterances with similar structure to what the model would like to generate , and is complementary to the topic-based Related Image Chats . To identify relevant information from the training utterances , we use the same dialogue features in the KNN search step , with one modification : we add the personality provided by the dataset . As utterances from speakers with the same personality are likely to be more related , this feature improves the quality of the fetched information .

Summary: This paper uses KNN-based retrieval method for extracting relevant information from different sources for the task of generating response in a dialogue task. They utilized different sources of information from Wikipedia, YFCC image set and dialogue utterances. The method is very straight forward and using nearest neighbors from external information sets as auxiliary information and encoding that information via same neural net to produce the encoded external information. The concatenation of encoded input and auxiliary information is concatenated to produce the output in the dialogue via decoder network. The novelty of the proposed method is incremental over Dinan et al. (2018). The KNN-based retrieval module can produce some drift in dialogue from the actual context of the input which can result in irrelevant response in dialogue. Some failure cases can show the quality changes with those drifts in dialogues.