Relational Memory-Augmented Language Models (2022)

Abstract: We present a memory-augmented approach to condition an autoregressive language model on a knowledge graph. We represent the graph as a collection of relation triples and retrieve relevant relations for a given context to improve text generation. Experiments on WikiText-103, WMT19, and enwik8 English datasets demonstrate that our approach produces a better language model in terms of perplexity and bits per character. We also show that relational memory improves coherence, is complementary to token-based memory, and enables causal interventions. Our model provides a simple yet effective way to combine an autoregressive language model and a knowledge graph for more coherent and logical generation.

## introduction a core function of language is to communicate propositions (e.g., who did what to whom). as such, language models need to be able to generate this information reliably and coherently. existing language models @xcite @xcite do not have explicit representations for such information and rely on it being implicitly encoded in their parameters @xcite @xcite . this encoding mechanism makes it difficult to interpret what the language models know and often leads to generating illogical and contradictory contents. for example, @xcite observe that existing language models rely heavily on word correlation and fall short of logical reasoning. this causes the model to hallucinate-for example, that barack obama's wife is hillary clinton based on the high co-occurrence of the two entities. in another example, lake and murphy (2020) notice that gpt-2 @xcite , * work completed during an internship at deepmind. 2019) states that unicorns have four horns, directly after speaking that unicorns have one horn. in this work, we explore ways to combine an autoregressive language model with a knowledge graph. we design a memory-augmented architecture that stores relations from a knowledge graph and investigate the effect of conditioning on this relational memory in an autoregressive language model. in contrast to existing token-based memory-augmented language models that store context-target pairs @xcite , our memory stores relation triples @xcite . relation triples form the basis of knowledge bases, empowering a wide range of applications such as question answering @xcite , machine reading @xcite , and reasoning @xcite . from a cognitive science perspective, we can consider the neural language model to be an instance of system 1, which performs fast inference and the symbolic relational memory as a world model to support slow and logical reasoning of system 2 @xcite . 1 we hypothesize that relational memory can improve performance and coherence of an autoregressive language model. given an observed context, we first run an entity tagger to identify entities in the context. we then use tf-idf @xcite to select salient entities. we retrieve relations (from a knowledge base) for the selected entities and design a gating function that allows the language model to adaptively combine information from extracted relations and observed textual context to predict the next token. existing knowledge bases such as freebase and wikidata can be used as a source of information from which to retrieve relations. however, they are often incomplete and do not contain relations that are suitable for the particular dataset that we want to work with. instead of using these predefined knowledge bases, we choose to perform open information extraction (openie) on each language modeling dataset to get relations. as a result, our model is able to move beyond simple co-occurrence statistics and generate text that is more grounded on real-world relations observed in a particular corpus. our main contributions are as follows: ## model an autoregressive language model defines the probability of a sequence of tokens p(x) = p(x 1 , . . . , x t ). it is common to factorize this joint probability as a product of conditional proba-bilities with the chain rule @xcite : our language model is based on transformer-xl ( §2.1) which is augmented with a relational memory ( §2.2). we discuss them in detail below. ## experiments our experiments seek to evaluate the effect of augmenting language models with a relational memory. we introduce datasets used for evaluation ( §3.1), discuss implementation details ( §3.2), and present our main results ( §3.3). we then show ablation studies and further analysis of our model ( §4). ## analysis in this section, we study several design choices of relational memory, including its knowledge source, input component, capacity, dynamic openie, entity scoring method used, and speed comparison. we then show quantitative and qualitative analysis results to better understand our model. ## conclusion we presented relationlm, a language model that is augmented with relational memory. we showed how to obtain relevant knowledge graphs for a given corpus and how to combine them with a state-of-the-art language model such as transformer-xl. we demonstrated that our model improves performance and coherence on wikitext-103, wmt19, and enwik8. we also performed a comprehensive analysis to better understand how our model works. our model provides a way to combine an autoregressive language model with general knowledge graphs.