Nonparametric Masked Language Modeling
Year: 2023

Abstract:
Existing language models (LMs) predict tokens with a softmax over a finite vocabulary, which can make it difficult to predict rare tokens or phrases. We introduce NPM, the first nonparametric masked language model that replaces this softmax with a nonparametric distribution over every phrase in a reference corpus. NPM fills in the [MASK] solely from retrieving a token from a text corpus. We show that NPM can be efficiently trained with a contrastive objective and an in-batch approximation to full corpus retrieval. Zero-shot evaluation on 16 tasks including classification, fact probing and question answering demonstrates that NPM outperforms significantly larger parametric models, either with or without a retrieve-and-generate approach. It is particularly better at dealing with rare patterns (word senses or facts) and predicting rare or nearly unseen words (e.g., non-Latin script). We release the model and code at github.com/facebookresearch/NPM.

heading: 1 Introduction
text: Current large language models, despite their wide use and impressive performance, are expensive to scale, difficult to update, and struggle with long-tail knowledge and patterns (Kandpal et al., 2022). Recent work follows a retrieve-and-generate approach to partially address these issues (Lewis et al., 2020; Izacard et al., 2022); however, their final predictions are still made by a parametric model. In particular, they still include a softmax over a finite vocabulary, which limits expressivity (Yang et al., 2018; Pappas et al., 2020) and can make them reluctant to predict rare or unseen tokens (e.g., Thessaloniki in Figure 1).
In this paper, we introduce NPM, the first NonParametric Masked Language Model that predicts tokens solely based on a nonparametric distribution over phrases in a text corpus (Figure 1). NPM consists of an encoder that maps the text
into a fixed-sized vector, and a reference corpus from which NPM retrieves a phrase and fills in the [MASK]. It, crucially, does not have a softmax over a fixed vocabulary, but instead has a fully nonparametric distribution over phrases. This is in contrast to a recent body of work that incorporates nonparametric components in a parametric model (Borgeaud et al., 2022; Izacard et al., 2022; Zhong et al., 2022b).
Training such a nonparametric model introduces two key challenges: (1) full corpus retrieval during training is expensive, and (2) learning to predict an arbitrary length phrase without a decoder is nontrivial. We address the first challenge by using inbatch approximations to full corpus retrieval (Wu et al., 2020; Zhong et al., 2022b), and the second by extending span masking (Joshi et al., 2020) and a phrase-level contrastive objective (Oord et al., 2018; Lee et al., 2021).
We perform zero-shot evaluation on 16 tasks including classification, fact probing and question answering. They include temporal shift and wordlevel translation tasks that highlight the need to predict new facts or rare phrases. We compare with
2097
a range of competitive baselines including encoderonly (Liu et al., 2019), encoder-decoder (Raffel et al., 2020), and decoder-only models (Zhang et al., 2022; Brown et al., 2020). We also compare with a retrieve-and-generate approach that feeds a concatenation of the input and passages to parametric models using off-the-shelf retrieval. Results show that NPM is significantly more parameter-efficient, outperforming up to 500x larger parametric models and up to 37x larger retrieve-and-generate models. It is particularly good at (1) predicting rare words (e.g., an entity split into multiple BPE tokens such as Thessaloniki) and (2) disambiguating word senses (e.g., cheap may indicate inexpensive or of very poor quality; Figure 1). Finally, our evaluation on an entity translation task demonstrates that NPM can predict a word consisting of characters that are extremely rare if not unseen (e.g., non-Latin script; Figure 1).
In summary, our contributions are as follows.
1. We introduce NPM, the first nonparametric masked language model that fills in the [MASK] solely from a phrase-level nonparametric distribution over a corpus.
2. We introduce a novel training scheme to train NPM on unlabeled data. We completely remove the softmax over the output vocabulary, enabling an effectively unbounded output space by predicting any n-gram.
3. Zero-shot evaluation on 16 downstream tasks shows that NPM outperforms significantly larger parametric models, are better on rare patterns, scale well, can be efficiently updated at test time, and can predict extremely rare if not unseen tokens (e.g., words in non Latin script).

heading: 2 Related Work
text: Language Models (LMs). Large LMs trained on a vast amount of text are shown to perform a wide range of downstream tasks in a zero-shot manner by converting a task into a cloze format (Radford et al., 2019; Brown et al., 2020). This is possible because a variety of knowledge is encoded in the parameters of the models. Recent work has scaled parametric LMs by adding more parameters (Brown et al., 2020; Rae et al., 2021; Chowdhery et al., 2022) which can be very expensive in practice. Moreover, such models struggle with predicting rare words or entities, and cannot be updated over time.
There has been a recent body of work that incorporates the nonparametric component with a para-
metric LM. We distinguish (1) work that concatenates retrieved text to the input and trains the model with a standard LM objective (Borgeaud et al. (2022); Izacard et al. (2022); so-called retrieveand-generate approaches) from (2) work that retrieves tokens from a large text corpus to estimate a probability distribution that is interpolated with the output distribution from a standard LM (Khandelwal et al. (2020); Yogatama et al. (2021); Zhong et al. (2022b); Lan et al. (2023); so-called kNN models). Our work is closely related to such a line of work and can be seen as an extreme version of the kNN approach with no interpolation. However, our work is the first that models a fully nonparametric distribution by entirely removing the softmax over a finite vocabulary. This offers a range of new functionalities, such as modeling a distribution over phrases, or predicting rare or unseen words.
Bottleneck in softmax. Most if not all language models use a softmax function that gives a categorical probability distribution over a finite vocabulary. Yang et al. (2018) showed that this softmax is a low-rank approximation of a high-rank output space, making the model less expressive. Pappas et al. (2020) discussed that a fixed output vocabulary makes language models resistant to adaptation to new domains and tasks. We share the motivation with such prior work and propose to use a nonparametric output space to address these issues. Moreover, although not explicitly explored in this paper, our work that completely removes the softmax over the vocabulary can make training more efficient, especially when the vocabulary is large (e.g., multilingual models (Conneau et al., 2020)).
Nonparametric models. In nonparametric models, the data distribution is not defined by a fixed set of parameters, but is rather a function of the available data (Siegel, 1957; Hollander et al., 2013). Having complexity that grows as the data grows, they are differentiated from parametric models whose complexity is bounded as a priori. Freeman et al. (2002) noted that the term nonparametric does not imply that they have no parameters, but rather that the number and nature of the effective parameters are flexible and can depend on the data.
Recent work in NLP has explored nonparametric inference without training (Khandelwal et al., 2020; He et al., 2021; Xu et al., 2022), or trained the nonparametric model on the labeled data for a specific downstream task (Seo et al., 2018, 2019; Lee et al., 2021). In contrast, our work trains a fully
Reference Corpus
nonparametric language model without the labeled data and performs a range of tasks zero-shot.

heading: 3 Method
text: We introduce NPM, the first NonParametric Masked Language Model. NPM consists of an encoder and a reference corpus, and models a nonparametric distribution over a reference corpus (Figure 1). The key idea is to map all the phrases in the corpus into a dense vector space using the encoder and, when given a query with a [MASK] at inference, use the encoder to locate the nearest phrase from the corpus and fill in the [MASK].
Encoder-only models are competitive representation models (Patel et al., 2022), outperforming the other two classes of models in classification tasks (Section 5.4). However, existing encoderonly models are unable to make a prediction whose number of tokens is unknown, making their use cases limited without fine-tuning. NPM addresses this issue, since it can fill in the [MASK] with an arbitrary number of tokens by retrieving a phrase.
We first describe inference of NPM assuming a learned encoder (Section 3.1), and then describe how we train the encoder to map the text into a good vector space (Section 3.2).

heading: 3.1 NPM: Inference
text: Overview. The encoder maps every distinct phrase in a reference corpus C into a dense vec-
tor space. At test time, the encoder maps the masked query into the same vector space and retrieves phrases from C to fill in the [MASK]. Here, C does not have to be the same as the training corpus, and can be replaced or scaled at test time without re-training the encoder.
In practice, there is a significant number of phrases in the corpus, and it is expensive to index all of them. We therefore use a technique from Lee et al. (2021) that represents a phrase with token representations of the start and the end of the phrase. In this approach, we index representations of each distinct token in C, and then at test time, use a k nearest neighbor search for the start and the end of the phrase, separately. Consider Figure 2 as an example. We represent a query with two vectors, qstart and qend. We then use each to retrieve the start and the end of the plausible phrases—in this case, c1 and c4, which are the start and the end of Thessaloniki, respectively.
Method. Formally, let C = {c1, · · · , cN} be a reference corpus with N tokens. We first map each token ci into a contextualized, h-dimensional vector ci ∈ Rh by feeding the text into the encoder and take the vector that corresponds to each token: c1...cN = Encoder(c1...cN ).
At inference time, NPM is given a query whose tth token is masked: q1...qt−1,[MASK], qt+1...qL. We replace [MASK] with two special tokens [MASKs][MASKe] and feed it into the encoder to obtain a list of h-dimensional vectors:
q1...qL+1 = Encoder(q1...qt−1,[MASKs],
[MASKe], qt+1...qL).
We then take the vector corresponding to [MASKs] and [MASKe] as qstart and qend, respectively.1
qstart = qt,q end = qt+1.
We then make a prediction via:
argmax v∗∈V∗
∑ i≤j I[v∗ = ci:j ]
(
exp(sim(qstart, ci)) + exp(sim(q end, cj)) ) ,
where V∗ is a set of possible n-grams defined by the vocabulary V and sim is a pre-defined similarity function that maps a pair of vectors into a scalar
1This allows obtaining two vectors without encoding the query twice, e.g., unlike Lee et al. (2021)
[MASKs][MASK e] . This is to obtain the start and
the end vectors for each span as we do at inference.
3.2.2 Training Objective
Key idea. We illustrate an example in Fig-
ure 3. The masked span is ‘the Seattle Sea-
hawks’, thus the model should retrieve a phrase
‘ the Seattle Seahawks’ from other sequences in
the reference corpus when it is given a query
like this at test time. Speci cally, we should
encourage the[MASKs] vector to be closer to
...the Seattle Seahawks... and the[MASKe] vector to
be closer to ...the SeattleSeahawks... , while being dis-
tant from other tokens. We train the model to do
so by approximating the full corpus as the other
sequences in the batch. Concretely, we train the
model to retrieve the start and the end of the span
‘ the Seattle Seahawks’ from other sequences in the
same batch. Note that our masking strategy ensures
that every masked span has a co-occurring span in
the batch (Section 3.2.1).
Obtaining vector representations. Consider the
-th sequence in the batch that consists ofL tokens,
x i = x i1:::x
i L . We denotêx i = x̂ i1:::x̂ i L as a conse-
quence of span masking overx i . Both x i andx̂ i
are fed into the encoder, and each token is mapped
into an h-dimensional vector:2
xi1 · · ·xiL = Encoder(xi1 · · ·xiL), x̂i1 · · · x̂iL = Encoder(x̂i1 · · · x̂iL).
Training objective. We consider a masked span in xi, represented with [MASKs][MASKe], denoted as x̂it, x̂ i t+1. We then denote g i t as the original n-gram that were replaced by x̂it, x̂ i t+1.
We now define the objective for this masked span, and the final objective is summed over all masked spans. The training objective for this masked span is defined as − ( log ∑ y∈Y+s (git) exp(sim(x̂ i t,y))∑
y∈Y+s (git)∪Y−s (git) exp(sim(x̂ i t,y))
+ log
∑ y∈Y+e (git) exp(sim(x̂
i t+1,y))∑
y∈Y+e (git)∪Y−e (git) exp(sim(x̂ i t+1,y))
) .
Here, sim(·, ·) is a similarity function defined in Section 3.1, and Y+s (git), Y−s (git), Y+e (git) and Y−e (git) are start positives, start negatives, end positives and end negatives of git, respectively, which are defined in the next paragraph. This objective follows a phrase-level contrastive learning objectives in prior work (Lee et al., 2021; Ram et al., 2021; Deng et al., 2021; Kulkarni et al., 2022) with an extension that allows multiple positives.
In-batch positives and negatives. The start positives and the end positives are the start and the end of the spans to be retrieved. The start negatives and the end negatives are tokens that are not the start positives and not the end positives, respectively. More formally: Y+s (git) = { xjm|git = xjm...xjm+|git|−1 & i ̸= j } , Y−s (git) = { xjm|git ̸= xjm...xjm+|git|−1 & i ̸= j } , Y+e (git) = { xjm|git = xjm−|git|+1...x j m & i ̸= j } , Y−e (git) = { xjm|git ̸= xjm−|git|+1...x j m & i ̸= j } .
Here, |git| indicates the length of the span git.

heading: 4 Training Details
text: Training data. We use English Wikipedia (August 2019) and an English portion of CC-News (Mackenzie et al. (2020), February 2019) for training, which contains 13B tokens in total. The data is segmented into sequences, each with up to 256 tokens.
2The unmasked sequence and the masked sequence may have different lengths before padding, but we pad them to have the same length.
Training. We use the model architecture and initial weights of RoBERTa large (Liu et al., 2019), consisting of 354M parameters. Training is done for 100,000 steps, using thirty-two 32GB GPUs. One batch consists of 512 sequences (131,072 tokens). We use an Adam optimizer (Kingma and Ba, 2014) with a learning rate of 3× 10−5, weight decay of 0.01 and 4, 000 steps of warm-up.
Batching. The choice of batching is important in in-batch approximations, as it determines the quality of positives and negatives. For instance, Zhong et al. (2022b) uses BM25 to ensure the sequences in the same batch are likely to share the same topic. With a pretraining corpus with billions of tokens, it can be significantly expensive to build a BM25 index. Therefore, we instead construct the batch by grouping sequences from the same document and assigning them to the same batch.3 This trick ensures that (a) positives (spans that share the string) are likely to share the context, reducing false positives, and (b) negatives are those that the model is likely to be confused with, thus training against them helps the model better identify positives. During training, we gather all sequences from multiple GPUs to increase the size of the effective batch and make in-batch approximation more effective.

heading: 5 Experiments: Closed-set Tasks
text: We perform zero-shot evaluation on closed-set tasks where a small set of candidates is given.

heading: 5.1 Evaluation Datasets
text: We include nine classification datasets that are known for not necessarily requiring factual knowledge: AGNews (Zhang et al., 2015), Yahoo (Zhang et al., 2015), Subj (Pang and Lee, 2004), SST2 (Socher et al., 2013), MR (Pang and Lee, 2004), Rotten Tomatoes (RT), CR (Hu and Liu, 2004), Amazon polarity (Amz, McAuley and Leskovec (2013)) and RTE (Dagan et al., 2005). The tasks range from topic classification and sentiment analysis to subjectivity classification and textual entailment. Statistics are provided in Appendix B.

heading: 5.2 Baselines
text: We compare with the encoder-only, the decoderonly and the encoder-decoder models with various sizes (354M to 175B parameters). We include RoBERTa (Liu et al., 2019) as the encoder-only,
3Documents that are not long enough to construct a batch are grouped with each other.
T5 (Raffel et al., 2020) as the encoder-decoder, and GPT-2/3 (Radford et al., 2019; Brown et al., 2020) as the decoder-only model. For the decoderonly models, we additionally apply PMI (Holtzman et al., 2021) for better calibration of the model output. We also compare with Shi et al. (2022) who use kNN inference using GPT-2 with PMI. In particular, (1) GPT-2 kNN uses kNN inference without training, and (2) GPT-2 kNN-LM interpolates distributions from GPT-2 and GPT-2 kNN.

heading: 5.3 Setup
text: We use the templates and verbalizers from Shi et al. (2022) for all models. When available, we use fuzzy verbalizers from Shi et al. (2022). We use a domain-specific reference corpus: a union of the English Wikipedia and CC News for AGN, Yahoo and RTE, a subjectivity corpus for Subj, and a review corpus for sentiment classification datasets. Their sizes vary from 15M tokens to 126M tokens. Details are in Appendix B. Fast similarity search is done using FAISS (Johnson et al., 2019) with the HNSW index. We use k = 4096 for inference.

heading: 5.4 Results
text: NPM outperforms baselines in the zero-shot setting (Table 1). We discuss the results in detail below.
Comparison between baselines. Among parametric models, RoBERTa achieves the best performance, outperforming larger models including
Sentiment analysis: {positive, negative} with fuzzy verbalizers
GPT-3. This is perhaps surprising, and is likely because bidirectionality of the encoder-only model plays a vital role, as claimed in Patel et al. (2022). The kNN-LM approach from Shi et al. (2022), which incorporates the nonparametric component to the parametric model, outperforms all other baselines. Nonetheless, solely relying on retrieval (kNN) performs poorly with GPT-2, suggesting that using kNN at inference only is limited.
Baselines versus NPM. NPM significantly outperforms all baselines, achieving consistently competitive performance over all datasets. This indicates that, even for tasks that do not explicitly require external knowledge, nonparametric models are very competitive.
Qualitative analysis. Figure 5 depicts predictions from RoBERTa and NPM on a sentiment
100 101 102 Model size (relative to 354M)
5
10
15
20
25
30 35 M ac ro -a ve ra ge d EM
T5
BM25+T5
OPT
BM25+OPT
GPT-3
BM25+GPT-3
NPM T-REx
100 101 102 Model size (relative to 354M)
5
10
15
20
25
30
M ac
ro -a
ve ra
ge d
EM
T5
BM25+T5
OPT BM25+OPT
GPT-3
BM25+GPT-3NPM T-REx UHN
100 101 102 Model size (relative to 354M)
5
10
15
20
25
30
M ac
ro -a
ve ra
ge d
EM
T5
BM25+T5
OPT
BM25+OPT
GPT-3
BM25+GPT-3
NPM T-REx Hard
100 101 102 Model size (relative to 354M)
0
5
10
15
20
25
M ac
ro -a
ve ra
ge d
EM
T5
BM25+T5
OPT
BM25+OPT
GPT-3
BM25+GPT-3
NPM Google RE
100 101 102 Model size (relative to 354M)
0
5
10
15
20
M ac
ro -a
ve ra
ge d
EM
T5
BM25+T5
OPT
BM25+OPT
GPT-3
BM25+GPT-3
NPM Google RE UHN
100 101 102 Model size (relative to 354M)
2.5
5.0
7.5
10.0
12.5
15.0
EM
T5
BM25+T5
OPT
BM25+OPT
GPT-3
BM25+GPT-3
NPM KAMEL
100 101 102 Model size (relative to 354M)
10
20
30
40
50
EM
T5
BM25+T5
OPT
BM25+OPT GPT-3
BM25+GPT-3
NPM
TQA
100 101 102 Model size (relative to 354M)
0
2
4
6
8
10
EM T5
BM25+T5
OPT
BM25+OPT
GPT-3
BM25+GPT-3
NPM NQ
Figure 6: Zero-shot results on knowledge tasks. The x-axis indicates the relative number of model parameters in log scale compared to RoBERTa large (354M). NPM outperforms significantly larger parameters models, either with or without BM25. See Table 8 in Appendix C for the raw numbers.
analysis task. The first example uses cheap to indicate inexpensive, and the second example uses cheap to indicate of very poor quality. RoBERTa predicts Positive to both, while NPM makes correct predictions by retrieving the context that uses cheap in the same context as the input.
We also find that representations from NPM lead to better word sense disambiguation. For instance, RoBERTa assigns a high similarity score between cheap (inexpensive) and cheap (of very poor quality). On the other hand, NPM successfully assigns a low similarity score between cheap and cheap, even though their surface forms are the same.

heading: 6 Experiments: Open-set Tasks
text: We include zero-shot evaluation on open-set tasks whose answer can be any arbitrary-length string.

heading: 6.1 Evaluation Datasets
text: We evaluate on seven datasets: T-REx and Google-RE from LAMA (Petroni et al., 2019), KAMEL (Kalo and Fichtel, 2022), Natural Questions (NQ, Kwiatkowski et al. (2019)), TriviaQA
(TQA, Joshi et al. (2017)), TempLAMA2219 and an entity translation task. In particular, TempLAMA requires probing knowledge with temporal updates, motivated by Dhingra et al. (2022) and Jang et al. (2022). The entity translation task involves a translation of an entity from English to other, non-Latin languages, requiring the model to predict extremely rare (if not unseen) characters. See Appendix B for details and statistics of all datasets.

heading: 6.2 Baselines
text: We compare with T5 (Raffel et al., 2020) as the encoder-decoder, and GPT-3 (Brown et al., 2020) and OPT (Zhang et al., 2022) as the decoder-only models. The encoder-only models are not applicable for open-set tasks since the number of tokens to predict is unknown.
Prior work found that a “retrieve-and-generate” approach that concatenates the input and passages from an off-the-shelf retrieval system is often helpful in knowledge-dependent tasks (Kandpal et al., 2022). We add them as baselines, using up to five passages from BM25 (Robertson et al., 2009).

heading: 6.3 Setup
text: For all datasets, we report Exact Match (EM). The LAMA test data is biased toward frequent entities because they are filtered to only include answers that are single tokens based on BERT (Devlin et al., 2019). Since we do not want our evaluation to be biased toward overly frequent entities, we report a micro-averaged accuracy over the data whose answers are 1, 2, 3 and 4+ grams, respectively. Other datasets do not have such filtering, therefore we report average EM.
As a reference corpus, we use the English Wikipedia from 08/01/2019, consisting of 810M tokens. For TempLAMA2219 , we use the English Wikipedia from 08/01/2022, consisting of 858M tokens.
For NPM, we find combining with sparse retrieval significantly helps, likely because dense retrieval and sparse retrieval capture complementary features (Karpukhin et al., 2020; Seo et al., 2019). In particular, we reduce the search space to the top 3 passages based on BM25 and perform dense search as done in Kassner and Schütze (2020).

heading: 6.4 Results
text: Figure 6 show results on five knowledge tasks.
First, performance of parametric models largely depends on the number of parameters, as it has been claimed in much of prior work (Brown et al., 2020; Kandpal et al., 2022). The retrieve-and-generate approach that combines parametric models with BM25 significantly improves performance.
NPM outperforms or is on par with significantly larger baselines across all datasets. It substantially outperforms all models on two LAMA datasets, including 500x larger GPT-3 either with or without BM25. On KML, TQA and NQ, NPM consistently outperforms 37x larger models with or
without BM25. This is impressive given that NPM is not trained on data with questions.
It is also worth noting that sparse retrieval is critical in NPM, e.g., without sparse retrieval, performance on LAMA-TREx drops from 34.5 to 16.1. We think this is because (1) sparse retrieval and dense retrieval capture complementary features, and (2) the removal of approximation in search improves search quality. We think future work can explore completely removing sparse retrieval, as has been done in Lee et al. (2021) to improve Seo et al. (2019).
Impact of the reference corpus size. Figure 7 reports the impact of the size of the reference corpus, from 41M tokens (5%) to 810M tokens (100%). Performance of NPM is highly correlated with the size of the reference corpus, strongly suggesting that using a larger reference corpus is important.
Results on temporal knowledge tasks. Table 2 reports results on TempLAMA. NPM retains its performance on the unchanged set (18.9 →19.5) and successfully updates its answers on the changed set (2.9 → 17.5). Its performance is significantly better than the performance of parametric models with up to 13B parameters, and is on par with a larger model with the retrieve-and-generate approach, which also successfully updates its answer by leveraging the updated corpus. This is in agreement with prior work that shows the model with a nonparametric component adapts to temporal updates by replacing the reference corpus at test time (Izacard et al., 2022). Nonetheless, the retrieve-and-generate approach is still significantly worse than NPM when the target entities are rare, which we show in the next paragraph.
Performance on rare entities. We break down the instances on LAMA and TempLAMA based on the number of BPE splits of the target entity, e.g., Thessaloniki is one word that is split into 4 BPE tokens, thus the number of splits is 3. Since BPE splits a word if they are rare, the number of BPE splits indicates the rarity of the entity. We compare NPM with GPT-3 and BM25+GPT-3 on LAMA, and BM25+T5 (770M and 3B) on TempLAMA, the two most competitive baselines on each dataset.
Figure 8 reports results. On LAMA, NPM outperforms GPT-3 fairly consistently, with larger gains as the number of BPE splits increases. On TempLAMA, while BM25+T5 is competitive on frequent entities with zero BPE split, it consistently lags behind NPM with ≥ 1 BPE splits. This suggests that NPM is particularly good at addressing rare entities, compared to not only parametric models without retrieval but also the retrieve-andgenerate approach.
Results in Entity Translation. Results on the entity translation task are shown in Table 3 (perlanguage results are reported in Table 10 of Appendix C). T5 and OPT struggle to perform the task, both with and without BM25 retrieval. In contrast, NPM performs well across all languages.
In order to better calibrate performance of NPM, we provide reference performance of models that are purposely trained on the multilingual data— mT5 (Xue et al., 2021) and BLOOM (Scao et al., 2022). NPM outperforms 3.4x larger mT5 and 20x larger BLOOM, and approaches 11x larger mT5, even though it is trained on English. We think strong cross-lingual transferability of NPM is likely because it can retrieve a phrase based on its surrounding context, even if it has not seen the exact word during training.

heading: 7 Conclusion
text: We introduced NPM, a nonparametric masked language model that replaces a softmax over the output vocabulary with a nonparametric distribution over a reference corpus. NPM can be efficiently trained using a contrastive objective and an in-batch approximation to a full corpus. Zero-shot evaluation on 16 tasks shows that NPM outperforms significantly larger parametric models. NPM is particularly good at rare patterns (word senses or facts), scaling and updating at test time, and predicting extremely rare if not unseen characters.
Limitation
Scaling through the inference corpus. The size of the reference corpus is an additional dimension for model scale in nonparametric models. In this paper, we scale the corpus up to nearly 1B tokens, which is still smaller than the training data of very large language models (Brown et al., 2020; Rae et al., 2021). We think future work can scale it fur-
ther using tools such as Distributed FAISS (Johnson et al., 2019) or ScaNN (Guo et al., 2020).
Significant memory usage. Using NPM saves GPU compute and memory compared to using models with more parameters. However, NPM requires more RAM and disk memory due to embeddings of a reference corpus. For instance, the largest corpus in our experiments (full English Wikipedia) requires 70GB of RAM and 1.4TB of disk memory. Future work can build more efficient NPM as done in prior work in nearest neighbor search (Jegou et al., 2010; Norouzi et al., 2012; Ge et al., 2014; Izacard et al., 2020; Yamada et al., 2021).
Exploration of larger vocabulary. Large vocabulary is known to lead performance gains (Conneau et al., 2020) but is bounded in memory costs. Previous work explored more efficient softmax approximations (Morin and Bengio, 2005; Chen et al., 2016; Grave et al., 2017). Our nonparametric training offers an alternative by removing the softmax over the vocabulary. With the RoBERTa architecture, increasing the vocab size by 2x makes the baseline training 50% more memory expensive, but does not increase the memory in training NPM. However, this paper does not include more systematic evaluation on the effect of large vocabulary. Future work can explore training NPM with a significantly larger vocabulary to further boost performance.
Extension for generation. Our paper evaluates NPM only on prediction tasks. It is currently nontrivial to use NPM for generation, since it is the encoder-only model. Future work can explore autoregressive generation as done in Patel et al. (2022) or use NPM for editing (Schick et al., 2022; Gao et al., 2022).
Extension to few-shot learning and fine-tuning. Our paper focuses on zero-shot evaluation only. Future work can extend NPM to a few-shot learning setup. In fact, fine-tuning NPM is significantly easier than fine-tuning larger models such as T5, OPT and GPT-3 which we compare NPM with, and can be explored in future work.
Better cross-lingual transfer. Our work explored cross-lingual transfer in a limited setup where the model is trained on monolingual data. We think future work can train multilingual NPM, and explore more comprehensive cross-lingual evaluation. In fact, nonparametric training may alleviate the burden of collecting large-scale multilingual
corpora since it makes the model less sensitive to the language coverage in the training data, and may lead to significantly better cross-lingual transfer, as we demonstrate in the entity translation task.
Limitation in speed. We find that search makes inference considerably slower than the counterpart without search. We think that (1) search can significantly be faster with better engineering (we use the default hyperparameters of the FAISS index with no tuning) or better index, and (2) the speed of NPM is still on par with the speed of significantly larger parametric models that NPM outperforms (see Table 4). Moreover, while not explored in this work, there has been work that improves inference speed (He et al., 2021; Alon et al., 2022) that can be applied to NPM. We leave improving inference speed to future work.

heading: Acknowledgements
text: We thank Ari Holtzman, Eric Wallace, Iz Beltagy, Jinhyuk Lee, Jungsoo Park, Mark Johnson, Noah Smith, Ofir Press, Patrick Lewis, Xiang Deng, Xinxi Lyu, Zexuan Zhong, UW-NLP members and anonymous reviewers for discussion and comments on the paper. This research was supported by NSF IIS-2044660, ONR N00014-18-1-2826, ONR MURI N00014- 18-1-2670, an Allen Distinguished Award and gifts from AI2. SM is supported by a J.P. Morgan fellowship.

heading: A Model Details
text: A.1 Details of NPM Approximation at inference. Given qstart and qend, we take the top k tokens with the highest similarity scores with each of them, and compute scores over spans composed by these tokens. Let c∗i:j be a span in C from the i-th token to the j-th token, and E(c) ∈ Rh be a vector corresponding to a token c ∈ C. We find the top k tokens for the start and the end:
cs1 , cs2 , · · · , csk = argTopk c∈C sim(qstart,E(c)), ce1 , ce2 , · · · , cek = argTopk c∈C sim(qend,E(c))
using a fast nearest neighbor search. We then define a set of candidate phrases C̃∗ as:   k⋃
i=1
lmax⋃
j=1
c∗si:si+j−1
  ∪   k⋃
i=1
lmax⋃
j=1
c∗ei−j+1:ei
  ,
and predict:
argmax v∗∈V∗
∑
c∗∈C̃∗ I[v∗ = c∗]expsim(q,E(c∗)),
where E(c∗) ∈ R2h is a vector corresponding to c∗, and V∗ is a set of any possible n-grams defined by the vocabulary V .
A.2 Training Details All implementation was done with PyTorch (Paszke et al., 2019), PyTorch Lightning4 and Huggingface Transformers (Wolf et al., 2020).
Masking. We use a masking ratio of 15% for all models, following the standard in prior work (Devlin et al., 2019; Liu et al., 2019; Joshi et al., 2020). We implement masking as follows: (1) we first identify all possible candidate spans (spans that positives are found from other sequences in the batch), (2) sample the length of spans to mask from a geometric distribution with a hyperparameter p = 0.5, and (3) mask the spans with respect to the sampled length until the masking budget has been spent. We do not mask more than 128 spans from one sequence, and do not mask the span if the same span has been masked for more than ten times within the batch in order to prevent repeatedly masking overly frequent spans.
4https://github.com/Lightning-AI/ lightning
For [MASKs] and [MASKe], we use the [MASK] vocab from the RoBERTa tokenizer. Note that it is not necessary to use different tokens for [MASKs] and [MASKe] since the Transformer can handle positional information.
A.3 A special case: NPM SINGLE
Along with NPM, we introduce NPM SINGLE, which outputs a nonparametric distribution over every single token in C, instead of a phrase. To some extent, NPM is a strict generalization of NPM SINGLE, and NPM SINGLE still has a problem that existing encoder-only models have, e.g., can only fill in the [MASK]with a single token. We however think NPM SINGLE can be useful for some applications, e.g., for fine-tuning, as existing encoder-only models are used for.
Inference. Given a reference corpus C = {c1, · · · , cN}, we construct N number of hdimensional vectors c1, · · · , cN ∈ Rh by feeding the text into the encoder. At inference time, given a query whose t-th token is [MASK], we feed it into the encoder:
q1..qL = Encoder(q1..qt−1,[MASK], qt+1..qL).
We take qt as a vector that represents the [MASK] token in the query. Finally, the prediction is made by aggregating the similarity scores to the tokens in C:
argmax v∈V
∑ c∈C I[c = v]exp(sim(qt,E(c))),
where E(c) ∈ Rh is a vector corresponding to c, and V is the vocabulary set.
In practice, since computing scores over all tokens in C is infeasible, an approximation is made by computing scores for the top k nearest neighbors only, and treating other tokens to have a similarity score of −Inf . More precisely:
c1, c2, · · · , ck = argTopk c∈C sim(qt,E(c))
are obtained by using an index (e.g., FAISS (Johnson et al., 2019)), and the following is returned as a prediction:
argmax v∈V
∑
1≤i≤k I[ci = v]exp(sim(qt,E(ci))).
Training. Let xi1...xiL be the i-th sequence in the batch, whose subset is replaced with [MASK] and converted to x̂i1...x̂ i L. Both the unmasked sequence and the masked sequence are fed into the encoder, and each token is mapped into an h-dimensional vector:
xi1 · · ·xiL = Encoder(xi1 · · ·xiL), x̂i1 · · · x̂iL = Encoder(x̂i1 · · · x̂iL).
The training objective is then defined as:
L∑
t=1
I[x̂t = [MASK]]l(xit, x̂it),
where l(xit, x̂ i t) is
−log ∑ y∈Y+(xit) exp(sim(x̂ i t,y))∑
y∈Y+(xit)∪Y−(xit) exp(sim(x̂ i t,y))
.
Here, sim(·, ·) is a similarity function defined in Section 3.1, and Y+(xit) and Y−(xit) are positives and negatives of xit—tokens from other sequences in the batch that share and do not the vocab, respectively.
Y+(xit) = { xjm|xit = xjm and i ̸= j } , Y−(xit) = { xjm|xit ̸= xjm and i ̸= j } .
A.4 Inference on closed-set tasks When applying NPM and NPM SINGLE on closedsetk tasks, we closely follow Shi et al. (2022) who adapts kNN-LM for zero-shot inference on classification tasks. We assume a fuzzy verbalizer: f : Y → Ṽ , where Y is a set of labels in the task and Ṽ ∈ V is a subset of the vocabulary V . The fuzzy verbalizer maps a label to a set of tokens that express the label, e.g., in a sentiment classification task, f(Positive) includes awesome or great, and f(Negative) includes terrible or broken.
NPM SINGLE is given a query vector q ∈ Rh and predicts:
argmax y∈Y
∑ c∈C I[c ∈ f(y)]exp
( sim(q,E(c))
τ
) ,
where E(c) ∈ Rh is a vector corresponding to c, and τ is a hyperparameter.
NPM is given a query vector q ∈ R2h and predicts:
argmax y∈Y
∑
c∗∈C∗ I[c∗ ∈ f(y)]exp
( sim(q,E(c∗))
τ
) ,
where E(c∗) ∈ R2h is a vector corresponding to c∗. Note that this is essentially equivalent to
argmax y∈Y
∑ c∈C I[c ∈ f(y)]exp
(
sim(qstart,E(c))
τ +
sim(qend,E(c))
τ
) .
We use τ = 5.0 for both NPM SINGLE and NPM.

heading: B Evaluation Details
text: Table 5 reports statistics and templates on each downstream task, and Table 6 reports statistics of the retrieval corpus used in experiments.
For closed-set tasks, we use templates and verbalizers provided by Shi et al. (2022) for most datasets, except two datasets. For RTE, we use the template from Artetxe et al. (2022). For Subj, we write our own template, motivated by Zhong et al. (2022a) that found Subj is mainly about differentiating a review and a summary. For open-set tasks, we use templates provided by the original authors, except NQ and TQA for which we use the templates from GPT-3 (Brown et al., 2020). Due to limited computation resource, we subsample the data to include up to 3,000 examples, following the standard from prior work (Zhao et al., 2021; Shi et al., 2022). For closed-set tasks, we use exactly the same set of data as Shi et al. (2022), and for open-set tasks, we use the same script to subsample the data. For LAMA T-REx and Google RE, we subsample up to 1,000 examples for each of 1, 2, 3 and 4+ grams. For the entity translation task, we subsample up to 1,000 examples per language.
The following is a more detailed description of open-set tasks used in Section 6.
LAMA (Petroni et al., 2019) is a factual probing benchmark that is designed to quantify the amount of factual knowledge in the model. It requires the model to predict the object given a subjectrelation tuple in a cloze format. We use two versions of LAMA (Petroni et al., 2019): (1) LAMA T-REx, derived from Elsahar et al. (2018) and (2) LAMA Google-RE, derived from the Google-RE corpus.5 For each version, we additionally consider the UHN (UnHelpfulNames) subset (Poerner et al., 2019)) where instances whose subject strongly hints the object by names (e.g., Apple Watch
5https://code.google.com/archive/p/ relation-extraction-corpus
and Apple) are excluded. We also consider the hard subset of T-Rex from Zhong et al. (2021).
Note that Petroni et al. (2019) only include triples whose object is one token based on BERT (Devlin et al., 2019); however, with a different pretrained model like RoBERTa, entities could be multiple BPE tokens. Entities that are splitted into multiple BPE tokens are more rare entities.
KAMEL (Kalo and Fichtel, 2022) is another factual probing task as LAMA but with a few key differences to make it more general and broad: (1) it includes a broader coverage of triples, (2) it removes the constraint that the object is one token based on BERT, (3) it includes objects with literal values, and (4) it has a question answering format.
Natural Questions (NQ, Kwiatkowski et al. (2019)) and TriviaQA (TQA, Joshi et al. (2017)) are two welll-studied open-domain question an-
swering datasets. We use the open-version of NQ (Lee et al., 2019) and TQA where the question is the only input and the model should use its knowledge to answer the question.
TempLAMA2219 is a task that requires probing knowledge with temporal updates. The task is first introduced by Dhingra et al. (2022) and Jang et al. (2022); however, we could not use either of existing data as their time split do not match our training. We therefore create the data by using a script provided by Dhingra et al. (2022) but using the 2019 and the 2022 dumps. We take Wikipedia triples whose relations are available for a template from either Petroni et al. (2019) or Dhingra et al. (2022). We then include triples whose object entities differ between the 2019 dump and the 2022 dump (due to the entity being updated), or only appear in the 2022 dump (due to the subject or the relation being
added) to the changed set. Otherwise, triples are included in the unchanged set. We additionally find that many triples are overly difficult because the fact is extremely niche and not really known. We thus filter the data to only include facts that appear in Wikipedia. Specifically, we include triples if the subject has a corresponding Wikipedia page and the object entity appears in that Wikipedia page.
Entity translation requires translating an entity from English to other languages that are not Latin based. While this is mainly to evaluate if the model can generate rare or unseen characters that are not in English, the entity translation task itself is a vital and challenging task in real applications such as machine translation (Babych and Hartley, 2003; Yan et al., 2018) and cross-lingual question answering (Clark et al., 2020; Asai et al., 2021). It is often beyond a series of simple translations of each word, or spelling out its pronunciation (Moore, 2003; Hassan et al., 2007; Sun et al., 2017). For instance, the Korean translation of Banpo Bridge in Figure 1 (반 포대교) is not the concatenation of the translations of Banpo and Bridge (반포다리).
We first identify a list of 15 non-Latin languages: Arabic (ar), Czech (cs), Greek (el), Hindi (hi), Hebrew (iw), Japanese (jp), Korean (ko), Malayalam (ml), Mongolian (mn), Polish (pl), Russian (ru), Tamil (ta), Thai (th), Turkish (tr), and Chinese (zh). We then implement heuristics to identify entities and their translations from English Wikipedia. Specifically, we parse the first paragraph of each Wikipedia article and pair the found translation with a topic en-
tity of the article. For instance, a Korean translation of Banpo Bridge is found from the first sentence of https://en.wikipedia.org/ wiki/Banpo_Bridge. Per-language statistics are reported in Table 7.

heading: C Additional Results
text: Full results on knowledge tasks. Table 8 reports full results on five knowledge tasks. See Figure 6 for an illustration, and Section 6.4 for discussion.
Comparison to few-shot GPT-3. Table 9 compares zero-shot NPM SINGLE and NPM with zeroand four-shot GPT-3. Our zero-shot models outperform 500x larger zero-shot GPT-3 and 7.6x larger 4-shot GPT-3, but lag behind 4-shot GPT-3 that is 19x or larger. We think future work can explore extending our models to a few-shot setup.
Additional qualitative results. Figure 9 depicts predictions from RoBERTa and NPM in topic classification, choosing a label between four candidates: health, computer, travel and politics. All three examples contain the word torch, but with different meanings, e.g., an infectious diseases, a tool, and a computer library. RoBERTa predicts health for all of them, while NPM predicts health, travel and computer, which are all correct predictions.
As in Figure 5, we find that representations from NPM enable better word sense disambiguation: the pairwise similarities between between different meanings of torch are significantly lower than the pairwise similarities between other tokens that share the meaning.
Entity translation given an oracle passage. We evaluate models on the entity translation task where an oracle passage—a passage that is guaranteed to contain the translation information—is provided to the model. Baselines prepend oracle passages to the input, as it does with the retrieve-and-generate approach. NPM uses oracle passages to restrict the search space.
Table 11 reports results. While performance overall increases compared to when the oracle passage is not provided, the overall comparison between models does not change from Table 10: (1) all monolingual models significantly suffer, except for a couple of languages that are derived from Latin; (2) NPM significantly outperforms all monolingual models; (3) NPM even outperforms 3.4x larger mT5 and 20x larger BLOOM, and approaches 11x larger mT5.
ACL 2023 Responsible NLP Checklist

heading: A For every submission:
text: 3 A1. Did you describe the limitations of your work?
Section following Section 6 (No section number).
3 A2. Did you discuss any potential risks of your work? Appendix E.
3 A3. Do the abstract and introduction summarize the paper’s main claims? Abstract and Section 1.
7 A4. Have you used AI writing assistants when working on this paper? Left blank.
B 3 Did you use or create scientific artifacts? Section 3 and 4.
3 B1. Did you cite the creators of artifacts you used? We include citations to the pretraining data and pretrained models (Section 3.3 and Appendix B) and evaluation datasets (Section 4.1, 5.1 and Appendix C).
7 B2. Did you discuss the license or terms for use and / or distribution of any artifacts? The license and terms for use will be released in the open-sourced repo, which we do not include in the submission in order to keep anonymity.
3 B3. Did you discuss if your use of existing artifact(s) was consistent with their intended use, provided that it was specified? For the artifacts you create, do you specify intended use and whether that is compatible with the original access conditions (in particular, derivatives of data accessed for research purposes should not be used outside of research contexts)? We discuss how we use the evaluation datasets in detail in Section 4, 5 and Appendix C.
7 B4. Did you discuss the steps taken to check whether the data that was collected / used contains any information that names or uniquely identifies individual people or offensive content, and the steps taken to protect / anonymize it? We use standardalized, research data/benchmarks, and checked that the authors of the original papers tried their best to ensure the data does not contain names or uniquely identifies individual people or offtensive content. Nonetheless, we’ve discussed them in "Potential Risk” (Appendix E).
3 B5. Did you provide documentation of the artifacts, e.g., coverage of domains, languages, and linguistic phenomena, demographic groups represented, etc.? We reported the language of the model/data and their domains in Section 3.3, 4, 5, Appendix B and Appendix C.
3 B6. Did you report relevant statistics like the number of examples, details of train / test / dev splits, etc. for the data that you used / created? Even for commonly-used benchmark datasets, include the number of examples in train / validation / test splits, as these provide necessary context for a reader to understand experimental results. For example, small differences in accuracy on large test sets may be significant, while on small test sets they may not be. We included statistics of datasets in Appendix C.
The Responsible NLP Checklist used at ACL 2023 is adopted from NAACL 2022, with the addition of a question on AI writing assistance.
C 3 Did you run computational experiments? Left blank.
3 C1. Did you report the number of parameters in the models used, the total computational budget (e.g., GPU hours), and computing infrastructure used? We include training details and related information in Section 3.3 and Appendix B.
3 C2. Did you discuss the experimental setup, including hyperparameter search and best-found hyperparameter values? We discuss detailed experimental setups in Section 3.3, 4 and 5, Appendix B and C. We evaluated models in a zero-shot setup, thus there is no hyperparameter tuning involved.
7 C3. Did you report descriptive statistics about your results (e.g., error bars around results, summary statistics from sets of experiments), and is it transparent whether you are reporting the max, mean, etc. or just a single run? Since all experiments are about zero-shot evaluation, all results are based on a single run and are deterministic.
3 C4. If you used existing packages (e.g., for preprocessing, for normalization, or for evaluation), did you report the implementation, model, and parameter settings used (e.g., NLTK, Spacy, ROUGE, etc.)? We described and cited packages that we use in the paper. For complete reproducibility, we will also open-source the code, which we do not include in the submission in order to keep anonymity.
D 7 Did you use human annotators (e.g., crowdworkers) or research with human participants? Left blank.
D1. Did you report the full text of instructions given to participants, including e.g., screenshots, disclaimers of any risks to participants or annotators, etc.? No response.
D2. Did you report information about how you recruited (e.g., crowdsourcing platform, students) and paid participants, and discuss if such payment is adequate given the participants’ demographic (e.g., country of residence)? No response.
D3. Did you discuss whether and how consent was obtained from people whose data you’re using/curating? For example, if you collected data via crowdsourcing, did your instructions to crowdworkers explain how the data would be used? No response.
D4. Was the data collection protocol approved (or determined exempt) by an ethics review board? No response.
D5. Did you report the basic demographic and geographic characteristics of the annotator population that is the source of the data? No response.

heading: Future Work
text: We think future work can explore extending our models to a few-shot setup. Additional qualitative results. Figure 9 depicts predictions from RoBERTa and NPM in topic classification, choosing a label between four candidates: health, computer, travel and politics. All three examples contain the word torch, but with different meanings, e. g. , an infectious diseases, a tool, and a computer library. RoBERTa predicts health for all of them, while NPM predicts health, travel and computer, which are all correct predictions. As in Figure 5, we find that representations from NPM enable better word sense disambiguation: the pairwise similarities between between different meanings of torch are significantly lower than the pairwise similarities between other tokens that share the meaning. Entity translation given an oracle passage. We evaluate models on the entity translation task where an oracle passage—a passage that is guaranteed to contain the translation information—is provided to the model. Baselines prepend oracle passages to the input, as it does with the retrieve-and-generate approach. NPM uses oracle passages to restrict the search space. Table 11 reports results. While performance overall increases compared to when the oracle passage is not provided, the overall comparison between models does not change from Table 10: (1) all monolingual models significantly suffer, except for a couple of languages that are derived from Latin; (2) NPM significantly outperforms all monolingual models; (3) NPM even outperforms 3. 4x larger mT5 and 20x larger BLOOM, and approaches 11x larger mT5. ACL 2023 Responsible NLP Checklist