NON-PARAMETRIC LANGUAGE MODELS
Year: 2022

Abstract:
Structural locality is a ubiquitous feature of real-world datasets, wherein data points are organized into local hierarchies. Some examples include topical clusters in text or project hierarchies in source code repositories. In this paper, we explore utilizing this structural locality within non-parametric language models, which generate sequences that reference retrieved examples from an external source. We propose a simple yet effective approach for adding locality information into such models by adding learned parameters that improve the likelihood of retrieving examples from local neighborhoods. Experiments on two different domains, Java source code and Wikipedia text, demonstrate that locality features improve model efficacy over models without access to these features, with interesting differences. We also perform an analysis of how and where locality features contribute to improved performance and why the traditionally used contextual similarity metrics alone are not enough to grasp the locality structure.

text: Structural locality is a ubiquitous feature of real-world datasets, wherein data points are organized into local hierarchies. Some examples include topical clusters in text or project hierarchies in source code repositories. In this paper, we explore utilizing this structural locality within non-parametric language models, which generate sequences that reference retrieved examples from an external source. We propose a simple yet effective approach for adding locality information into such models by adding learned parameters that improve the likelihood of retrieving examples from local neighborhoods. Experiments on two different domains, Java source code and Wikipedia text, demonstrate that locality features improve model efficacy over models without access to these features, with interesting differences. We also perform an analysis of how and where locality features contribute to improved performance and why the traditionally used contextual similarity metrics alone are not enough to grasp the locality structure.

heading: 1 INTRODUCTION
text: Language models (LMs) predict a probability distribution over sequences, and are most widely studied to model and generate natural languages (Bengio et al., 2003; Merity et al., 2018; Baevski & Auli, 2018; Brown et al., 2020). Advances in LMs benefit many natural language processing downstream tasks, such as machine translation (Bahdanau et al., 2015), dialog systems (Sordoni et al., 2015), question answering (Yang et al., 2019; Raffel et al., 2019), and general representation learning for natural language (Devlin et al., 2018; Liu et al., 2019). Recently, LMs have also been adopted to model sequences other than text, such as source code written in programming language (Hindle et al., 2016; Hellendoorn & Devanbu, 2017; Alon et al., 2020; Karampatsis et al., 2020), which can enable useful downstream tasks like code completion (Raychev et al., 2014).
Most current neural LMs are based on parametric neural networks, using RNN (Mikolov et al., 2010) or Transformer (Vaswani et al., 2017) architectures. These models make predictions solely using a fixed set of neural network parameters. Recently, more and more neural LMs also incorporate non-parametric components (Grave et al., 2017; Guu et al., 2018; He et al., 2020; Khandelwal et al., 2020), which usually first select examples from an external source and then reference them during the prediction. For example, Khandelwal et al. (2020) model the token-level probability by interpolating the parametric LM probability with a probability obtained from the nearest context-token pairs in an external datastore. Using such non-parametric components in LMs is beneficial because the model no longer needs to memorize everything about the language in its parameters.
For such non-parametric LMs, one important concept is a distance metric between the current context and other contexts in the datastore. One example of such metric is the `2 distance between context vectors calculated by the parametric model (Khandelwal et al., 2020). This distance can be used in both retrieval and probability calculation; items in the datastore that are less distant from the current context are more likely to be retrieved and have a higher influence on the final probability. However, given that non-parametric datastores are typically very large, containing a myriad of contexts from disparate sources, calculating a metric that accurately reflects semantic similarities is non-trivial; as we demonstrate in experiments, there is much room for improvement in current practice.

heading: Nobel
text: In this paper, we argue that the relevance of contexts may be correlated with not only contextual distance, but also structural characteristics of the underlying data. Specifically, we take advantage of a property we dub structural locality, the propensity of text to be divided into local groups sharing common hierarchical attributes. This property is ubiquitous across many kinds of texts and can provide additional information on how closely related two different examples are to each other. Throughout this paper, we will provide two case-studies of this phenomenon. First, in the domain of programs written in source code, if two source files originate from the same project, they are more likely to be related than files from other projects, and even more so if they are from the exact same package (Hellendoorn & Devanbu, 2017). Second, in natural language, two sections of Wikipedia text may be more related if they fall within the same topical domain, are from similarly titled sections, or even are from the same article (as in Figure 1). Notably this locality often manifests itself at different levels, such as the levels of “project”, “subdirectory”, and “file” cited above for source code.
In this paper, we hypothesize that by using multiple levels of structural locality, we can better calibrate the distance metrics used to retrieve examples from non-parametric datastores, thereby improving LM performance. Specifically, we propose a simple-yet-effective approach that can easily be applied to non-parametric LMs: we use different levels of structural locality to define functions that modify the contextual distance metrics used by the non-parametric module.
We evaluate our method on two drastically different domains: Java programming language source code, and natural language Wikipedia articles, achieving noticeable LM performance gains in both by adding just 5 & 7 parameters respectively. Moreover, we perform an in-depth analysis showing how the traditionally used contextual similarity metrics alone are not enough to grasp the locality structure, providing evidence for why adding the locality features is indeed useful. We also compare programming languages and natural languages to highlight several interesting differences in terms of how, and how much, the locality helps improve LM performance.

heading: 2 NON-PARAMETRIC LANGUAGE MODELS
text: Given a linguistic context consisting of a sequence of tokens ct = (w1, ...wt−1), autoregressive parametric LMs estimate p(wt|ct; θ), the probability distribution over the next token wt. Such parametric LMs store information regarding the language being modeled in the parameters θ. The size of θ is fixed in advance based on the hyperparameters of the model architecture, in recent years typically a neural network (Grave et al., 2016; Baevski & Auli, 2018; Dai et al., 2019; Brown et al., 2020). In contrast, a non-parametric LM’s number of parameters is not determined by just the model architecture, but also by the underlying data used to train the model. While non-parametric LMs using Bayesian statistics have existed for some time (Wood et al., 2011; Shareghi et al., 2017; He et al., 2020), they have recently seen increased prevalence through the introduction of neural LMs that retrieve relevant examples from an external datastore (Hashimoto et al., 2018; Guu et al., 2018). In particular, we focus on kNN-LMs (Khandelwal et al., 2020), a variety of such models that uses a nearest neighbor retrieval mechanism to augment a pre-trained parametric LM, achieving impressive results without any additional training.
Neural network-based LMs usually map the context c to a fixed-length vector representation, with a trained function f(c). In kNN-LMs, the non-parametric component consists of a collection (D) of contexts for the kNN to retrieve from. Denoting these contexts and their corresponding next token as
During inference, the parametric component of the LM generates the output distribution over next tokens pLM (wt|ct; θ) and the corresponding context representation f(ct), given the test input context ct. Then the non-parametric component of the LM queries the datastore with f(ct) representation to retrieve its k-nearest neighbors N according to a distance function d(·, ·). We can then compute a probability distribution over these neighbors using the softmax of their negative distances. The model aggregates the probability mass for each vocabulary item across all its occurrences in the retrieved targets. This distribution is then interpolated with the parametric LM distribution pLM to produce the final kNN-LM distribution:
pkNN(wt|ct) ∝ ∑
(ki,vi)∈N
1wt=vi exp(−d(ki, f(ct))) (2)
p(wt|ct; θ) = λpkNN(wt|ct) + (1− λ)pLM(wt|ct; θ) (3) In our experiments, we follow Khandelwal et al. (2020) in setting the interpolation factor λ to 0.25.

heading: 3 DEFINING STRUCTURAL LOCALITY
text: We define structural locality as a categorical feature calculated between a pair of contexts (ci, cj) in a collection of data, that describes whether the pair share some common, potentially hierarchical, attributes (e.g., the section title of a Wikipedia article section, or the directory path of a source code file). For each domain, a set of hierarchical attributes {l0, l1, ..., ln} can be defined based on prior knowledge of the domain. We denote lk(ci, cj) ∈ {0, 1} as the boolean locality feature value for the context pair, representing whether ci and cj share the same hierarchical attributes lk. Here, l0 is reserved for “no locality”, in case the pair shares none of the attributes. Without loss of generality, we set a constraint that ∑ k lk(ci, cj) = 1, as new features can be introduced by conjunction and negation of the attributes if needed.
Specific Instantiations. We instantiate these features on our two case studies of Wikipedia text and Java source code, as summarized in Table 1.
In Wikipedia, for every context ci, we define four mutually exclusive hierarchical attributes, l0 − l3. We calculate these features based on the Wikipedia article and section titles, using simple pattern matching. We then link each article to a set of categories (one article may belong to multiple categories) using the knowledge graph WikiData,1 by aggregating all the category entities involving two properties: P31 (instance of) and P279 (subclass of). The criterion for “same section title” is exact string match (Hayashi et al., 2020). If there is at least one common category between the sets of categories for two articles, the pair is assigned the “same article category”.
For Java source code, we define 3 mutually exclusive attributes, l0 − l2 based on the location of the code. For each source file, we use the full file path to obtain the two attributes: project name and sub-directory path.2 The criterion for both “same project” and “same subdirectory” is exact string match. Note that these features are strictly hierarchical, hence only two features are used to capture specific locality here.
An Aside: Connections to Domain Adaptation. Domain adaptation typically refers to reusing existing information about a given problem (e.g., data or model) to solve a task in a new domain.
1https://www.wikidata.org/ 2For example, full path .../Journal.IO/src/main/java/journal/io/api/DataFile.java has project
Journal.IO and sub-directory src/main/java/journal/io/api/ for package journal.io.api.
Domain adaptation for neural models generally focuses on fine-tuning models on in-domain data (Sennrich et al., 2016; Chu et al., 2017) or making direct modifications to the model to consider domain information (Britz et al., 2017) or latent topic features (Khudanpur & Wu, 2000; Mikolov & Zweig, 2012; Wang & Cho, 2016). Most of these methods do not natively support new test-time contexts that were not seen at training time. In comparison, one immediate advantage of non-parametric LMs is the ability to adapt to different domains at test time without re-training (Merity et al., 2016; Grave et al., 2016; 2017; Khandelwal et al., 2020). For example, some adaptive LMs (Grave et al., 2016; 2017) make use of the previous hidden states of test documents dynamically during inference. Similarly, our proposed locality features do not require re-training on the training set.
Note that within the scope of this paper, although connected, the proposed structural locality is a different concept from domain. We consider domains as higher-level classifications describing the text where one example belongs to one domain label; e.g., a section about Kim Kardashian’s early life belongs to a category of texts describing celebrities. One the other hand, with the structural locality, a user could define multiple levels of locality: to that same section, we can assign not only the domain label, but also, the section title “Early Life”. The lightweight nature of our model combined with non-parametric LMs also makes adding more levels of features straightforward, as the features only need to be calculated for the top nearest neighbors, and the number parameters that need tuning in our proposed method (Section 5) is only about twice the number of locality features.

heading: 4 STRUCTURAL LOCALITY AND NEAREST NEIGHBORS
text: In this section, we examine the relationship between distances derived from neural LM features d(f(ci), f(ct)), structural locality features l(ci, ct), and the accuracy of the next-word prediction wi. Specifically, the underlying assumption of the kNN-LM is that less distant contexts will be more likely to accurately predict the next word wt. We would like to test whether this correlation between distance d(·) holds uniformly across different locality levels l(·), or if locality provides additional information indicative of whether a particular context is useful for predicting wi beyond just that provided by the neural representations.
Data. We use two different corpora from different domains to examine this question.
WIKITEXT-1033 is a standard language modeling benchmark (Merity et al., 2016) consisting of natural language text from English Wikipedia. It contains a 250K token, word-level vocabulary, with 103M tokens in the training set and 250K tokens in both the validation and test sets.
JAVA GITHUB4 is a programming language corpus containing Java source code from Github (Allamanis & Sutton, 2013) that is widely used in source code modeling (Hellendoorn & Devanbu, 2017; Karampatsis et al., 2020). It contains 1.44B tokens from 13,362 projects in the training split, 3.83M tokens from 36 projects in the validation split and 5.33M tokens from 38 projects in the test split. The splits are separated by whole projects. The dataset is tokenized with byte-pair encoding (Sennrich et al., 2015) using the vocabulary from Karampatsis et al. (2020) with 2,000 subtokens.
Base Model. As the neural model used to calculate context features, we follow Khandelwal et al. (2020),5 train an LM with the exact architecture and optimization described by Baevski & Auli (2018): a decoder-only Transformer (Vaswani et al., 2017), with 1024 dimensional hidden states for the WIKITEXT-103 dataset and 512 for JAVA GITHUB. We set the number of retrieved nearest neighbors to be analyzed to 1024, and the distance metric to `2 following the default.
Datastore. To capture the effect of our proposed locality features, the datastore should ideally be both closely related to the test examples, sufficiently large to ensure precise kNN retrieval performance for a wide range of contexts, and not too sparse in terms of the prevalence of locality features.
For WIKITEXT-103, we include the training set, as well as the validation/test set (excluding the text currently being evaluated) in the datastore. For the JAVA GITHUB, due to the relatively large size of the validation/test set, and the unwieldy size of the training set, we include only the validation/test set (also excluding the current file).
3https://blog.einstein.ai/the-wikitext-long-term-dependency-language-modeling-dataset/. 4https://zenodo.org/record/3628665. 5https://github.com/urvashik/knnlm
Analysis. Consider k nearest neighbor contextsNt = {cr|r = 1...k} retrieved for any test context ct in the test set C, ordered by the ascending distance: ∀r : d(cr, ct) < d(cr+1, ct). We define r ∈ [1, k] as the “rank” for the retrieved context cr. To study the quality of the retrieved contexts, we calculate the number of correctly retrieved tokens, defined as #{wr = wtgold} across C.
We plot in Figure 2, from left to right: (1) Negative distances {−d(cr, ct)|cr ∈ Nt, ct ∈ C} grouped into bins, vs. the retrieval accuracy of the bin avg(#{wr = wtgold}). (2) Rank r vs. the retrieval accuracy at rank r, avg(#{wr = wtgold}). (3) Rank r vs. the average negative distance avg(−d(cr, ct)) at rank r. All of the plots are grouped by different locality levels l0 to ln. Naturally, the left-most sub-figures reflect that the lower the (negative) distance, the lower the accuracy on both datasets. Yet, interesting, on the Wikipedia domain (Figure 2a), as the negative distance gets close to 0 (perfect match), the retrieval accuracy for the next word does not always increase; the accuracy values in this range have very high variance and all 4 levels of locality show no clear trend. This partly indicates that context-based distance is imperfect, regardless of locality. Even so, at slightly lower distances, the trends stabilize and largely show a consistent picture: more specific locality features, especially those involving the same category (l1&l3) yield better predictions than the locality-insensitive component for identical distances. This is especially significant at higher ranked retrievals (middle sub-figure), where contexts that share the same section title and the same article category are substantially more likely to share the same completion. This suggests that the proposed locality features are not fully represented by, or correlated with the original distance metric, and thus implies that there is room for improvement by incorporating these features.
In the Java source code domain (Figure 2b), we generally observe that the retrieval accuracy is much higher than that in the Wikipedia domain, suggesting that the kNN is doing a better job retrieving relevant contexts. This is largely due to higher repetitiveness of source code (Hindle et al., 2016); as we show later, the base Transformer model also performs much better here than on natural language text (Section 6.2). We also observe a more pronounced locality effect here: at the same distances close to 0 and at the same rank, neighbors that are local to the current context have far higher accuracy, indicating usefulness of locality features in the source code domain as well. However, as we can see
from the (right-most) plot of rank versus the negative distance, the average distances of the neighbors with higher locality levels are also significantly smaller than the distance of those without locality. This suggests that the distance in the Java source code domain already correlates well with the level of locality, which may render incorporating locality features less beneficial. We study the precise benefit under one instantiation of this model next.

heading: 5 INCORPORATING STRUCTURAL LOCALITY IN NON-PARAMETRIC LMS
text: Now that we have demonstrated that locality is additionally indicative of next-word prediction accuracy beyond context distance, we propose a method to incorporate this information into the nonparametric retrieval module. In the case of kNN-LMs (Section 2), recall that pkNN is calculated based on the softmax of the negative distance −d(f(ci), f(ct)). Assuming locality features {l0, l1, ..., ln} for each pair (ci, ct) consisting the retrieved nearest neighbor and the current inference context ct, we modify the formulation of pkNN (Equation 3) to consider these features as below:
pkNN(wt|ct; {θn}) ∝ ∑
(ki,vi)∈N
1wt=vi exp(−g(ki, ct; {θn})) (4)
g(ki, ct; {θn}) = gn(d(ki, f(ct)); θn) if ln(ci, ct) = 1. (5) where gn(d(·, ·); θn) is a learnable function of the distance of the nearest neighbors, with parameter θn for each type of locality feature ln. One can view function g(·) as a “modified” distance for nearest neighbors after taking locality information into consideration. In our experiments, we adopt a linear form of g(·): gn(d(·, ·);wn, bn) = wnd(·, ·) + bn (6) We omit the bias for g0(·) by setting b0 = 0 to remove one free parameter from the model and potentially make optimization easier.6 To learn these functions, a user needs to provide just a small sample of annotated data in the same domain, as there are only 2n+ 1 parameters to optimize. In our experiments, we use the validation split for optimization. The parameters are trained to minimize the negative log-likelihood of the kNN prediction of the gold token:
argmin {θn} − log pkNN(wt = wtgold |ct; {θn}) (7)
To optimize the parameters, we use the Adam (Kingma & Ba, 2014) optimizer with a learning rate of 0.0001 on the validation set for 200 epochs. It converges within 20 minutes for both datasets.

heading: 6 HOW DOES STRUCTURAL LOCALITY IMPROVE LANGUAGE MODELING?

heading: 6.1 EXPERIMENTAL SETUP
text: Baselines. Since we base our model on kNN-LMs, this model will be our most directly comparable baseline. We also compare our model to the underlying parametric LM (Baevski & Auli, 2018), without the kNN module. For the JAVA GITHUB dataset, we additionally compare to the recent state-of-the-art model from Karampatsis et al. (2020) on code language modeling, which uses BPE and LSTMs. In all experiments, the maximum number of tokens per input sample is 3,072.
Evaluation. We evaluate the performance of the LM with the standard perplexity metric and token prediction top-k accuracy on the held-out data.7 The top-k accuracy is calculated by checking if the ground truth token is among the predicted top-k list. This metric, primarily for k ∈ {1, 5} (with more k values in Appendix A.2), is commonly used to evaluate predictive models of source code (Hindle et al., 2016). In order to more easily incorporate and analyze the locality features, and also following Karampatsis et al. (2020), we split the evaluation dataset into independent test examples to evaluate, where each of the example is an atomic unit in the locality hierarchy. For JAVA GITHUB, each test example is a source code file, and for WIKITEXT-103, each example is an article section.8
6We also experimented with an adaptive variant that conditioned the weights and biases ({wn}, {bn}) on the current context representation f(ct) parameterized by a MLP. However, this did not result in significant improvement over directly optimizing w and b (Appendix A.3).
7For JAVA GITHUB, the perplexity is calculated on full tokens by aggregating the likelihood of subtokens. The accuracy is calculated that all subtokens in a full token have to be predicted correctly.
8Note that because we predict WIKITEXT-103 section-by-section instead of article-by-article the perplexity numbers reported here are somewhat worse than other works. Article-by-article calculation is not inherently

heading: 6.2 RESULTS
text: The high-level results are shown in Table 2. At first glance, we can already see that the two datasets vary greatly in predictability. With a similar Transformer architecture, performance on JAVA GITHUB is much better than on the WIKITEXT-103 across the board, partly due to the rigid nature of programming language syntax. With a Transformer model, we achieved a strong state-of-the-art language model on Java code, with low perplexity and very high prediction accuracy (∼75%).
By adding kNN module onto the Transformer-based LMs, perplexity and accuracy in both domains improves by a large margin. This is expected and in line with previous experiments on kNNLMs (Khandelwal et al., 2020). The Wikipedia domain enjoys less relative improvement in perplexity (16%) than the Java source code domain (29%). This is particularly interesting, considering that the datastore used for WIKITEXT-103 contains both the current held-out split and the training data (∼100M contexts), compared to that for JAVA GITHUB with only the current held-out split (∼5M contexts). This reflects the fact that source code is known to benefit strongly from project- and package-specific locality (Tu et al., 2014; Hellendoorn & Devanbu, 2017).
Adding proposed locality features and finetuning the parameters on the validation set improves the performance further on both datasets, albeit with a smaller relative gain. This confirms our hypothesis that incorporating locality into the non-parametric retrieval-based LMs is beneficial. We also see that locality features in the Wikipedia domain result in fairly consistent gains, while the Java source code domain sees especially strong accuracy improvements. This echoes our analysis of the source code corpus in Section 4, where we found that distance was generally
strongly correlated with accuracy, but that locality was particularly informative at low distances. There, it may help discriminate between top-ranked completion candidates (as also shown later in Tab. 4). It is notable that despite the fact that the perplexity and accuracy on JAVA GITHUB are already very strong with the vanilla Transformer, we still see a noticeable relative error reduction of 4.9% by adding locality levels information.
We next study how locality features guide towards a “better” distance distribution among nearest neighbors. We plot the relationship between the nearest neighbor ranks and “modified” distance g(ki, ct) in Figure 3. Table 3 shows the specific learned parameters for each level of g(·, ·). Evidently, the biases and weights vary accordingly with locality levels, as the model tries to “correct” the original distance by emphasizing more local contexts. Compared with the original negative distance −d(ki, f(ct)) depicted in Figure 2, the negative modified distance is more separated between the different locality levels on either dataset, showing the relative importance of different locality more clearly. We analyze several alternative approaches to parameterizing locality in Appendix A.3.
For WIKITEXT-103, comparing Figure 3a with Figure 2a, we can see that with the original distance different localities cluster together, and the modified distance separates them much better. We can also
incompatible with our proposed method, but it would require additional implementation to use different locality features for different locations in the output. Hence, we used section-by-section calculation for expediency.
see that if two contexts have the same article category and the same section title, then their distance on average is the closest, closely followed by those sharing article categories only. On the other hand, contexts that only share section titles are not as closely related. This is intuitively reasonable; the education section for a computer scientist and a musician can be very different in content, even if sharing some commonalities, like the phrases used to describe courses, grades, school locations, etc. This also underscores the usefulness of explicit locality features in providing interpretable insights into the relevance of domain knowledge.
For JAVA GITHUB, comparing Figure 3b with Figure 2b, we can see that the original distance is already more separated between different locality levels than that of WIKITEXT-103, again suggesting better learned representations (in terms of locality sensitivity) for the Java domain. However, the model still benefits somewhat from the contexts that are under the same subdirectory (more so than just the same project), especially for top nearest neighbors: the gap for ranks higher than 80 is more pronounced with the modified distance. This again verifies our hypothesis about the hierarchical nature of structural locality. It also indicates potential practical applications – if this model were deployed in a code editor, one could obtain representations of the files in the same sub-directory as the current file and use them, along with the proposed locality features, to bias auto-complete results.
Table 4 shows a randomly sampled test context from each domain where pkNN for the gold token increases after using locality features. We can see that the nearest neighbor search using context representations performs reasonably well at capturing patterns and themed phrases, especially closer to the last token, finding two very similarly rated candidates. However, in both examples, the second retrieved candidate has a wrong target token. Informed by locality features – in the WIKITEXT-103 example, a matching section and category for the first candidate – the more “local” context enjoys a large boost in probability, while the non-local one’s decreases slightly. We present additional examples in Appendix A.1. The JAVA example demonstrates the same effect; the second retrieved example shows resemblances in variable name and template structures, but the fact that the project is focused on Google API rather than Twitter API makes the original retrieval undesirable.

heading: 7 CONCLUSION
text: In this paper, we propose a novel9 method of incorporating structural locality into non-parametric LMs that reference retrieved examples from a datastore. We evaluate this approach in both a natural language and programming language domain, and empirically explore the similarities and differences of how structural locality affects LMs in these settings. The improvement in perplexity and prediction
9Previous work (Hellendoorn & Devanbu, 2017) made the observation that source code files from the same GitHub repository or sub-directory tend to be relatively similar, but did not include an empirical analysis of this effect. Rather, their observation was backed up by improved performance of their n-gram language model with multiple tiered caches. Our work gives more fine-grained insights into this phenomenon, expands the applicability to neural models and new domains, and proposes a more generalized formulation for encoding multiple localities across multiple domains. See Appendix A.4 for a detailed discussion of the connection to previous work and novelty.
accuracy across both domains show the effectiveness and ubiquity of such locality information. Besides language modeling, we also envision that the method can benefit other applications that could be enhanced using user-defined prior domain knowledge such as conditional generation or representation learning using retrieved information.
Limitations. Our method applies to settings where locality effects are present, there is sufficient data to build a reliable datastore for each locality level, and that locality is not already meaningfully captured by the model. While this may not apply to every domain, these features are common: besides source code & Wikipedia, domains including books (features: authorship & dewey decimal system information), research papers (venue, research area), product manuals (kind, sections), and online discussions (time, topic) are all plausible candidates. The features in our studied domains were selected based on availability and prior knowledge of the domain (e.g., for Java, Hellendoorn & Devanbu (2017)). While they did provide measurable improvements and were natural to interpret, these may not be the optimal choice, and other options are worth investigating. It is also possible that LM improvements will eventually lead to learned context representations that almost perfectly capture the relevant locality information. However, we believe this to be unlikely: in many practical settings, there is some inherent ambiguity in partial contexts that cannot be solved with the surface text only. For instance, in Java source code files, it is common to declare a package, which will obviously match perfectly based on the first few tokens (e.g., package org.) with many other contexts. Yet given the scoped nature of this declaration, locally retrieved continuations are inherently far more useful.

heading: ACKNOWLEDGEMENTS
text: We thank Uri Alon for the helpful discussions and thorough feedback. We also thank the reviewers for the helpful conversations during the revision period of the paper. This work was supported in part by the National Science Foundation under Grant #1815287, and a gift from Amazon AWS.

heading: ETHICS STATEMENT
text: There are several ethical considerations regarding our proposed method. First, while language models have a large variety of positive uses in generating natural language (Li et al., 2021) or source code (Allamanis et al., 2018), there is also a potential for dual use. For example, previous work has cited the potential to generate fake news (Zellers et al., 2019) or undesirable/defamatory content (Wallace et al., 2019). Our methodology improves the accuracy of language models, which has the potential to increase their applicability not only in positive use cases but also in ethically questionable scenarios. In order to mitigate these risks, methods to detect machine generated content may be employed, although these are not perfect remedies (Zellers et al., 2019).
Second, because our methodology explicitly references the training corpus in generation it may increase the likelihood of copying content more-or-less verbatim from the training text. This raises potential issues of copyright violation (Chen et al., 2021) or privacy violation (Carlini et al., 2021). However, at the same time, because non-parametric models increase traceability through direct references to the training corpus, it also provides a tool to identify the provenance of the original content, providing a tool to identify such verbatim copying compared to standard parametric models.
Finally, there has been much recent interest in the energy and environmental impact of large language models (Strubell et al., 2019). Due to the necessity to look up information in a datastore, nonparametric models have additional computational overhead compared to parametric models. However, at the same time, as noted by Khandelwal et al. (2020), non-parametric models also provide a tool for quickly adapting to new domains through the use of domain-specific datastores, obviating the necessity for domain-specific fine-tuning. Our work takes this a step further, allowing models to leverage locality of the datastores, potentially making this an even more attractive alternative for efficient adaptation.

heading: REPRODUCIBILITY STATEMENT
text: The source code package containing a README document on how to reproduce the results and analysis and experiment scripts is available in the paper’s supplementary material. The details about the dataset used, model hyper-parameters, and analysis performed are described in Section 4 and Section 6.1. All experiments are conducted on a single machine with a 48 core CPU and 8 NVIDIA V100 32GB GPU. For WIKITEXT-103 we use the pretrained model provided by (Khandelwal et al., 2020) for fair comparison. For JAVA GITHUB the Transformer model is trained until it converges, requiring approximately 2 days. The datastore size is about 5GB for JAVA GITHUB and 351GB for WIKITEXT-103.

heading: A APPENDIX

heading: A.1 ADDITIONAL EXAMPLES
text: We provide additional examples on WIKITEXT-103 in Table 5.

heading: A.2 ADDITIONAL RESULTS ON TOKEN PREDICTION ACCURACY
text: We show additional results on top-k (k = 10, 20) accuracy and relative error reduction (RER) on two datasets in Table 6.

heading: A.3 ALTERNATIVE FORMULATIONS TO LEARN PARAMETERS FOR LOCALITY FEATURES
text: An alternative way to incorporate locality features into the model is an adaptive variant that conditions the weights and biases ({wn}, {bn} in Equation 6) on the current context representation f(ct)

heading: Dataset Model Top-10 RER Top-20 RER

heading: Dataset Model Dev PPL Test PPL
text: parameterized by a MLP:
[w0 ... wn b1 ... bn] T =MLP (f(ct)) (8)
In our experiments, we used a two-layer MLP with ReLU activations, with 64 hidden units and 0.3 dropout rate during training. The perplexity results compared with directly optimizing weights and biases ({wn}, {bn}) are shown in Table 7. We find that contextualizing the parameters does not result in significant improvements over directly optimizing w and b, and sometimes makes the performance even worse. This is perhaps because the context vector space is very large (512-1024 dimensions) compared to the relatively few data points from the validation set used to train.
In Section 6.2, we discuss the effect of learned parameters for each locality level. Observing that the bias terms (bi) and weights (wi) vary according to the locality levels in the learned parameters and to study the weights of the non-local level w0, we freeze all weights except for non-local weights (wi>0) to 1 and only optimize bias terms and the weight for the non-local level (w0). This is to exacerbate the effect of bias on different locality levels. The learned parameters are shown in Table 8. We see similar results where the bias terms vary aggressively to modify the “distance” with different levels of locality, and the weights for the non-local level are less than 1, lowering the importance of those non-local retrieved candidates. It’s worth mentioning that for JAVA GITHUB these learned biases are much larger in amplitude than before, to compensate for the small scale weights learned before (only around 0.03). However, the perplexity results on both datasets are slightly worse than the full optimization setting that we use in the main experiments (19.33 vs. 19.16 in WIKITEXT-103 and 2.15 vs. 2.13 in JAVA GITHUB).

heading: A.4 CONNECTION WITH RELATED WORK AND NOVELTY
text: Previous work (Hellendoorn & Devanbu, 2017) has made the observation that source code files from the same GitHub repository or sub-directory tend to be relatively similar, but did not include an empirical analysis of this effect. Rather, their observation was backed up by an improved performance of their n-gram language model with multiple tiered caches. Our work improves on this in a number of ways, including 1. directly examining the internal representations of a neural language model, 2. demonstrating that the internal representations do not sufficiently capture structural locality
features, 3. providing efficient ways to compensate for this disconnect, leading to improved language modeling performance, and 4. showing that this carries over to Wikipedia, which has not been previously examined in this way. As a result, our work both gives more fine-grained insights into this phenomenon and expands the applicability to neural models and new domains. In addition, our work proposes a more generalized formulation for encoding multiple localities across multiple domains than the one proposed in Hellendoorn & Devanbu (2017), which treated locality as strictly nested (e.g. project→ sub-directory→ file). Our formulation in Eq. 5 can encode more general hierarchies, such as the lattice we used in the Wikipedia case:
same section & category
same section same category
any section & category

Review summaries:
This work concerns itself about utilizing structural locality inherent in real-world datasets in improving the effectiveness of non-parametric language models. It makes a claim that a) structural locality is not implicitly fully captured by the distance metric used in non-parametric language models and further that b) explicitly plugging in structural locality into non-parametric language models can improve their effectiveness. It validates this claim first by doing analysis of two datasets with the help of custom locality functions and then by plugging in the locality functions into a non-parametric language model with the help of learnable parameters.  

The paper is about modelling structural locality in non-parametric language models. The key hypothesis is in modelling not only the co-occurrence characteristics but also structural characteristics such as locality. The paper explains the key claims via case studies conducted on source code data and Wikipedia datasets.  The model paradigm is based on non-parametric language models. A key difference between the non-parametric model and the parametric counterpart is that in the non-parametric model the model parameters are not only determined by the model architecture but also the underlying data.  Structural locality, which is different from just co-occurrence counts, models the structural relationships between pairs of items, e.g., whether they belong to the same or different directory in the case of source code.  The optimisation model is presented in Equation 7 where the authors need a small sample set emanating from the same domain to train the model. The authors then conduct experiments to demonstrate that the method improves upon existing works. Both qualitative and quantitative experimental results are shown.

The authors propose an approach to complement context by adding 'locality' information in examples present in external stores of non-parametric language models. The locality information captures the hierarchical structure, and deems two contexts more similar (or having less distance) if they share common hierarchical structure/attributes. The authors conduct experiments on source code as well as natural language articles. They analyze the results to point out reasons for improvement, and also highlight differences between the domains.