Capturing Structural Locality in Non-parametric Language Models

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 . 1 INTRODUCTION . 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 . Nobel . 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 . 2 NON-PARAMETRIC LANGUAGE MODELS . 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 . 3 DEFINING STRUCTURAL LOCALITY . 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 .

Summary: 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.