On-The-Fly Information Retrieval Augmentation for Language Models (2020)

Abstract: Here we experiment with the use of information retrieval as an augmentation for pre-trained language models. The text corpus used in information retrieval can be viewed as form of episodic memory which grows over time. By augmenting GPT 2.0 with information retrieval we achieve a zero shot 15% relative reduction in perplexity on Gigaword corpus without any re-training. We also validate our IR augmentation on an event co-reference task.

## introduction we are interested in exploring the value of long term episodic memory in language modeling. for example, a language model can be used in january to assign a probability distribution over the statements that will appear in the newspaper in march. but one month later, in february, the distribution over the predictions for march should be updated to take into account factual developments since the previous prediction. long term episodic memory should be taken into account when assigning a probability to a statement. here we take a simple approach in which a pretrained gpt language model @xcite is zero-shot augmented with an episodic memory consisting simply of a corpus of past news articles. conceptually the past news articles are viewed as additional training data which can be legitimately accessed when evaluating on future text. in our most basic experiment we calculate the probability of a future article by first calculating the probability of its first k sentences using the pre-trained gpt model. we then use the first k sentences as a query in an information retrieval system to extract a relevant past article. we then insert the past article following the first k sentences when calculating the probability of the remainder of the future article using the same pre-trained gpt model. this is a zero-shot augmentation in the sense that there is no additional training or fine tuning of the pre-trained model. our results show that this augmentation significantly reduces perplexity. we also present various other experiments including results on fine-tuning the model in the presence of the memory and the effect of this memory on event co-reference. ## related work various language models have utilized external knowledge or long contexts @xcite @xcite @xcite @xcite . but these papers do not address the question of whether additional context or external knowledge is useful as a zero-shot augmentation of large scale pre-trained nlp models. the value of external knowledge has previously been demonstrated for nlp tasks such as natural language inference @xcite , language generation @xcite , knowledge base completion @xcite and question answering @xcite @xcite . however, all those prior works assume the model is small and trained from scratch. as large scale pre-trained models have become more powerful it is not immediately clear whether external resources can still add value. the only work we know of on using external resources in modern large scale models is @xcite where a human curated external lexical resource is used to improve bert. our approach bears some resemblance to neural cache models @xcite . however, neural cache models store past hidden states as memory and accesses them through a dot product with the current hidden states. this is different from retrieving knowledge from a corpus-sized memory. our approach is also somewhat related to memory networks @xcite . memory networks have a memory module which can be learnt jointly with other components. it has shown success in applications such as machine reading comprehension @xcite and visual question answering @xcite @xcite . significant progress in memory networks has been achieved in both architecture @xcite @xcite and model scale @xcite . several papers have formulated, and experimented with, scalable memory networks -memory networks that employ some method of efficiently reading and writing to very large neural memories. this is done with approximate nearest neighbor methods in @xcite and with product keys in @xcite . these large memories are used to provide additional model capacity where the memory contents are trained over a large data set using gradient descent training, just as one would train the parameters of a very large network. it is shown in @xcite that it is possible to insert a large memory as a layer in a transformer architecture resulting a model where the same number of parameters and the same performance can be achieved with half the layers and with much faster training time than a standard transformer architecture. here, however, we are proposing zero-shot augmentation with an external data source used as an episodic memory. the use of key-value memories in @xcite is particularly similar to our model. keyvalue memories were used there in treating a corpus of wikipedia movie pages as a memory for answering questions about movies. as in our system, articles were extracted using word based information retrieval. each article was encoded as a vector which was then given to a question answering architecture. this was shown to improve on automated knowledge base extraction from the same corpus but was still not competitive with human curated knowledge graphs for movies. here we give the text of the retrieved article directly to the language model architecture and focus on augmenting large scale language models. ## model we use the pre-trained transformer gpt 2.0 @xcite . let w w and w p be the subword and position embeddings respectively. let m denote the total number of layers, for a token at time step t, the m-th layer's hidden state h m t is given by: that's to say, for the first k sentences, we directly feed them to gpt to obtain their representations. for remaining sentences, their representations are conditioned on both the first k sentences and the retrieved context c. ## experiments we focus on two tasks: document level language modelling and event co-retrieved . in both tasks we take a document as input and use first k sentences to query the memory. to calculate the perplexity of a document, we compute the log-probability of a document by multiplying byte level probability, @xcite ; sam: sparse access memory @xcite ; kvm: key value memory @xcite ; lmn: large memory network @xcite . memory size is measured in their own words. for information retrieval we use lucene due to its simplicity. given a query document we first do sentence and word tokenization and then use the first k sentences to retrieve top 20 retrieved documents with the default tf-idf distance metric provided by lucene. since too distant document pairs are uninformative and too related document pairs tends to be duplicates of the test article, we further filter those top ranked documents by time stamp, news source and cosine similarity. more specifically, we choose the highest ranked retrieved document that simultaneously satisfies the following three conditions: it comes from a different news source; it appears earlier but within two weeks time window of the test document, and the bag of word cosine similarity between the test and the retrieved cannot be larger than 0.6α where α is the largest bag of word cosine similarity between the test article and any retrieved articles. to support fine-tuning experiments we constructed a corpus of pairs of a query article and a cached retrieved 1 https://github.com/huggingface/pytorch-transformers document. we split the dataset into train/dev/test by query document's time stamp. the train/dev/test size is: 79622,16927,8045. for zero-shot experiments we use the test set of 8045 articles. we do experiments with @xmath0, 2, 5}. to check the quality of query-retrieved pairs, we randomly sample 100 pairs from dev set and compute the bag of word cosine similarity between the two documents. the mean cosine similarity is 0.15. we also manually inspect them: we ask two nlp researchers to annotate the query-retrieved pair as "bad" or "ok" independently, i.e., if two documents are almost duplicates or totally unrelated, then it's "bad", otherwise, it's "ok". among 100 pairs, 83 pairs are "ok", 17 pairs are "bad" due to irrelevance. the cohen's kappa coefficient between two annotations is 0.94. ## conclusion in this paper we propose a method to augment a pre-trained nlp model with a large episodic memory. unlike previous work, we use information retrieval to handle a large external corpus of text and feed retrieved documents directly to language models. evaluation results on language modelling and event co-reference show the promise of our method. to the best of our knowledge, this is the first work that augments pre-trained nlp models with large episodic memory. in principle, the memory-augmented gpt-2 can be used as a variant of gpt-2 for any downstream tasks, such as glue tasks @xcite , although we have not experimented with that here.