Adaptive Semiparametric Language Models (2021)

Abstract: We present a language model that combines a large parametric neural network (i.e., a transformer) with a non-parametric episodic memory component in an integrated architecture. Our model uses extended short-term context by caching local hidden states—similar to transformer-XL—and global long-term memory by retrieving a set of nearest neighbor tokens at each timestep. We design a gating function to adaptively combine multiple information sources to make a prediction. This mechanism allows the model to use either local context, short-term memory, or long-term memory (or any combination of them) on an ad hoc basis depending on the context. Experiments on word-based and character-based language modeling datasets demonstrate the efficacy of our proposed method compared to strong baselines.

## introduction human language processing is facilitated by complex systems interacting together. a core component that enables such a process is human memory. memory in humans consists of specialized systems, which form a basis for intelligent behaviors @xcite @xcite . for language processing, working (short-term) memory is a temporary storage that can be used to comprehend sentences and follow conversations. episodic (long-term) memory stores individual experience and events. semantic memory stores facts and knowledge about words and concepts. 1 in artificial language processing systems (e.g., language models), a popular approach to design a better model is by encoding all of the desired knowledge (to produce grammatical sentences, process long text, remember events, etc.) in the weights of a large parametric neural network via end-to-end training. we see an increasingly larger transformer become a better language model @xcite @xcite . in this scale approach, the knowledge is implicitly represented in the weights of a parametric neural network, and it is not straightforward to interpret whether a model contains a particular knowledge without asking the model to produce a response-for example, via a cloze-style question @xcite or a prompt @xcite . an alternative strategy is to design a modular architecture that separates memory storage and computational processing, where each module has a clear purpose. recent progress in memoryaugmented neural networks has given rise to many variants of memory-augmented transformer language models that fall under this category. for example, attempts to incorporate extended local context to a neural network-such as those found in neural cache @xcite , transformer-xl @xcite , compressive transformer @xcite , performers @xcite , longformer @xcite , and reformer @xcite )-can be seen as models of working memory. models of episodic memory include knn-lm @xcite and architectures that are designed for more complicated tasks such as question answering (de masson d' @xcite and machine translation @xcite . in machine learning and natural language processing, memory-augmented neural networks is used to refer to all types of memory systems. in this paper, inspired by the modular design of human memory systems, we present a language model architecture (spalm) with storage modules that resemble working and episodic memory systems, which we combine with a large parametric neural network that is responsible for computation ( §2). our hypothesis is that encouraging each component to focus on a specific function (e.g., storing long-term information, capturing extended context, modeling local information) facilitates easier training that produces an overall better language model. 2 specifically, we follow transformer-xl @xcite to capture extended context by caching hidden states in a temporary short-term memory. for long-term context, we use a persistent key-value database and perform sparse retrieval with (approximate) k-nearest neighbors. in contrast to previous language models that either interpolate output probabilities @xcite @xcite or use input concatenation @xcite to combine information from different sources, we design a context-dependent gating mechanism to incorporate local, extended, and global context. we discuss similarities and differences to related work in §3. in language modeling, many tokens can be predicted from their local context without requiring long-term information. our model can adaptively decide whether the current (local) context is enough, or whether it needs to use information from the short-term and/or long-term memory. in §4, we compare spalm with strong baselines-including transformer-xl and knn-lm-on word-based and character-based language modeling. our positive results establish the benefit of the proposed architecture. they also indicate the generality of our approach and its potential applicability to other sequence modeling tasks. we analyze how spalm uses long vs. short-term context ( §5) to better understand how the model operates when making predictions. we conclude by discussing limitations and future directions ( §6). ## model we consider a language model that takes as input a sequence of words @xmath0, . . . , x t } and outputs a probability distribution of the next word 2 we note that spalm is not intended to be a model of human language processing system. we merely take inspirations from human memory systems to design a better artificial language model. figure 1 : our language model architecture has three main components: (i) a transformer that processes the current local context, (ii) a short-term memory module that stores hidden states from an extended context, (iii) and a key-value (hidden state-output token) database that stores compressed long-term context. at each timestep, our model combines the current context and short-term memory with a mechanism similar to transformer-xl. it then retrieves a set of past output tokens that are used in a similar context from the longterm memory module. these past output tokens are then encoded and aggregated to a single vector that represents long-term information. we use a contextdependent gate to combine information from multiple sources for making a final prediction. p(x t+1 | @xmath1; w). given a corpus of t words, the log likelihood of the corpus is: spalm consists of three main components: (i) a large parametric neural network in the form of a transformer to process local context, (ii) a short-term memory to store extended context, and (ii) a non-parametric episodic memory module that stores information from long-term context. we integrate these components in a single architecture with a gating mechanism. figure 1 shows an illustration of our model, which we discuss in detail below. ## comparisons to previous work knn-lm. there are several language models that are related to our proposed method. the closest one is knn-lm @xcite , which is another language model that is augmented with a nearest neighbor retrieval mechanism. knn-lm is an ensemble technique that is designed to be used only at evaluation time. in knn-lm, a pretrained language model (e.g., a transformer) is combined with another retrieval-based language model by interpolating their probabilities: the interpolation weight λ is tuned at the corpus level on a development set. while this post hoc integration method used by knn-lm has its merit (e.g., very practical, fast to incorporate to any model since it does not require additional training), our focus is on designing a model that combines short-term and long-term memory at the architecture level. our motivation is twofold. first, interpolating the language model weights at the corpus level forces the model to use the same interpolation weight λ for p lm and p knn for each token in the corpus. it cannot adaptively combine short-term and long-term information at the token level based on the context. in addition, λ needs to be tuned on an extra development set. foot_3 spalm, on the other hand, is able to adjust the weights placed on m t and h r t when constructing z t differently for different tokens. second, we believe that integration of different memory modules at the architectural level is a more natural approach that could help pave the way for applications with other memory sources (e.g., knowledge bases, images, videos)-where the memory output is not in the same space as the prediction output (i.e., words) and an interpolation technique cannot be used. we compare with knn-lm in our experiments. since interpolating model probabilities is an ensembling technique that is independent of the architecture, we also show that our language model can be furher ensembled with p knn if necessary. cache-based language models and pointer networks. cache-based language models @xcite store pairs of hidden states and output tokens from previously seen tokens (within a limited context length) in a cache. the best variant of the method uses an interpolation (ensemble) method similar to knn-lm to combine information from the cache and the backbone language model. this class of models temporarily stores m past hidden states (typically, in the order of thousands), so it is a working-memory model as opposed to long-term memory. in addition, they also rely on interpolating probabilities of a backbone language model and a cache component (similar to knn-ln when the cache size is unbounded). ## experiments we use word-based and character-based english language model datasets-wikitext 103, wmt, and enwik8-to evaluate our proposed method. we provide descriptive statistics in ## analysis we have demonstrated the efficacy of our proposed method on three language modeling tasks. in this section, we analyze the model to gain more insights into how it works.