Augmenting Transformers with KNN -Based Composite Memory for Dialog (2021)

Abstract: Various machine learning tasks can benefit from access to external information of different modalities, such as text and images. Recent work has focused on learning architectures with large memories capable of storing this knowledge. We propose augmenting generative Transformer neural networks with KNN-based Information Fetching (KIF) modules. Each KIF module learns a read operation to access fixed external knowledge. We apply these modules to generative dialog modeling, a challenging task where information must be flexibly retrieved and incorporated to maintain the topic and flow of conversation. We demonstrate the effectiveness of our approach by identifying relevant knowledge required for knowledgeable but engaging dialog from Wikipedia, images, and human-written dialog utterances, and show that leveraging this retrieved information improves model performance, measured by automatic and human evaluation.

## introduction machine learning approaches to various tasks, such as game-playing or dialog, are often dependent on external information. this information can take multimodal 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 in a particular context, and incorporate them into standard archi-tectures on each task. in this work, we focus on human-machine dialog and how to efficiently retrieve external knowledge that is relevant to the dialog. we consider two scenarios and for each scenario, retrieve two types of knowledge: (i) knowledge about similar dialog contexts and (ii) external knowledge used to ground the conversation into real world information. knowledge about similar dialog contexts allows for a hybrid retrieval/generative approach to dialog where the system response is generated based not only on a representation of the current dialog context and of the relevant world knowledge, but also based on a response retrieved from a similar dialog context. the retrieved knowledge can be viewed as providing information about structure and dialog sentences, or utterances: which response is likely given a similar context? external knowledge is also retrieved to improve the semantic content of the dialog model. in one scenario, wizard of wikipedia @xcite , general topics are provided to crowdworkers, who are asked to have in-depth and specific conversations about these topics by referencing specific wikipedia sentences as knowledge. in this scenario, external knowledge is retrieved from a pre-selected set of wikipedia sentences associated with the current dialog topic. retrieval aims to select the sentence that is most relevant at each step of the dialog and thereby to ground system responses in relevant world knowledge (e.g., by referring to star wars when talking about science fiction). in the other scenario, engaging imagechat @xcite , crowdworkers are provided with images and asked to have a conversation 82 transactions of the association for computational linguistics, vol. 9, pp. 82-99, 2021. https://doi.org/10.1162/tacl a 00356 action editor: masaaki nagata. submission batch: 6/2020; revision batch: 9/2020; published 3/2021. c 2021 association for computational linguistics. distributed under a cc-by 4.0 license. 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 dialog datasets that challenge generative models to leverage external information to write coherent, on-topic responses. both of our chosen tasks require models to leverage external information, such as information from wikipedia or images, to engage in the conversation. we show that relevant information can be identified from hundreds of thousands of candidates in a multimodal, multi-knowledge-source setting to improve the performance of generative dialog models. further, the output of the kif modules is interpretable as specific human-readable know-ledge elements are selected, allowing users to better understand the information the generative model conditions upon when writing the subsequent utterance. on both datasets, we achieve state-of-the-art results compared to generative models and find there is no statistically significant difference in the interestingness or human preference of our model output compared to stateof-the-art retrieval models. ## related work we discuss related work on learning to incorporate external knowledge into neural networks and efficiently access relevant information. we then describe work in generative dialog that incorporates knowledge. ## knn-based information fetching modules broadly, the kif module assumes an encoder model m can access inputs @xmath0, x 2 , . . . , x n }. for example, x can be a collection of sentences, and x i represents an individual sentence. in a setting without additional supporting information, the encoder will process an input x i and produce the encoder output m (x i ). if x i is a sequence such as a sentence, then m (x i ) is a representation of the variable size of the sequence length by the fixed size encoder m 's hidden size. however, in many tasks, additional information is present, represented as @xmath1, e 2 , . . . , e m }. we encode each element of x and e into a vector representation using the encoder. to identify the closest information in e that is relevant to x i , our general approach will be to use knn by comparing the representation of x i with the representation of each element in the set e. knn is a fully differentiable operation @xcite , so can be incorporated in a straightforward way into neural models. the most relevant information in e will then be available in the model. we display a kif-augmented model 1 in figure 1 and describe how the kif module operates. one challenge to overcome is that the representation of all elements of the knowledge source e are pre-computed and kept fixed, creating m (e)-we do not backpropagate to affect the embeddings of the pre-encoded knowledge. in the early stages of training, the model receives large amounts of loss, which would affect the quality of the pre-encoded embeddings if we backpropagated to them. further, encoding the fixed external knowledge once and re-using it allows for greater scalability. however, this lack of backpropagation can introduce a mismatch where ws stands for weighted sum), then concatenated to m ′ (x i ). full description with notation can be found in section 3. to circumvent this misalignment, we learn a mapping operator f e (m ′ (x i )) that trains to map elements of the model's representation of x, or m ′ (x), into the additional information representation space m (e). concretely, f e (m ′ (x i )) is a multilayer perceptron with relu nonlinearities. from the input elements of x, f e (m ′ (x i )) learns representations 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 (as m ′ would equal m ). however, as the model changes significantly during the training process, the nonlinear mapping capability of f e (m ′ (x i )) 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 with inner product. then, the relevant elements identified by knn are re-encoded by m ′ . for example, if element e j is retrieved by kif, it would produce m ′ (e j ). we use the optimized faiss library for knn search, which can conduct billion-scale knn efficiently on gpus. the knn output for an element x i is produced by using faiss to search for the k nearest representations to f e (m ′ (x i )) in m (e). note that as the encoders m and m ′ produce output representations of variable length (for example, in the case where x i is a variable length sequence, such as a sentence), we average across the length dimension to produce a fixed-size representations r to conduct the knn search. then, the kif module output for an element x i is the set of all re-encoded representations of the knn-retrieved knowledge: these elements are weighted by their normalized nearest neighbor scores and then summed. this is subsequently concatenated to m ′ (x i ) to form the final encoder output: this can be easily extended to using multiple modules simultaneously. for instance, two sources of external information, e 1 and e 2 , 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 the encoded input m ′ (x i ). the kif output dimensionality is the same size as the hidden size of m ′ (x i ), so they can be directly concatenated. 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 finegrained 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. kif1 i and kif2 i denote the kif module from equation (4) applied to e 1 and e 2 , respectively. this concatenation represents the output of the encoder m ′ and can be used for various purposes, such as providing the encoder output to a decoder in a sequence to sequence model. ## applying kif to dialog tasks we describe how to apply kif to the task of generative dialog, a setting where models must generate engaging and on-topic responses. we investigate dialog for two reasons: first, dialog 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 dialog utterances from the training set. using kif, we can incorporate the benefits of retrieval architectures into generative, knowledge-based models. ## results we describe the results of incorporating kif modules into transformer networks. we display an example conversation between a human and our model in figure 4 , and show the top scoring wikipedia knowledge and training utterance fetched by kif modules. we compare to various baselines using automatic and human evaluation, and discuss our experiments. we present various ablation settings to understand the key features that make our method function. ## conclusion we present a knn-based information fetching module that learns to identify relevant information from external knowledge sources by learning a mapping-based read operation. kif modules benefit from the scalability and efficiency of knn search, enabling computation with large external memories. we show in the context of two dialog datasets that relevant knowledge can be identified and incorporated to create more engaging, high-quality dialog.