Nearest Neighbor Speculative Decoding for LLM Generation and Attribution

Authors: minghan  li
xilun  chen
ari  holtzman
beidi  chen
jimmy  lin
wen-tau  yih
victoria  xi
  lin
  cohere
meta  fair
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 07:25:03

Abstract: Large language models (LLMs) often hallucinate and lack the ability to provide attribution for their generations. Semi-parametric LMs, such as kNN-LM, approach these limitations by refining the output of an LM for a given prompt using its nearest neighbor matches in a non-parametric data store. However, these models often exhibit slow inference speeds and produce non-fluent texts. In this paper, we introduce Nearest Neighbor Speculative Decoding (NEST), a novel semi-parametric language modeling approach that is capable of incorporating real-world text spans of arbitrary length into the LM generations and providing attribution to their sources. NEST performs token-level retrieval at each inference step to compute a semi-parametric mixture distribution and identify promising span continuations in a corpus. It then uses an approximate speculative decoding procedure that accepts a prefix of the retrieved span or generates a new token. NEST significantly enhances the generation quality and attribution rate of the base LM across a variety of knowledge-intensive tasks, surpassing the conventional kNN-LM method and performing competitively with in-context retrieval augmentation. In addition, NEST substantially improves the generation speed, achieving a 1.8x speedup in inference time when applied to Llama-2-Chat 70B. Code will be released at https://github.com/facebookresearch/NEST/tree/main.

Main paper: 
 introduction: 
 Large language models (LLMs) have demonstrated strong potential as multi-task solvers, excelling in a wide range of applications (Brown et al., 2020; Chowdhery et al., 2022; Touvron et al., 2023a; Anil et al., 2024) . Despite their advanced capabilities, LLMs frequently encounter the problem of hallucination, particularly when dealing with long-tail knowledge that is less represented in their training data (Kandpal et al., 2023; Asai et al., 2023a) . To address this limitation, retrieval augmentation incorporates information retrieval and nearest neighbour search from a non-parametric data store to enhance evidence-based and situated reasoning with LLMs. The resulting semi-parametric LMs exhibit a reduced tendency to generate unsupported content (Khandelwal et al., 2020; Borgeaud et al., 2022; Shi et al., 2024a,b; Asai et al., 2023a) .However, the effectiveness of retrieval-augmented language models (RALMs) in ensuring accurate and reliable content generation varies. The widely used in-context retrieval-augmentation (RA) regime (Ram et al., 2023; Shi et al., 2024a,b) softly biases the LM output distribution by prepending retrieved content to the input, which does not reliably guarantee faithful attribution of information. Approaches such as kNN-LM (Khandelwal et al., 2020) modify the LM output with a non-parametric 2 Background 
problem definition: 
 Given an input x, a mixture model M predicts the output y consisting of segments {y 1 , y 2 , ..., y T }. In our setting, M may produce multiple tokens at a time step t, and therefore y t indicates the t-th segment consisting of at most n tokens where 1 ≤ |y t | ≤ n. Let {w (1) t , w(2) t , ..., w (n) t } be the tokens in segment y t , we use p M (w|x, y <t ) to denote the distribution of the next token, and use p M (w = w(1) t |x, y <t ) to denote the probability of w(1) t of the next segment y t .The mixture model M involves a pre-trained LM and key-value datastore (K, V) that enables approximate nearest neighbors search without further training or fine-tuning.Key-value datastore. To create the datastore (K, V) using the LM for a corpus D, let f (•) be the mapping from input sequence c to the hidden states h of the LM at some fixed layer. Let w be the next word of c in the corpus D. For a sample (c i , w i ) in D after segmentation, we define the i-th key-value pair (k i , v i ) in (K, V) as (h i , w i ), where h i = f (c i ). The whole datastore is thus defined Figure 1 : The NEST approach first locates the tokens in the corpus using the LM hidden states. The retrieval distribution p k-NN is dynamically interpolated with p LM based on the retriever's uncertainty λ t . The token and its n-gram continuation are then selected from the mixture distribution p M , while the final span length is determined by speculative decoding to remove undesired tokens. The spans incorporated in the final generation provide direct attribution and amortize the generation latency. as the set of all possible key-value pairs in D:(K, V) = {(h i , w i )|(c i , w i ) ∈ D}.(1)The size of the datastore (K, V) is proportional to the total number of tokens in corpus D. This brings difficulty in scaling the size of the corpus and the model, which may require massive storage space and computational resources.Probability interpolation. During inference, the language model outputs the token distribution p LM (w|x, y <t ), together with the hidden state q t . The model uses q t as a query to search the datastore (K, V) and retrieve the r-nearest neighbors π according to the similarity s(q, k) between a query q and a key k. The final non-parametric distribution p k-NN (w|x, y <t ) is computed using a softmax function over the similarity of all retrieved neighbors: p k-NN (w|x, y <t ) ∝ (ki,vi)∈πEQUATIONwhere µ is the inverse temperature. We use 1/ dim(q t ) for µ in practice where dim(q t ) is the hidden state dimension. This is similar to computing attention in the Transformer model (Vaswani et al., 2017) . For similarity s(q, k), we follow Khandelwal et al. (2020) and use the negative squared ℓ 2 distance. Items not in π are assigned with 0 probability based on the indicator function I w=vi .Finally, the next token is sampled from the mixture distribution p M of the non-parametric distribution p k-NN and the parametric distribution p LM using a fixed hyper-parameter λ ∈ [0, 1]:EQUATION3 Nearest Neighbor Speculative Decoding 
two-stage k-nn search: 
 As mentioned in Section 2.2, maintaining a token-level key-value store can be expensive in terms of both latency and storage. To provide a better trade-off between latency and accuracy, we adopt the two-stage design, which is widely applied in information retrieval and search engines.First-stage passage retrieval Given the corpus D, we segment the documents into separate passages of less than m tokens each. We then encode the corpus and use a hybrid retrieval system to select the relevant passages, as dense retrievers are good at handling semantics in queries (Karpukhin et al., 2020) and sparse retrievers are good at lexical matching (Sciavolino et al., 2021) .Second-stage k-NN token search After obtaining the top-b retrieved passages {d 1 , d 2 , ..., d b } at time step t, we use the encoder f (•) of LM to encode the prefixes of all tokens as keys as shown in Figure 1 . The key-value datastore (K, V) therefore is created on the fly. Similarly, we use the negative squared ℓ 2 distance as the similarity function and q t as the queries to search for the top-r nearest neighbors π in (K ′ , V ′ ).The two-stage design provides a trade-off between search latency and accuracy and the passage-level index only takes a fraction of the token-level index in Section 2.2. In addition, the first-stage passage search also acts as a filter to remove deceptively similar tokens in non-relevant contexts.Similar to Equation (3), we linearly interpolate the language model's distribution p LM and nonparametric distribution p k-NN using a coefficient λ t for a time step t in generation. The difference is that we use the token retrieval score to compute λ t :EQUATION)where σ is the sigmoid function and the min-max ratio expresses the uncertainty of the k-NN component. We use the sigmoid activation to re-center and re-scale this uncertainty, where α is the offset and τ is the scale for the sigmoid function. We refer to this method as Relative Retrieval Confidence (RRC).If the downstream task does not involve generation, such as perplexity evaluation and multi-choice tasks, our method will end at Equation ( 4). The mechanisms introduced in the following sections are only applied to generation, including token/span selection and post-hoc revision.Directly sampling tokens from the mixture distribution p M might escalate the exposure bias since the seemingly coherent tokens might be retrieved from completely different sources. To maintain coherence, we extend the context of the current sampled token by using its n-gram continuation in the corpus. Given the current time step t, we first select the next token w t from the mixture distribution p M . However, the sampled token w t may correspond to multiple retrieved w i (i.e., the value v i ), in the neighbors π which have different n-gram continuations. We use a simple max-pooling strategy2 to select the starting token w(1)EQUATIONThe corresponding n-gram for time step t is {wEQUATIONt , ..., wt } where n is fixed hyper-parameter. The final output is determined by the interpolation coefficient λ t in Equation (4):y t = w t , if λ t > δ; {w (1) t , w (2) t , ..., w (n) t }, otherwise. (6)where δ is a threshold and y t is the segment output at time step t.Despite the dynamic selection, the hyper-parameter n is hard to control over different tasks. To produce spans with adaptive length, we take inspiration from Leviathan et al. (2023) , where we use M to revise the proposed n-gram. However, the proposal distribution q(w|x, y <t ) is unknown besides the first token w(1) t . Therefore, we use a relaxed version of speculative decoding that upper bounds the acceptance probability. The probability of accepting the token w (i) t in a span is:EQUATIONwhere γ ∈ (0, 1] is the relaxation factor, which is referred to as "leniency" by Leviathan et al. (2023) . The smaller γ is, the less often M rejects the draft. If token w (i)t is rejected, we will remove all the tokens from w (i)t to w (n)t , and then re-select a token w (i) t from the distribution p M without going through the span selection. The computation for processing multiple tokens can be parallelized and NEST can thus maintain the latency or even accelerate the generation. Moreover, suppose all tokens in the draft are not rejected. In that case, we will directly fetch the n-gram's continuation in the corpus and use it for the next draft proposal until rejection, removing the reliance on the hyper-parameter n.Once the n-gram is accepted, the corresponding parts are masked in the corpus and will never be used again in this generation. This is to prevent the k-NN component from repetitively retrieving the same segments in a small key-value store (K ′ , V ′ ). We provide the complete procedure in Algorithm 1. 
experiments: 
 We evaluate NEST and other baselines on various tasks including text completion, question-answering, fact-verification, and multi-choice tasks, providing a comprehensive picture of factuality, fluency, and attribution of NEST in different domains. In all experiments, we focus on evaluating instructionfollowing models. We use Llama-2-chat under a zero-shot setting, where we remove the few-shot demonstrations from the instructions to simulate the realistic usage of these models.Text completion. WikiText-103 (Merity et al., 2017) is a standard benchmark for language modeling, extracted from the set of verified articles on Wikipedia. Pile of Law (Henderson et al., 2022) is a growing dataset of legal and administrative data. We use the datasets3 from Huggingface and further split the test data into validation and test sets. For language modeling, we report the perplexity score. For free-form generation, we report ROUGE-1, 2, L (Lin, 2004) and MAUVE (Pillutla et al., 2021) .Question answering. We select four knowledge-intensive question-answering datasets, including Natural Questions (NQ) (Kwiatkowski et al., 2019) , TriviaQA (TQA) (Joshi et al., 2017) , HotpotQA (HQA) (Yang et al., 2018) , and MedMCQA (MQA) (Pal et al., 2022) . Since the in-context demonstrations are removed for free-form generation, we use answer-level recall (i.e., Hit@1) (Karpukhin et al., 2020) which checks if the output contains any correct answers instead of exact match.Fact verification. We evaluate a biography-generation task (Min et al., 2023) and Truth-fulQA (Lin et al., 2022) which is a benchmark for testing false beliefs or misconceptions. We use FACTSCORE (Min et al., 2023) for biography. For TruthfulQA, we follow Lin et al. (2022) which uses the difference between the max similarity to a true reference answer and the max similarity to a false reference answer for BLEU and ROUGE-1.Closed-set tasks. MMLU (Massive Multitask Language Understanding) (Hendrycks et al., 2021) benchmark covers 57 subjects across STEM, the humanities, the social sciences, and more. We report the macro accuracy for each domain.Knowledge Sources. Wikipedia (CC BY-SA 3.0): For tasks except text completion on Pile of Law, we use the Wikipedia 2021 dump released by Izacard et al. (2024) as the knowledge source and follow the same pre-processing procedures in RA-DIT (Lin et al., 2024) , yielding ∼33M passages with each less than 200 tokens. Pile of Law (CC BY-NC-SA 4.0): We use the training split from Huggingface and select only the English data. We then follow the same procedure applied in Wikipedia, yielding a corpus containing ∼15M passages after filtering. More details are provided in Appendix A.Inference setting. kNN-LM and NEST share the same first-stage retriever. We use DRAGON+ (Lin et al., 2023) and BM25 (Robertson and Zaragoza, 2009) to encode the segments into dense and sparse vectors, respectively. Given the input, we query both the dense and sparse indexes at the same time and retrieve their corresponding top-(b • l) passages. We linearly interpolate the similarity scores between the two search results (also known as fusion) and sort them before selecting the top-b passages. The number of passage candidates b is set to be 40 and the scaling factor l is set to be 100. For RA, we use the top-3 passages in the prompt due to the context window limit. We further combine NEST and RA since they are independent methods. Greedy decoding is used during generation. More details about retrieval, decoding, and hyper-parameters are described in Appendix B.Evaluation setting. For text completion tasks and perplexity evaluation, we use 128 tokens as the prefix and the consecutive 256 tokens as the target. For the other tasks, we use 128 tokens as the max generation length for question answering and 512 for fact verification. For retrieval-based models, only the prefix will be used for retrieval. Hyper-parameters of all baselines and NEST are tuned on the dev set of WikiText-103, NQ, and Biography. Each baseline uses the same hyper-parameters for all tasks evaluated. We first tune the related hyper-parameters for perplexity and then tune the rest for generation metrics to reduce the search space. More details are provided in Appendix B.Base LMs. We evaluate publicly available, instruction-tuned language models, Llama-2-chat series4 , with model sizes ranging from 7B, 13B to 70B.Two-Stage kNN-LM. We apply the two-stage strategy described in Section 3.1 to kNN-LM as well, where we retrieve the top-b passages and encode a key-value datastore (K ′ , V ′ ) on the fly.In-Context Retrieval Augmentation (RA). A common retrieval-augmentation method is adding the retrieved evidence into the prompt. We perform retrieval given the only input instead of retrieving new passages every k step due to the expense of refreshing the kv-cache.Table 1 shows the main results of NEST and other baselines. For language modeling, RA-NEST is able to achieve the lowest complexity on both WikiText-103 and Pile of Law. For text completion, RA has the best MAUVE scores and ROUGE scores in Wikitext-103 while RA-NEST works better for 7B and 13B models on Pile of Law. We observe that for legal documents, quoting the exact clauses from the source might be more favourable compared to Wikipedia.For question-answering, RA-NEST tends to work better for smaller models (7B and 13B) in general. The gap between base LMs and other methods diminishes for 70B LMs, which is consistent with previous work where retrieval is found most useful for smaller models (Borgeaud et al., 2022) .For fact-verification, NEST is able to outperform the base LMs but underperform RA in terms of the FACTSCORE. RA-NEST is able to outperform RA for the 70B model. The degradation for RA-70B is caused by generating shorter claims which is punished by the FACTSCORE. On TruthfulQA, the semi-parametric LMs consistently outperform base LMs and RAs where in-context retrieval seems to have a negative effect on the scores. This is because TruthfulQA is an adversarial dataset containing difficult questions where in-context RA is more susceptible to the "evidence" in the prompt (e.g., astrology and myths). In contrast, NEST only interpolates the results at the output level and therefore performs better in this case. The combination RA-NEST is also affected by the in-context retrieval.For closed-set tasks, NEST is comparable to RA and RA-NEST manages to achieve the best macro scores on average. Overall, NEST is able to outperform base LMs and kNN-LM's on most tasks while being on par with RA. The combination of RA and NEST further improves over the two methods on some tasks. Despite the limited improvement, we will show that NEST is able to provide better attribution and latency in the following sections.Wikitext-103 Pile of Law PPL(↓) MAUVE RG-1 RG-2 RG-L Avg. Len PPL(↓) MAUVE RG-1 RG-2 RG-L Avg. Len Llama-2-Latency breakdown. The combination of dynamic span selection and relaxed speculative decoding can improve the latency of the LLM generation by quick draft proposal and processing multiple tokens at a time step. Figure 2a shows the latency breakdown of a NEST-70B model (α = 0.3, τ = 0.1, δ = 0.5) for different relaxation factors on the Biography validation data. The latency experiment is done on 8×A100 GPUs (for model parallelization) and 32 CPU threads (for search). The batch size is set to 1. We use internal, research-purpose implementation of the base Llama-2-chat model which did not optimize for latency. As we can see, the LM encoding time takes about half of the latency, while the sum of the others takes the rest. Noticeably, the cost of passage search and token index building stay relatively constant per query, while the others are related to the number of tokens processed per time step. Still, even with extra retrieval overheads, the slowest NEST model is faster than the base LM, showing the efficacy of span selection and speculative decoding.Latency-accuracy trade-off. To understand why NEST can accelerate generation, we first show the latency-accuracy trade-off by tuning the relaxation factor in Figure 2b . The smaller γ is, the less often NEST rejects a segment retrieved from the corpus, which enables more tokens to be processed in parallel. The average proposed span length in Figure 2a can increase from 5 tokens to 35 tokens at each time step as the relaxation factor gets smaller. Combined with Figure 2a , we can reach the conclusion that fetching longer spans from the corpus results in lower generation latency per query.For the accuracy, the FACTSCORE on Biography validation data shows that there is a sweet spot around γ = 5e -2 where both low latency and high accuracy can be achieved at the same time. Figure 2 : Latency-accuracy trade-off and breakdown on Biography using Llama-2-Chat 70B+NEST.As the relaxation factor γ decreases, NEST tends to accept longer spans from the corpus. We choose γ = 5e -2 in our main experiments, which accelerates the generation and improves the FACTSCORE. (Attribution) Length: Average length of consecutive spans in the generation that are taken from the same document. Green: Segments taken from the corpus. Gray: Reference.One of the most important features of NEST is providing attribution directly at a span level, where the reference for the corresponding statement is accurate since it is directly taken from the corpus.Table 2 shows the attribution ratio, average attributed span length, and two examples for analysis.For NQ and Biography tasks, depending on the model and hyper-parameters in Equation ( 4) and ( 7), the ratio of tokens that can be traced back to the corpus ranges from 33.2% to 95.5%. In addition, it is more desirable to have consecutive segments that come from the same source so that consistent attribution can be provided, and the average length of spans taken from the corpus ranges from 3.0 to 27.9 tokens. This feature provides span-level attribution for most claims in the LLM generation.To our knowledge, neither of the baselines can achieve the same granularity and preciseness for the attribution as NEST. We provide more analyses on sensitivity and ablation for NEST in Appendix C.5 Related Work 
retrieval-augmentation: 
 Retrieval Augmentation involves external knowledge sources to improve the effectiveness of language models on knowledge-intensive tasks. Chen et al. (2017) propose DrQA which combines extractive models and independent retrievers for open-domain question-answering. Follow-up works on retrievalaugmentation such as REALM (Guu et al., 2020) , RAG (Lewis et al., 2020) , and Atlas (Izacard et al., 2024) further combine the retrieval component in pre-training and fine-tuning for downstream knowledge-intensive tasks. Asai et al. (2024) further divide them into three categories:Input augmentation. REPLUG (Shi et al., 2024a) and in-context RALM (Ram et al., 2023) propose to pre-pend the retrieved passages in the prompts for zero-shot factual support. Recently, Self-RAG (Asai et al., 2023b) leverages special tokens to perform adaptive retrieval and different critics to iterative refine the RALM's output. RA-DIT (Lin et al., 2024) retrofits LLMs with retrieval capabilities via instruction fine-tuning.Intermediate fusion. RETRO (Borgeaud et al., 2022) employs a novel attention mechanism to incorporate multiple pre-processed text fragments in intermediate layers for more efficient integration of retrieved results. This approach has been successfully applied to larger decoder-only language models as demonstrated by RETRO++ (Wang et al., 2023b) and InstructRetro (Wang et al., 2024) .FiD (Izacard and Grave, 2021) applies similar an encoder-decoder structure in a zero-shot manner and achieves better effectiveness at a document level.Output integration. kNN-LM (Khandelwal et al., 2020) pioneers this direction and proposes to interpolate the retrieval distribution and LM's prediction. Follow-up works further propose adaptive interpolation methods which involve training (He et al., 2021; Bhardwaj et al., 2023) and excessive tuning (Drozdov et al., 2022) . Another line of work proposes to joint train the phrase encoder and LM to expand the vocabulary dynamically using the retrieved phrases, such as Copy-Generator (Lan et al., 2023) and its follow-up work (Cao et al., 2024) . Martins et al. ( 2022) proposes a chunk-based kNN machine translation model which retrieves chunks of tokens from the datastore.Speculative decoding (Leviathan et al., 2023; Chen et al., 2023; Miao et al., 2023; Spector and Re, 2023 ) is an acceleration method that leverages a small model to generate drafts for a large model to evaluate. The latency is improved as the larger model can process multiple tokens in parallel at each time step. Recently, REST (He et al., 2024) proposes to draw multiple drafts from a datastore and leverages a prefix trie tree to compute the proposal distribution, which is the closest concurrent work. Yang et al. (2023) also utilizes prefix matching to select draft sentences from a datastore, and keep the continuation of the draft sentence as long as the token matches with the model generation.In general, speculative decoding can be categorized as an unbiased self-revision method. In comparison, NEST changes the LM output distribution through interpolation with a non-parametric probability distribution. Previous work focusing on fact-checking follows a similar idea to generate factually consistent texts with a set of evidence via post-hoc editing, such as FRUIT (Iv et al., 2022) and PEER (Schick et al., 2022) . Recently, RARR (Gao et al., 2023) leverages more complex planning with LLMs to verify the retrieved evidence and generate attribution reports. 
limitations: 
 While being able to directly retrieve segments from the corpus and apply them in the generation, the output of NEST might still contain factual errors depending on the accuracy of the first-stage passage retrieval and the second-stage token retrieval. Moreover, as a plug-and-play method, our main goal is to provide a flexible solution that can combine different LLMs and data stores in zero-and few-shot manners. Without further fine-tuning, the integrated system might be sub-optimal and the results can be better if it is fine-tuned on appropriate tasks. Lastly, such semi-parametric LMs may not improve the ability of in-context learning, since the demonstrations in the prompts are unlikely to appear in any contexts that can be found in the database. An observation from preliminary experiments is that the current neural retrievers do not have the capability to process the in-context few-shot information, where techniques such as query reformulation might be needed for parsing the demonstrations.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: We provide details about datasets, inference, evaluation, hyper-parameters, and baseline in Section 4 and Appendix B. Our method is a training-free method and therefore does not involve training-related hyper-parameters such as training steps and learning rate.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
conclusion: 
 This paper presents NEST, an inference-time revision method for LMs that improve their factuality and attribution through nearest neighbor speculative decoding. Leveraging two-stage k-NN search, relative retrieval confidence, dynamic span selection, and relaxed speculative decoding, NEST improves both validation perplexity and free-form generation quality on nine different tasks. Its effectiveness can be further improved when combined with in-context retrieval augmentation. With these results, we demonstrate that NEST is capable of generating text grounded to real-world sources in low latency while maintaining fluency.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [No] Justification: This paper does not include significant tests considering the performance gap between the proposed approach and the baselines.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
broader impact: 
 The ability to copy real-world texts from existing data stores is useful for finding the source of the claim (credibility), preventing hallucination (factuality), as well as protecting copyright (risk management). It helps to resolve the dispute that often happens in AI tools by acknowledging the contents that are borrowed from existing human works (e.g., arts, books, and other creative content). Meanwhile, the information on the Internet is mixed and it is important to filter out false and sensitive information before directly injecting them into the generation.Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: Computational resources and hardware information for inference are reported in Section 4.5 and Appendix A. Our method does not involve training.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: Our paper conforms with the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: Section 8.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our method is a general framework that can be applied to different language models and knowledge sources, and therefore does not require safeguards.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: Except the internal code base, all the existing assets including code and data used in this paper are properly cited in Section 4.2 and Appendix B.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: This paper does not introduce new assets.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: This paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: This paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Two-stage k-NN search For the first-stage passage search, we use a Faiss (Douze et al., 2024) dense index and Pyserini (Lin et al., 2021) BM25 index for efficient search. For the dense index, we first use DRAGON+ to encode each passage in the corpus into a single vector, and then use Faiss (index string "IVF65536,PQ256") to cluster the vectors into 65536 centroids and quantize them into 256 codes of 8 bits each. For the sparse index, we use the default hyper-parameters and the "optimize" option in Pyserini to reduce the index size. For the approximate nearest neighbor retrieval, we use nprobe= 4096. During passage search, we retrieve 4000 passages from each index and keep the similarity score for fusion. The fusion coefficient η is determined by the relative confidence of dense and sparse retrievers similar to Equation (4). We set the dense coefficient η dense = 1top-100(s dense (q, d))/ max s dense (q, d) and same for the sparse coefficient η sparse . The final interpolation coef is η = 0.7 * (1 -η sparse ) + 0.3 * η dense . The fusion score for each document s(q, d) = η * s dense + (1 -η) * s sparse . If a document is missing in either dense or sparse retrieval results, we set its score to the minimum similarity of the dense/sparse retrieval results. The first-stage search is done on RAM and CPUs with 32 threads. The final Wikipedia dense index size is about 8.96GB, and the sparse index size is about 3.48GB on disk.For the second-stage token search, we use the LLM to encode the sequence and use the input to the final layer's feed-forward network after layer normalization as the key and query vectors following Khandelwal et al. (2020) . We retrieve the top-1024 tokens using the squared ℓ 2 distance and compute the non-parametric probability according to Equation (2).Rest of NEST For relative retrieval confidence, we set α = 0.3, τ = 0.1 for all Wikipedia-based tasks and α = 0.2, τ = 0.1 for Pile of Law for all model sizes in Equation ( 4). For dynamic span selection, we set the n-gram length to be 64 and δ = 0.5 for all model sizes and all tasks in Equation ( 6). For relaxed speculative decoding, we set γ = 5e -4 for Pile of Law tasks for all model sizes in Equation ( 7). For Wikipedia-based tasks, we set γ = 5e -4 for the 7B model, γ = 5e -3 for the 13B model, and γ = 5e -2 for the 70B model. For RA-NEST, all models use the same γ = 5e -1 for all tasks except Pile of Law which still uses γ = 5e -4 We observe that as the model gets stronger, using larger γ which leads to more rejection, is more beneficial to generation quality. The complete NEST procedure is provided in Algorithm 1.Datasets. We sample subsets of WikiText-103, NQ, and Biography as dev sets for hyper-parameter tuning. We use the validation sets of WikiText-103 (CC BY-SA 3.0), NQ (Apache License 2.0), TriviaQA (Apache License 2.0), MedMCQA (MIT License), HotpotQA (Apache License 2.0), MMLU (MIT License), and Biography (MIT License) for validation. TruthfulQA (Apache License 2.0) only has the test set. We finally test all datasets shown in Table 1 . For HotpotQA and MedMCQA, we do not have access to the test set and therefore the validation results are reported. For Biography, we use the labeled data that have human annotation as the validation and dev set, and the unlabeled data as the test set. In the original FACTSCORE paper, the authors use InstructGPT (Ouyang et al., 2022) to perform fact decomposition before verification. We hereby train our own decomposition model by further fine-tuning Llama-2 7B using publicly available datasets (Chen et al., 2022; Liu et al., 2023; Malaviya et al., 2024) . For fact verification, we use the option retrieval+llama+npm to evaluate the decomposed atomic facts.Inference and prompts. For language modelling and text completion, we use a context length of 128 tokens and a max generation length of 256 tokens. For the other tasks, we use 128 tokens as max generation length for question answering and 512 for fact verification. We remove all the in-context demonstrations from the prompt to test the zero-shot effectiveness of our model. We use greedy decoding in our experiments as the randomness in sampling can undermine factuality.Regarding the prompts we use for evaluation, for MMLU, we compare the perplexity of each option concatenated with the question and select the one with the minimum perplexity. Hyper-parameters and baselines. For the base LM, we do not tune the hyper-parameters released with the original Llama-2-chat models. For the in-context retrieval augmented baseline, we select the top-3 retrieved passages. For kNN-LM, we follow Equation (3) and use an interpolation coefficient of 0.7 for Wikipedia-based tasks and 0.9 for Pile of Law. For NEST and RA-NEST, we first tune the hyper-parameters in Equation ( 4) on language modelling tasks using perplexity. We then fix those hyper-parameters and then tune the rest of the parameters in Equation ( 6) and ( 7) on generation tasks. All hyper-parameters in the above methods are tuned on the dev sets of WikiText-103, NQ, and Biography.The following analyses are performed on the validation set of WikiText-103, NQ, and Biography data with the Llama-7B-chat model.Number of retrieved passages and tokens Khandelwal et al. (2020) show that increasing the size of the database and the number of tokens can improve the perplexity with proper hyper-parameter setting. We also verify whether our two-stage k-NN search and RRC approach follow the same trend. Figure 3a shows the validation perplexity on WikiText-103. For a fixed number of passages, the perplexity decreases as the number of tokens increases; for a fixed number of tokens, the perplexity decreases about 0.5 ∼ 1.0 as the number of passages doubles. However, as NEST needs to encode the retrieved passages on the fly, the latency also increases linearly w.r.t. the number of passages. Therefore, we set the passage number to be 40 and the token number to be 1024 in the main experiments.Interpolation coefficient Figure 3b shows the sensitivity of the hyper-parameters α (offset) and τ (temperature) in Equation (4) on WikiText-103. When τ is big, λ t is close to a uniform distribution and therefore the offset α does not have a big impact on the perplexity. When τ is small, the impact of α is enlarged and the sweet spot is achieved around τ = 0.1 and α = 0.4.Threshold for dynamic span selection Figure 3c shows how threshold δ in Equation ( 6) affects the generation on NQ. A bigger δ means selecting the span instead of a token more often. We can see that the answer-level recall on NQ first increases and then decreases as we increase the value of δ, where the sweet spot is around δ = 0.5.Wiki./ROUGE-1 NQ/ALR Bio./FS Table 3 shows a progressive ablation of NEST on WikiText-103, NQ, Biography. As mentioned in Section 3.1, it is extremely expensive to encode billion-token corpus with billion-parameter models. Therefore, we directly start with the two-stage implementation of kNN-LM and gradually add the methods applied in NEST. As we can see, adding the RRC component gives the first effectiveness boost. The second dynamic span selection method does not seem to increase the effectiveness, yet it is crucial to give consistent attribution for consecutive spans and tokens. The last relaxed speculative decoding method further improves the final generation quality.Answer: [Yes] Justification: The code base is released at https://github.com/facebookresearch/ NEST/tree/main. All the datasets we used are publicly available as discussed in Section 4.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.We used a slightly different implementation to ensure the sampled token is in π. Please see the code here: https://github.com/facebookresearch/NEST/blob/main/models/knn_transformers.pyhttps://huggingface.co/datasets/pile-of-law/pile-of-law/tree/mainhttps://huggingface.co/meta-llama/Llama-2-70b-chat-hf 


Decision: The paper presents a competitive technique to incorporate real-world text spans of arbitrary length into LM generations, thus enhancing also attribution.
All major concerns from reviewers appear resolved during rebuttal.

Reviews: review: paper_summary: paper_summary: The paper introduces NEST, a novel semi-parametric language modeling approach that enhances the generation quality and attribution of Large Language Models (LLMs) by incorporating real-world text spans. NEST employs a two-stage k-NN search and speculative decoding, achieving improved performance and reduced inference time.
strength_weakness: strength_weakness: The paper could benefit from a more detailed comparison with state-of-the-art methods.

The generalizability of NEST across different domains and languages needs further exploration.

The potential impact of NEST on the in-context learning ability of LLMs is not thoroughly discussed.

The paper lacks a comprehensive analysis of error rates and statistical significance.
questions: questions: How does NEST compare with other advanced language models in terms of handling long-tail knowledge?

What are the implications of using NEST for models pre-trained with different objectives or datasets?

Could the authors elaborate on the impact of NEST on the diversity and creativity of LLM generations?

On page 2, the authors mention a 1.8× speedup; could they specify if this is consistent across different model sizes?

How does NEST handle potential biases in the retrieved text spans from the corpus?

Can the authors provide more insight into the decision process behind the choice of hyperparameters?

Is there a risk of overfitting the corpus used for training the key-value datastore?
limitations: limitations: N/A
paper_summary: paper_summary: This work presents Nearest Neighbor Speculative Decoding (NEST), a technique to better inject real-world text spans into the output of existing language models. NEST is a kNN-LM approach adding an initial passage retrieval step. During inference, NEST uses Relative Retrieval Confidence (RRC) for confidence-based interpolation, dynamically extends selected tokens to include text spans when confidence is high, and employs a relaxed speculative decoding process that accepts only highly probable token spans.

According to the evaluation conducted by the authors with Llama 2 chat, NEST significantly outperforms both the base LM and standard kNN-LM, in terms of speed and accuracy.
strength_weakness: strength_weakness: - It’s unclear to me what are the most important contributions of this work. What is really new and has the most impact on the results? Even with the help of the related work section I cannot answer this question, hence my low confidence score.
questions: questions: It’s unclear to me what are the most important contributions of this work. What is really new and has the most impact on the results? Even with the help of the related work section I cannot answer this question, hence my low confidence score
limitations: limitations: The limitations are pointed out and sufficiently discussed.
paper_summary: paper_summary: This paper:

* Introduces NEST, a semi-parametric language modeling approach that integrates real-world text spans into language model generations.
* Enhances generation quality and reduces latency by using token-level retrieval and speculative decoding.
* Outperforms conventional kNN-LM and competes well with in-context retrieval methods.
* Demonstrates significant improvements with various models.
strength_weakness: strength_weakness: * Performance heavily relies on the accuracy of the first-stage passage retrieval and second-stage token retrieval.
* The retrieval process may still be complex and resource-intensive for practical deployment.
* The gains in performance and efficiency are less pronounced in larger models.
questions: questions: * How scalable is the NEST approach in real-world scenarios?
* What measures are taken to mitigate the impact of noise and errors introduced during the data retrieval stages?
* How does NEST address potential biases that may arise from the retrieved data?
* How well does NEST adapt to the evolving nature of language?
limitations: limitations: * The overall effectiveness of NEST is contingent upon the precision of both stages of text retrieval.
* The approach may struggle with ensuring that the retrieved content is contextually relevant and free from biases.
paper_summary: paper_summary: The paper presents a new semi-parametric language modeling approach that can incorporate text spans from a datastore into LLM-based generation improving both quality and attribution of generated texts. They propose a two-step approach that requires constructing an on-the-fly token-level datastore based on a small number of retrieved passages from a passage-level datastore. The current token is generated from the mixture distribution between the base LLM and the retrieved token-level distribution interpolated using a relative retrieval confidence score capturing the uncertainty of the token retriever. Furthermore, the approach enables extending the generation from token to an n-gram span based on a speculative decoding that rejects or accepts the tokens in the continuations, thus enabling span-level generation, improving efficiency.
strength_weakness: strength_weakness: 1. Multiple-token nearest neighbor generation has been proposed in prior work but was not discussed or compared, see: [1] Chunk-based Nearest Neighbor Machine Translation (https://aclanthology.org/2022.emnlp-main.284/).
questions: questions: See weaknesses.
limitations: limitations: Yes

Comments: title: Respectfully Asking to Reconsider Your Rating
comment: Thank you again for your feedback. We're happy to address any remaining concerns you may have and to provide further clarifications to ensure a thorough understanding of our research.
Thank you again for your feedback. 

We would like to follow up on our rebuttal to ensure that we have thoroughly addressed your initial concerns, and to respectfully ask for a reconsideration of the overall rating. Specifically, we would like to provide additional clarification on the two areas of concerns raised.

(1) **Novelty:** Our proposed approach, NEST, primarily consists of two novel techniques: (1) a confidence-adjusted, token-level retrieval score to extract text spans from real-world corpora as draft inputs, and (2) a relaxed-speculative decoding procedure to seamlessly integrate these drafts into the LLM generation process, rejecting the uncertain suffixes. We demonstrated empirically that this approach allows attribution of the generated text directly to the source, and yields significant improvements in both generation speed and quality compared to the base LLMs and kNN-LM with two-stage retrieval. kNN-LM with two-stage retrieval is a stronger baseline compared to the standard kNN-LM. Reviewers nNiG, vMVQ, and kPHS have also recognized the novelty and strengths of our approach.

(2) **Which technique had the most impact on performance:** We included an ablation study in the appendix of our submission (**Appendix C.2, Table 3**). The results show that both adjusting the interpolation score based on token-level retrieval confidence (+4.6 ROUGE-1 for WikiText-103 and +6.8 FactScore on Biography) and performing token rejection using relaxed speculative decoding (+2.3 ROUGE-1 for WikiText-103 and +5.2 FactScore on Biography) significantly contribute to the performance improvements. 

We plan to move the ablation study to the main text to provide clearer insights for future readers with additional main content space.
| Models (7B)                     | Wiki./ROUGE-1 | NQ/ALR | Bio./FS |
|---------------------------------|---------------|--------|---------|
| kNN-LM (two-stage)              | 20.1          | 40.8   | 34.8    |
| + Relative Retrieval Confidence | 24.7          | 44.4   | 41.6    |
| + Dynamic Span selection        | 24.5          | 44.6   | 41.6    |
| + Relaxed speculative decoding  | 26.8          | 45.4   | 46.8    |
I have read the author's response and other reviews, and I will keep my rating unchanged.