Nearest Neighbor Speculative Decoding for LLM Generation and Attribution

ArXiv: 2405.19325

🎯 Pitch

Nest makes retrieval-augmented text generation faster than standard decoding—achieving a 1.8× speedup on Llama-2-70B—while simultaneously forcing outputs to copy verifiable spans from a corpus, boosting factuality by over 21% on biography generation.

---

1. Executive Summary

This paper introduces Nearest Neighbor Speculative Decoding (Nest), a novel semi-parametric language modeling approach that incorporates real-world text spans of arbitrary length into LLM generations while providing source attribution. Evaluated on nine knowledge-intensive benchmarks—including WikiText-103, Natural Questions, Biography, and MMLU—using zero-shot Llama-2-Chat models (7B, 13B, 70B), Nest combines confidence-based output interpolation (a dynamic mixing coefficient derived from token retrieval uncertainty, called Relative Retrieval Confidence), dynamic span selection (extracting n-gram continuations from the corpus when retrieval confidence exceeds a threshold), and relaxed speculative decoding (an approximate rejection sampling procedure that accepts span prefixes likely under the mixture distribution) to improve both generation quality and attribution over standard kNN-LM methods. Nest achieves a 1.8× speedup in inference time on Llama-2-Chat 70B, a 42.3% improvement in ROUGE-1 on WikiText-103, and a 21.6% improvement in FActScore on Biography, establishing that direct span-level attribution with low latency is achievable when the corpus retrieval and speculative acceptance mechanisms are calibrated to the base LM's uncertainty.

2. Context and Motivation

The Core Problem: LLM Hallucination Without Verifiable Provenance

The fundamental issue this paper tackles is that large language models (LLMs) hallucinate—they generate plausible-sounding but factually incorrect content—and critically, they provide no intrinsic mechanism for attributing their claims to sources. When an LLM states "The song 'Does He Love You' is a duet between Reba McEntire and Linda Davis," a user has no way to verify whether this claim was invented by the model or faithfully reproduced from a real document. This is not merely an inconvenience; in high-stakes domains like medicine, law, or journalism, the inability to trace generated text back to a source renders LLM outputs effectively unusable without extensive human verification.

The problem is particularly acute for long-tail knowledge—facts that appear infrequently in training data. Kandpal et al. (2023) demonstrated that LLMs struggle disproportionately with such knowledge, precisely where retrieval from an external corpus would be most valuable. The paper frames this as a dual failure: LLMs both get facts wrong and cannot show their work.

Why This Problem Matters

The importance extends along several dimensions:

Trust and verifiability. When an AI system generates a factual claim, downstream users—whether doctors reviewing a suggested diagnosis, journalists fact-checking a generated article, or students using an AI tutor—need to verify the provenance of that claim. Without attribution, every LLM output must be independently fact-checked, which largely defeats the purpose of automation. Nest's central promise is to make verification trivial: if a span of text is green-highlighted with a citation, the user can trace it directly to a specific passage in the corpus.

Latency in retrieval-augmented systems. Even when attribution is theoretically possible (e.g., through in-context retrieval augmentation that prepends retrieved passages to the prompt), existing methods impose significant inference-time overhead. The LM must process all retrieved passages as part of its context, expanding the already-large key-value cache and increasing per-token generation cost. By generating potentially many tokens at each inference step (through the speculative decoding mechanism), Nest actually reduces latency relative to a base LM, achieving a 1.8× speedup on the 70B model. This is a striking inversion of the usual tradeoff: Nest provides both better attribution and lower latency, rather than sacrificing one for the other.

Faithfulness through copying. A subtle but important point: when an LM generates text from its parametric knowledge, even factually correct statements are paraphrased—the model produces tokens that capture the gist of something it learned during training, but the exact wording is novel. Paraphrasing introduces the risk of subtle factual distortions. By directly copying spans from a trusted corpus, Nest ensures that attributed portions of the output are literally present in the source material, eliminating the paraphrasing error mode entirely. This is a fundamentally different guarantee from in-context retrieval augmentation, where the model may attend to retrieved passages but still generates novel token sequences that may drift from the source.

Where Existing Approaches Fall Short

The paper identifies three categories of prior work on retrieval-augmented language models (RALMs), each with distinct limitations:

In-context retrieval augmentation (input augmentation). The most widely adopted approach—represented by REPLUG (Shi et al., 2024a), in-context RALM (Ram et al., 2023), and Self-RAG (Asai et al., 2023b)—prepends retrieved passages to the model's input prompt. The LM then conditions on this expanded context during generation. While effective at biasing the output toward the retrieved information, this approach suffers from several weaknesses:

• No guarantee of faithfulness. The LM may ignore the retrieved passages entirely, attend to irrelevant portions, or combine information from multiple sources in ways that introduce new errors. The retrieval influences the output distribution but does not constrain it—the model can still hallucinate even when presented with correct evidence.
• Context window competition. Each retrieved passage consumes valuable context window space, limiting the number of passages that can be provided and potentially crowding out task instructions, conversation history, or other important context.
• Latency penalty. Processing a longer prompt prefix (with all retrieved passages) increases the computational cost of every subsequent token in the generation, since the key-value cache for the full context must be maintained and attended to.

Intermediate fusion (RETRO, InstructRetro). Approaches like RETRO (Borgeaud et al., 2022) and its successors integrate retrieved information into intermediate layers of the transformer through specialized cross-attention mechanisms. While powerful, these methods require architectural modifications and retraining from scratch—they cannot be applied as a plug-and-play wrapper around off-the-shelf LMs like the Llama-2-Chat models used in this paper. This limits their practical adoption, as most practitioners work with existing pretrained models rather than training new architectures.

Output integration (kNN-LM and variants). The kNN-LM framework (Khandelwal et al., 2020) pioneered the approach of interpolating the LM's output distribution with a non-parametric distribution derived from nearest-neighbor tokens in a corpus. At each inference step, the model's hidden state is used as a query to search a key-value store of (context representation, next token) pairs, producing a retrieval distribution $p_{\text{k-NN}}$ that is mixed with the LM's parametric distribution $p_{\text{LM}}$ via a fixed interpolation coefficient $\lambda$:

$p_{\mathcal{M}}(w|x,y_{<t}) = \lambda \cdot p_{\text{LM}}(w|x,y_{<t}) + (1-\lambda) \cdot p_{\text{k-NN}}(w|x,y_{<t})$

While this provides more direct attribution (you can point to the specific retrieved token), kNN-LM suffers from several critical shortcomings that Nest specifically addresses:

• Degraded generation quality for open-ended text. Wang et al. (2023a) showed that kNN-LM harms the fluency and coherence of generated text. The problem arises because the retrieval distribution is token-level and myopic—each token is retrieved independently based on similarity of the current context, without considering whether consecutive retrieved tokens come from a coherent source. The result is a "patchwork quilt" of tokens drawn from unrelated documents, which reads unnaturally even when individual tokens are appropriate.
• Fixed interpolation coefficient. The standard kNN-LM uses a single global $\lambda$ for all tokens and all contexts. This is manifestly suboptimal: when the LM is highly confident about the next token (e.g., completing a common phrase), the parametric distribution should dominate; when the LM is uncertain (e.g., recalling a specific factual detail that may have been poorly learned during training), the non-parametric distribution should dominate. A fixed $\lambda$ cannot adapt to this variation.
• Computational scaling challenges. The original kNN-LM requires a datastore containing every (context, next-token) pair in the entire corpus. For a corpus the size of Wikipedia (billions of tokens), this datastore becomes prohibitively large to store and search efficiently at inference time. The two-stage retrieval design Nest introduces is directly motivated by this scalability limitation.

Copy-generator approaches (CoG). Lan et al. (2023) proposed the Copy Generator (CoG), which jointly trains a phrase encoder and LM to dynamically expand the vocabulary using retrieved phrases. While related in spirit to Nest's span-copying mechanism, CoG requires joint training of the retrieval and generation components, making it incompatible with off-the-shelf LMs. Nest implements a conceptually similar idea—copying multi-token spans rather than individual tokens—but does so without any training, relying instead on the speculative decoding mechanism to validate that copied spans are acceptable under the mixture distribution.

The Fundamental Tension: Accuracy vs. Fluency in Semi-Parametric LMs

Underlying all of these limitations is a deeper tension that the paper confronts directly. Semi-parametric LMs face a fluency-factuality tradeoff:

• Parametric generation (the base LM) produces fluent, coherent text because the model has learned the syntactic and semantic patterns of natural language through pretraining. However, it hallucinates facts, especially for long-tail knowledge.
• Non-parametric retrieval (the corpus) provides accurate facts but at the cost of coherence. Blindly pasting tokens from potentially unrelated documents into the generation destroys fluency, as kNN-LM's poor open-ended generation performance demonstrates (Wang et al., 2023a).

Nest's design can be understood as an attempt to resolve this tension by intervening at two points:

1. During token selection: The confidence-based interpolation adaptively weights the parametric vs. non-parametric distributions based on the retriever's uncertainty, so the model leans on the corpus when it "knows what it doesn't know" and defaults to its own distribution otherwise.
2. During span acceptance: The relaxed speculative decoding procedure evaluates candidate spans in context, using the full mixture distribution conditioned on all previously accepted tokens. A span is only accepted if it is likely under this distribution, meaning the model verifies that the corpus text fits naturally into the generation. This provides a coherence filter that pure kNN-LM lacks.

How This Paper Positions Itself

The paper positions Nest as a training-free, plug-and-play method that combines the best aspects of three lines of work while addressing their individual weaknesses:

• From kNN-LM: It inherits the direct output-level attribution via token retrieval, but replaces the fixed interpolation coefficient with a dynamic, confidence-based mixture.
• From Copy Generator: It inherits the idea of copying multi-token phrases rather than individual tokens, but eliminates the training requirement by using speculative decoding for validation rather than learned copy mechanisms.
• From speculative decoding (Leviathan et al., 2023): It borrows the rejection sampling procedure that enables multi-token generation with quality guarantees, but adapts it to operate on corpus-retrieved spans (rather than drafts from a smaller model) and relaxes the acceptance criterion to account for the unknown proposal distribution.

The paper explicitly frames Nest as an output integration method in the taxonomy of Asai et al. (2024), contrasting it with input augmentation (in-context RA) and intermediate fusion (RETRO). A key empirical claim is that these approaches are complementary: Nest can be combined with in-context RA (the "RA-Nest" configuration in the experiments) to achieve gains beyond either method alone, suggesting that output-level and input-level retrieval augmentation address different failure modes of the base LM.

The paper also positions itself relative to the concurrent work REST (He et al., 2024), which also uses a datastore for speculative decoding. The critical difference is that REST uses the datastore purely for draft generation (maintaining the original LM's output distribution), while Nest actively modifies the output distribution through interpolation with the retrieval distribution. This makes REST an acceleration-only method, whereas Nest simultaneously improves factuality, attribution, and speed. The paper thus argues that Nest occupies a unique position in the design space: it is simultaneously a retrieval-augmentation method (changing what the model generates) and an inference-acceleration method (changing how fast it generates), unified through the relaxed speculative decoding framework.

3. Technical Approach

3.1 Reader Orientation

Nest is a training-free inference-time method that wraps around an off-the-shelf language model and, at each generation step, retrieves relevant tokens from a corpus, interpolates the LM's predictions with the retrieval distribution using a dynamic confidence-based coefficient, and then copies multi-token spans from the corpus—validating them through an approximate speculative decoding procedure—so that the generated text contains verbatim source segments with direct attribution. The system solves the problem that standard kNN-LM degrades fluency while providing no span-level attribution and operates at high latency, by adaptively deciding when to trust the retriever (via Relative Retrieval Confidence), what to copy (via dynamic span selection from n-gram continuations), and how much to accept (via relaxed speculative decoding that filters out corpus text unlikely under the mixture distribution).

3.2 Big-Picture Architecture (Diagram in Words)

The Nest system has six major components, arranged in a pipeline that executes at every generation step:

1. First-Stage Passage Retriever — a hybrid dense (Dragon+) plus sparse (BM25) search system that, given the input prefix, retrieves the top-$b$ most relevant passages from the corpus. This prunes the search space from billions of tokens to a few thousand, making token-level search tractable. It runs once per query (or periodically), not at every token.

2. On-the-Fly Token-Level Key-Value Store Builder — uses the LM's own encoder $f(\cdot)$ to encode every prefix in the retrieved passages, producing a key-value store $(\mathcal{K}', \mathcal{V}')$ where keys are hidden states (the input to the final layer's feed-forward network after layer normalization) and values are the corresponding next tokens. This store is rebuilt for each generation step as the set of retrieved passages may change.

3. Token-Level k-NN Retriever — at each inference step, encodes the current generation context to produce a query vector $q_t$, searches $(\mathcal{K}', \mathcal{V}')$ for the top-$r$ nearest neighbors using negative squared $\ell_2$ distance, and computes a non-parametric next-token distribution $p_{\text{k-NN}}(w|x, y_{<t})$ via softmax over the similarity scores.

4. Confidence-Based Interpolator — computes a dynamic interpolation coefficient $\lambda_t$ from the ratio of minimum to maximum retrieval similarity scores (the Relative Retrieval Confidence), then mixes the LM's parametric distribution $p_{\text{LM}}$ with the non-parametric $p_{\text{k-NN}}$ to produce the mixture distribution $p_{\mathcal{M}}$. When the retriever is uncertain (similarity scores cluster together), $\lambda_t$ is high and the LM dominates; when the retriever is confident (one match is much stronger than others), $\lambda_t$ is low and the retrieval distribution dominates.

5. Dynamic Span Selector — samples (or greedily selects) the next token from $p_{\mathcal{M}}$, then—if $\lambda_t$ is below a threshold $\delta$—extends that token into an n-gram by following its continuation in the source passage from the corpus. This produces a candidate multi-token span for potential acceptance.

6. Relaxed Speculative Decoder — evaluates each token in the candidate span sequentially against the mixture distribution $p_{\mathcal{M}}$ conditioned on all previously accepted tokens. A token is accepted with probability proportional to how its mixture probability compares to the maximum probability token at that position (times a relaxation factor $\gamma$). The first rejected token truncates the span, and a new token is sampled from $p_{\mathcal{M}}$ to replace it. This ensures corpus text is only incorporated when the full mixture model considers it likely.

Information flows as follows: a prompt enters → first-stage retriever selects top-$b$ passages → at each generation step, the LM encodes the current prefix to produce $q_t$ and $p_{\text{LM}}$ → token-level search uses $q_t$ to retrieve neighbors and compute $p_{\text{k-NN}}$ → confidence-based interpolator mixes the distributions to produce $p_{\mathcal{M}}$ → dynamic span selector picks the best token and optionally extends it to an n-gram → relaxed speculative decoding accepts or rejects each token in the span → accepted tokens are appended to the output, rejected tokens trigger resampling from $p_{\mathcal{M}}$ → accepted spans are masked in the corpus to prevent repetition.

3.3 Roadmap for the Deep Dive

• First, the two-stage $k$-NN search architecture, because it is the prerequisite that makes token-level retrieval scalable and establishes how the key-value store is constructed—this is infrastructure shared by both the standard $k$NN-LM baseline and Nest.
• Second, the confidence-based output interpolation mechanism (Relative Retrieval Confidence), because it is the core adaptive mixing strategy that dynamically decides how much to trust the retriever versus the LM at each step—this replaces the fixed $\lambda$ of standard $k$NN-LM and is a necessary input to all downstream mechanisms.
• Third, the dynamic span selection procedure, because it extends single-token retrieval into multi-token copying and establishes how spans are extracted from the corpus—this depends on both the interpolation coefficient (to decide whether to copy a span) and the retrieval scores (to decide which span to copy).
• Fourth, the relaxed speculative decoding algorithm, because it provides the quality-control filter that validates copied spans under the mixture distribution and determines the final accepted span length—this builds on the mixture distribution from the interpolator and the candidate spans from the selector.
• Fifth, the full algorithm walkthrough (what happens at one inference step from start to finish), to tie all components together into a concrete procedural narrative.
• Sixth, the design rationale and key hyperparameter choices explained together, to clarify why each parameter exists, how it was set, and what tradeoffs it governs.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a systems-and-methods paper whose core idea is that corpus-retrieved spans can be incorporated into LM generation with quality, attribution, and latency improvements if three mechanisms work in concert: (1) a two-stage retrieval architecture that makes token-level search tractable, (2) a dynamic interpolation that adapts to retriever uncertainty, and (3) a relaxed speculative decoding procedure that validates copied spans under the full mixture distribution.

---

Two-Stage $k$-NN Search Architecture

The original $k$NN-LM (Khandelwal et al., 2020) requires storing a key-value pair for every token in the entire corpus. For each token position in the corpus, the LM encodes the preceding context to produce a hidden state (the key), and the actual next token becomes the value. Formally, for a corpus $\mathcal{D}$:

$(\mathcal{K}, \mathcal{V}) = \{(h_i, w_i) \mid (c_i, w_i) \in \mathcal{D}\}$

where $h_i = f(c_i)$ is the hidden state produced by the LM's encoder $f(\cdot)$ for context $c_i$ (the sequence preceding token $w_i$), and $w_i$ is the actual next token in the corpus. The size of this datastore is proportional to the total number of tokens in $\mathcal{D}$, which for Wikipedia-scale corpora runs into billions of entries. Storing and searching this datastore at every inference step is computationally prohibitive—it requires either massive RAM (to hold all keys in memory) or expensive disk-based approximate nearest neighbor search with high latency.

Nest introduces a two-stage design that is widely used in information retrieval systems (search engines, document retrieval) but had not been applied to $k$NN-LM inference:

First-stage passage retrieval. The corpus $\mathcal{D}$ is segmented into passages of at most $m$ tokens each (where $m=200$ in the paper's implementation). A hybrid retrieval system combining dense and sparse retrieval is used to identify the most relevant passages given the input $x$:

• Dense retrieval uses Dragon+ (Lin et al., 2023), a state-of-the-art dense retriever that encodes each passage into a single fixed-length vector. These vectors are indexed using FAISS (Douze et al., 2024) with an IVF-PQ structure (inverted file with product quantization): 65,536 centroids and 256 quantization codes of 8 bits each. At query time, the input is encoded by the same Dragon+ model, and the top-$(b \cdot l)$ passages are retrieved by approximate nearest neighbor search with nprobe=4096 (the number of Voronoi cells to search).
• Sparse retrieval uses BM25 (Robertson and Zaragoza, 2009) with Pyserini (Lin et al., 2021), a traditional lexical matching algorithm that scores passages based on term frequency and inverse document frequency. This captures exact keyword matches that dense retrieval may miss—complementary strengths, since dense models handle semantic similarity and sparse models handle lexical overlap.
• Fusion. Both retrievers return their top-$(b \cdot l)$ passages along with similarity scores. The scores are linearly interpolated using a weighting coefficient $\eta$ that is itself adaptive: $\eta_{\text{dense}} = 1 - \text{top-100}(s_{\text{dense}}(q, d)) / \max s_{\text{dense}}(q, d)$, meaning that when the dense retriever's top scores are highly concentrated (one passage dominates), the dense component gets lower weight, and similarly for sparse. The final fusion coefficient is $\eta = 0.7 \cdot (1 - \eta_{\text{sparse}}) + 0.3 \cdot \eta_{\text{dense}}$. The fused similarity for each passage is $s(q, d) = \eta \cdot s_{\text{dense}} + (1 - \eta) \cdot s_{\text{sparse}}$. Passages missing from one retriever's results are assigned the minimum similarity from that retriever's returned set. After fusion and sorting, the top-$b$ passages are selected (where $b=40$ and $l=100$ in all experiments).

The scaling factor $l=100$ means that 4,000 passages are retrieved from each index before fusion, but only 40 survive to the second stage. This over-retrieval is a standard technique in two-stage retrieval: the first stage casts a wide net to maximize recall (finding all potentially relevant passages), and the second stage (token-level search) provides precision by examining the actual token-level matches. The passage-level index sizes are modest: the Wikipedia dense index is approximately 8.96 GB on disk, and the sparse index approximately 3.48 GB. This is orders of magnitude smaller than what a full token-level index for the same corpus would require.

Second-stage token-level $k$-NN search. Once the top-$b$ passages $\{d_1, d_2, ..., d_b\}$ are identified, a token-level key-value store $(\mathcal{K}', \mathcal{V}')$ is constructed on the fly. Unlike the original $k$NN-LM which pre-builds and stores this for the entire corpus, Nest builds it per inference step from only the retrieved passages. The procedure:

1. For each retrieved passage $d_i$, tokenize it into $\{w_1^{d_i}, w_2^{d_i}, ..., w_m^{d_i}\}$.
2. Use the LM's encoder $f(\cdot)$ to encode each prefix of the passage: $h_j^{d_i} = f(w_1^{d_i}, ..., w_j^{d_i})$ for $j = 1$ to $m-1$. This produces a hidden state for every position except the last (since there is no "next token" for the final position).
3. Add each $(h_j^{d_i}, w_{j+1}^{d_i})$ pair to the key-value store $(\mathcal{K}', \mathcal{V}')$. The key is the hidden state after processing the prefix up to position $j$; the value is the token that actually followed in the corpus.

The encoder $f(\cdot)$ is defined as the input to the final layer's feed-forward network after layer normalization, following exactly the same representation choice as Khandelwal et al. (2020). This layer was chosen in the original $k$NN-LM work after empirical comparison of representations from different layers, with the final-layer pre-FFN representation providing the best tradeoff between semantic abstraction (higher layers capture more contextual meaning) and token-predictiveness (closer to the output, so more directly useful for retrieving likely next tokens).

At inference time, the current generation prefix $y_{<t}$ is encoded by the same $f(\cdot)$ to produce the query vector $q_t$. This query is used to search $(\mathcal{K}', \mathcal{V}')$ using negative squared $\ell_2$ distance as the similarity function:

$s(q, k) = -\|q - k\|_2^2$

The top-$r$ nearest neighbors are retrieved (where $r=1024$ in all experiments; this is the maximum number of tokens considered). The non-parametric next-token distribution is then:

$p_{\text{$k$-NN}}(w \mid x, y_{<t}) \propto \sum_{(k_i, v_i) \in \pi} \mathbb{I}_{w = v_i} \exp(\mu \cdot s(q_t, k_i))$

where $\pi$ is the set of $r$ retrieved neighbors (each a key-value pair), $\mathbb{I}_{w = v_i}$ is an indicator that is 1 only when the token being scored matches the retrieved value $v_i$ and 0 otherwise, $s(q_t, k_i)$ is the negative squared $\ell_2$ distance between the query and the $i$-th key, and $\mu = 1/\sqrt{\text{dim}(q_t)}$ is the inverse temperature (scaling the exponent to account for the dimensionality of the hidden state, analogous to the scaling in Transformer attention from Vaswani et al., 2017).

What this computes, operationally: For each token $w$ in the vocabulary, the formula sums the exponentiated similarities of all retrieved neighbors whose value matches $w$, then normalizes across the vocabulary. This means that if multiple retrieved neighbors all point to the same token (e.g., "the" appears as the next token in many similar contexts), their similarities accumulate, giving that token higher total probability. Tokens that do not appear among the $r$ retrieved neighbors receive probability 0—the distribution is sparse, restricted to at most $r$ distinct tokens.

Why this form: The softmax-over-similarities is mathematically equivalent to doing attention over the retrieved key-value pairs, where the query is $q_t$, the keys are the $k_i$, and the values are the one-hot indicator of $v_i$. The inverse temperature $\mu$ controls the sharpness: higher $\mu$ makes the distribution more peaked around the most similar neighbor, lower $\mu$ makes it more uniform across neighbors. The $\ell_2$ distance is used rather than cosine similarity because Khandelwal et al. (2020) empirically found it performed better for this task—possibly because the magnitude of the hidden state vector encodes information about token predictability (high-norm states correspond to high-confidence predictions) that is lost when cosine similarity normalizes away magnitude.

Why two stages are necessary: A single-stage token-level search over all tokens in Wikipedia would require storing and searching billions of key-value pairs. The passage filter reduces this to searching over approximately $b \times (\text{avg tokens per passage}) \approx 40 \times 200 = 8,000$ tokens per step—a reduction of roughly six orders of magnitude. The tradeoff is that some relevant tokens may be missed if the correct passage is not in the top-$b$, but the hybrid dense+sparse retrieval with aggressive first-stage recall (retrieving 4,000 passages before pruning to 40) mitigates this. The paper's perplexity results (Appendix C, Figure 3a) show that increasing the number of passages beyond 40 yields diminishing returns while linearly increasing latency, confirming 40 as a reasonable operating point.

---

Confidence-Based Output Interpolation (Relative Retrieval Confidence)

The standard $k$NN-LM (Equation 3 from the paper) uses a fixed global interpolation coefficient $\lambda \in [0, 1]$ to mix the parametric LM distribution and the non-parametric retrieval distribution:

$p_{\mathcal{M}}(w \mid x, y_{<t}) = \lambda \cdot p_{\text{LM}}(w \mid x, y_{<t}) + (1 - \lambda) \cdot p_{\text{$k$-NN}}(w \mid x, y_{<t})$

When $\lambda = 1$, the model is purely parametric (the base LM). When $\lambda = 0$, it is purely non-parametric (only the retrieved neighbors matter). A fixed $\lambda$ treats all tokens and all contexts identically, which is clearly wasteful: when the LM is highly confident about a common function word (e.g., "the" after "in"), there is no benefit to consulting the retriever; when the LM is uncertain about a rare entity, the retriever's vote should carry more weight.

Nest replaces the fixed $\lambda$ with a dynamic, per-step coefficient $\lambda_t$ computed from the retrieval scores themselves. The key insight is that the distribution of retrieval similarities encodes how confident the retriever is—if all retrieved neighbors have similar scores, the retriever is uncertain about which token to recommend; if one neighbor has a much higher score than all others, the retriever is confident.

The Relative Retrieval Confidence (RRC) formula is:

$\lambda_t = \sigma\left(\left(\frac{\min_i |s(q_t, k_i)|}{\max_i |s(q_t, k_i)|} - \alpha\right) / \tau\right)$

where $\sigma(\cdot)$ is the sigmoid function ($\sigma(x) = 1/(1 + e^{-x})$), $s(q_t, k_i)$ is the negative squared $\ell_2$ distance for the $i$-th retrieved neighbor, $\min_i |s(q_t, k_i)|$ is the minimum absolute similarity among the top-$r$ retrieved neighbors (i.e., the worst match), $\max_i |s(q_t, k_i)|$ is the maximum absolute similarity (the best match), $\alpha$ is an offset hyperparameter that shifts the sigmoid's inflection point, and $\tau$ is a temperature hyperparameter that controls the steepness of the sigmoid.

Operational meaning of the min-max ratio: Because $s(q_t, k_i)$ is negative squared distance, $|s(q_t, k_i)|$ is the absolute distance—larger values mean worse matches. The ratio $\min_i |s(q_t, k_i)| / \max_i |s(q_t, k_i)|$ is the distance to the best match divided by the distance to the worst match among the top-$r$. If all retrieved neighbors are similarly good matches, this ratio is close to 1, meaning the retriever cannot distinguish between them, and the resulting $\lambda_t$ will be large (the LM should dominate). If one neighbor is dramatically better than the others, the ratio is close to 0, the retriever is confident, and $\lambda_t$ will be small (the retrieval distribution should dominate).

The sigmoid maps this ratio from the range $(0, 1]$ to the range $(0, 1)$, with $\alpha$ controlling where the transition happens and $\tau$ controlling how sharp the transition is. For example, with $\alpha = 0.3$ and $\tau = 0.1$ (the settings used for all Wikipedia-based tasks), $\lambda_t \approx 0.95$ when the ratio is 0.4 (retriever uncertain), and $\lambda_t \approx 0.05$ when the ratio is 0.1 (retriever confident). The specific values $\alpha = 0.3, \tau = 0.1$ were tuned on the WikiText-103 validation set for perplexity; the sensitivity analysis in Appendix C (Figure 3b) shows that perplexity is relatively flat near this region, indicating the method is not brittle to these choices.

Why absolute values in the ratio: The use of absolute similarities $|s(q_t, k_i)|$ rather than the raw signed similarities is important. The raw similarities are negative (since they are negative squared distances), so taking absolute values converts them to positive distances. The ratio of minimum to maximum absolute distance is always in $(0, 1]$, providing a normalized uncertainty measure that is invariant to the overall scale of the distances. Without the absolute values, the ratio of raw similarities could be any positive number (since the maximum would be less negative than the minimum), making the sigmoid input hard to calibrate.

Why sigmoid rather than a hard threshold: A hard threshold (e.g., "if ratio < 0.3, use $\lambda = 0$; else use $\lambda = 1$") would create discontinuities in the generation and would not capture the continuum of retriever confidence. The sigmoid provides a smooth interpolation between trusting the LM and trusting the retriever, which is important because real queries exist on a spectrum of retriever certainty. The temperature $\tau$ controls the smoothness: smaller $\tau$ makes the transition sharper (more like a threshold), larger $\tau$ makes it broader (more gradual blending). $\tau = 0.1$ produces a relatively sharp transition, meaning Nest mostly operates in a "mostly LM" or "mostly retriever" regime rather than a 50-50 blend.

Why this form over learned or heuristic alternatives: Prior work explored adaptive interpolation through training (He et al., 2021; Bhardwaj et al., 2023) or extensive per-task tuning (Drozdov et al., 2022). Training requires task-specific data and model modifications; excessive tuning breaks the zero-shot plug-and-play goal. RRC is a heuristic, but it is a principled one—it directly measures the retriever's internal uncertainty using only the retrieval scores that are already computed, requiring no additional forward passes, no training, and no per-task tuning beyond the two hyperparameters $\alpha$ and $\tau$. The ablation in Appendix C (Table 3) shows that adding RRC to a two-stage $k$NN-LM baseline improves WikiText-103 ROUGE-1 from 20.1 to 24.7, NQ answer-level recall from 40.8 to 44.4, and Biography FActScore from 34.8 to 41.6—confirming that dynamic interpolation is substantially better than a fixed $\lambda$.

---

Dynamic Span Selection

Standard $k$NN-LM selects one token at a time from the mixture distribution. This is problematic for two reasons: (1) consecutive tokens may be retrieved from unrelated documents, producing incoherent text (the "patchwork quilt" problem identified by Wang et al., 2023a), and (2) even when the retrieved token is correct, the attribution is limited to individual tokens, which is far less useful than attributing a complete factual claim to a source passage.

Nest extends token-level selection to multi-token span selection by following the n-gram continuation of a retrieved token in the source corpus. The procedure at time step $t$ is:

Step 1: Select the next token. The next token $w_t$ is selected from the mixture distribution $p_{\mathcal{M}}$ (using greedy decoding in the paper's experiments; sampling would also work). This token may correspond to multiple retrieved neighbors $\pi$ that happen to have the same value—for example, the token "United" might appear as the next token in several different retrieved contexts, each with its own similarity score.

Step 2: Choose the best-matching instance. Among the retrieved neighbors whose value equals $w_t$, the one with the highest $p_{\text{k-NN}}$ probability is selected:

$w_t^{(1)} = \underset{\{w_i \mid w_i = w_t, w_i \in \pi\}}{\text{argmax}} \ p_{\text{$k$-NN}}(w = w_i \mid x, y_{<t})$

This max-pooling strategy ensures that when multiple corpus contexts could have produced the same token, Nest picks the one that the retriever considers most similar to the current generation context. This instance determines where in the corpus the span will be drawn from.

Step 3: Extract the n-gram continuation. Given the selected instance $w_t^{(1)}$ and its position in the source passage, Nest extracts the next $n-1$ tokens from that passage, forming the n-gram $\{w_t^{(1)}, w_t^{(2)}, ..., w_t^{(n)}\}$. These are simply the tokens that appeared after $w_t^{(1)}$ in the original corpus—they are not individually retrieved or scored. The hyperparameter $n$ is set to 64 in all experiments, meaning Nest can potentially copy up to 64 tokens at once from the corpus.

Step 4: Decide whether to use the span or just the token. The decision is based on the interpolation coefficient $\lambda_t$ computed by RRC:

$y_t = \begin{cases} w_t, & \text{if } \lambda_t > \delta \\ \{w_t^{(1)}, w_t^{(2)}, ..., w_t^{(n)}\}, & \text{otherwise} \end{cases}$

where $\delta$ is a threshold hyperparameter set to 0.5 in all experiments regardless of model size or task.

Operational meaning: When $\lambda_t > \delta$, the RRC mechanism has determined that the LM is more trustworthy than the retriever at this step, so Nest generates only the single token $w_t$ from $p_{\mathcal{M}}$ without copying a span. When $\lambda_t \leq \delta$, the retriever is confident enough that Nest attempts to copy the n-gram continuation from the corpus. The threshold $\delta = 0.5$ means that span copying is attempted only when the retriever's confidence leads to an interpolation that gives the non-parametric distribution more than 50% weight in the mixture. The sensitivity analysis in Appendix C (Figure 3c) shows that answer-level recall on NQ peaks around $\delta = 0.5$—lower values (requiring even more retriever confidence) copy too rarely, while higher values copy too aggressively.

A critical implementation detail: Nest uses a "slightly different implementation" to ensure the sampled token is actually present in $\pi$. The standard procedure samples $w_t$ from $p_{\mathcal{M}}$, but since $p_{\mathcal{M}}$ is a continuous mixture of a full-support LM distribution and a sparse $k$-NN distribution, it's possible that the sampled token appears in $p_{\text{LM}}$ but not among the retrieved neighbors. In that case, no n-gram continuation can be extracted. The paper notes this in a footnote referencing the code base, but the exact resolution (rejection sampling? restricting the sample space to tokens in $\pi$?) is not specified in the text—the code would need to be consulted for the precise implementation.

Why copy n-grams rather than retrieve each token independently: Copying n-grams ensures that the attributed text is verbatim from a single source, providing coherent multi-token attribution. If Nest retrieved each token independently, it would suffer from the same incoherence problem as standard $k$NN-LM. The n-gram extraction also amortizes the retrieval cost: after one token-level search, Nest can potentially generate up to $n$ tokens without additional retrieval operations, which is one source of the latency improvement.

Why max-pooling among matching instances: When multiple retrieved neighbors have the same token value, they likely come from different contexts in different passages. Max-pooling selects the most contextually similar source, which maximizes the chance that the subsequent n-gram is relevant to the current generation. Alternatives like random selection or averaging would introduce unnecessary noise.

---

Relaxed Speculative Decoding

The dynamic span selector proposes an n-gram of up to 64 tokens, but this n-gram is extracted blindly—there is no guarantee that tokens $w_t^{(2)}, ..., w_t^{(n)}$ are actually appropriate continuations of the current generation, since they come from a different context in the corpus. Nest needs a mechanism to accept or reject these proposed tokens, accepting a prefix that fits naturally and truncating the rest.

This is essentially the same problem that speculative decoding (Leviathan et al., 2023) solves: given a draft of multiple tokens proposed by a fast but approximate model, use the target model (which is more accurate but slower) to verify and selectively accept the draft tokens in parallel. In standard speculative decoding, a small "draft" model proposes tokens, and the large "target" model evaluates them. Here, the corpus acts as the draft source (proposing the n-gram), and the mixture model $\mathcal{M}$ acts as the target (evaluating the proposal).

However, standard speculative decoding requires knowing the proposal distribution $q(w \mid x, y_{<t})$ to compute the correct acceptance probability:

$P(\text{accept } w) = \min\left(1, \frac{p_{\text{target}}(w)}{q(w)}\right)$

For Nest, the proposal distribution for tokens beyond the first ($w_t^{(2)}, w_t^{(3)}, ...$) is unknown—these tokens were simply copied from the corpus continuation, not generated by a known distribution $q$. Nest therefore uses a relaxed version of speculative decoding that upper-bounds the acceptance probability by replacing the unknown $q(w)$ with a conservative overestimate: the maximum probability assigned by $\mathcal{M}$ to any token at that position.

The acceptance probability for the $i$-th token in the span is:

$P(\text{accept token } w_t^{(i)}) = \min\left(1, \ \frac{p_{\mathcal{M}}(w = w_t^{(i)} \mid x, y_{<t}, w_t^{(1)}, w_t^{(2)}, ..., w_t^{(i-1)})}{\gamma \cdot \underset{w}{\max} \ p_{\mathcal{M}}(w \mid x, y_{<t}, w_t^{(1)}, w_t^{(2)}, ..., w_t^{(i-1)})}\right)$

where $p_{\mathcal{M}}(w = w_t^{(i)} \mid \cdot)$ is the mixture distribution's probability for the proposed token conditioned on all previously accepted tokens, $\max_w p_{\mathcal{M}}(w \mid \cdot)$ is the maximum probability the mixture assigns to any token at this position (i.e., the probability of the most likely token under $\mathcal{M}$), and $\gamma \in (0, 1]$ is the relaxation factor (called "leniency" in Leviathan et al., 2023).

Operational meaning of the acceptance criterion: The numerator is the mixture model's assessed probability of the proposed token given the full generation context. The denominator is $\gamma$ times the probability of the single most likely token at this position. If the proposed token is the most likely token (or close to it), the ratio is high and the token is likely accepted. If the proposed token is implausible under $\mathcal{M}$, the ratio is low and the token is likely rejected. The relaxation factor $\gamma$ makes the test more permissive: when $\gamma < 1$, the denominator is reduced, making acceptance more likely. At $\gamma = 1$, the test is strictest (acceptance only if the proposed token is exactly the most likely or very close). At $\gamma \rightarrow 0$, everything is accepted (degenerate, no filtering).

The procedure is sequential within the span: tokens are evaluated one at a time, and the first token that fails acceptance (deterministically, by checking if the acceptance probability exceeds 0.5—a simplification of the stochastic acceptance in standard speculative decoding) causes truncation. All tokens from that position to the end of the n-gram are discarded, and a new token $w_t^{(i)}$ is sampled from $p_{\mathcal{M}}$ conditioned on the accepted prefix, without going through span selection (to prevent infinite loops). If all $n$ tokens are accepted, Nest fetches the next $n$ tokens from the same corpus passage and continues the speculative evaluation, removing the reliance on the hyperparameter $n$ and enabling arbitrarily long copied spans.

Why not standard speculative decoding: Standard speculative decoding requires knowing $q(w)$, the proposal distribution, for every token in the draft. For $w_t^{(1)}$, this is known: it's p_{\text{k-NN}}(w=w_t^{(1)} \mid x, y_{<t}), since the first token was selected based on the retrieval distribution. But for subsequent tokens, the corpus continuation has no associated probability under any known distribution. The relaxation replaces $q(w)$ with $\max_w p_{\mathcal{M}}(w)$, which is always an overestimate (since $\max_w p_{\mathcal{M}}(w) \geq p_{\mathcal{M}}(w)$ for any specific $w$, and the true $q(w)$ would likely be smaller). This makes the acceptance test more conservative (it accepts less often than if the true $q(w)$ were known), but it errs on the side of quality—it prefers to reject and resample from $\mathcal{M}$ rather than accept tokens that might degrade coherence.

Why $\gamma < 1$: The relaxation factor compensates for the conservative overestimation in the denominator. Without it, the test would reject almost all tokens because the denominator (the max probability) is typically much larger than the numerator (the probability of one specific token). $\gamma$ effectively shrinks the denominator, making it possible for plausible tokens to be accepted. The specific values used: $\gamma = 5 \times 10^{-4}$ for the 7B model, $\gamma = 5 \times 10^{-3}$ for the 13B model, and $\gamma = 5 \times 10^{-2}$ for the 70B model on Wikipedia tasks. Larger models use larger $\gamma$, meaning they are stricter about acceptance—the paper observes that "as the model gets stronger, using larger $\gamma$ which leads to more rejection, is more beneficial to generation quality." Stronger models can be more selective because their own parametric distribution is more reliable.

Why accept deterministically at 0.5 threshold: Standard speculative decoding accepts with probability equal to the ratio (stochastic acceptance), which preserves the property that the output distribution is exactly $p_{\text{target}}$. Nest uses deterministic acceptance at threshold 0.5, which is faster (no random sampling needed) but sacrifices the exact distribution-matching guarantee. The paper acknowledges this is an approximation: it's "a relaxed version of speculative decoding that upper bounds the acceptance probability." In practice, greedy decoding is used throughout the experiments, so the distribution-matching property is less critical than for stochastic sampling.

Accepted span masking: Once an n-gram is accepted, the corresponding text in the corpus is masked (marked as used) and will not be retrieved again for the remainder of the generation. This prevents the $k$-NN component from repeatedly retrieving the same passage, which would be a particular risk since the corpus only contains the top-$b$ retrieved passages (a small subset of the full corpus). Without masking, the model could get stuck in a loop, repeatedly retrieving and copying the same text.

---

Complete Algorithm Walkthrough: One Inference Step

To make the interaction of all components concrete, here is what happens at a single inference step $t$ during generation, following Algorithm 1 in Appendix A:

Initialization (done once per query):

1. The input prompt $x$ is passed to the first-stage retriever $R$, which queries both the dense (Dragon+) and sparse (BM25) indexes, retrieves 4,000 passages from each, fuses their scores using the adaptive $\eta$, and returns the top-$b=40$ passages $\{d_1, ..., d_{40}\}$.
2. Each passage is tokenized and encoded by the LM's $f(\cdot)$ to produce a key-value store $(\mathcal{K}', \mathcal{V}')$ containing approximately 40 × 200 = 8,000 (key, next-token) pairs. This store is rebuilt for each generation step (though the implementation may cache the encoding of the retrieved passages since $b$ passages are fixed for the query, only $q_t$ changes per step).

Per-step generation loop:

1. Query encoding: The LM processes the current prefix $y_{<t}$ and produces the hidden state $q_t = f(y_{<t})[-1]$ (the last hidden state). Simultaneously, the LM's output layer produces the parametric distribution $p_{\text{LM}}(w \mid x, y_{<t})$.

2. Token search: $q_t$ is used to search $(\mathcal{K}', \mathcal{V}')$ for the top-$r=1024$ nearest neighbors using negative squared $\ell_2$ distance. The scores $\{s_1, s_2, ..., s_{1024}\}$ and values $\{v_1, v_2, ..., v_{1024}\}$ are returned.

3. Non-parametric distribution computation: $p_{\text{k-NN}}(w \mid x, y_{<t})$ is computed as a softmax over the similarity scores, with each neighbor's similarity accumulated for its corresponding value token. Only tokens appearing in $\pi$ have non-zero probability.

4. Confidence-based interpolation: The minimum and maximum absolute similarities among the top-$r$ are used to compute $\lambda_t$ via the RRC formula. The mixture distribution is formed: $p_{\mathcal{M}}(w) = \lambda_t \cdot p_{\text{LM}}(w) + (1 - \lambda_t) \cdot p_{\text{k-NN}}(w)$.

5. Dynamic span selection: The next token $w_t$ is greedily selected (argmax of $p_{\mathcal{M}}$). Among the retrieved neighbors with value $w_t$, the one with maximum $p_{\text{k-NN}}$ probability is identified. If $\lambda_t \leq \delta = 0.5$, the n-gram continuation of length up to 64 is extracted from the corpus; otherwise, only $w_t$ is output and the step ends here.

6. Relaxed speculative decoding: For $i = 1$ to the proposed span length:

   • The LM computes $p_{\mathcal{M}}(w \mid x, y_{<t}, w_t^{(1)}, ..., w_t^{(i-1)})$. This requires a forward pass through the LM for each new token position, but these can potentially be batched or parallelized.
   • The acceptance probability is computed as the ratio of $p_{\mathcal{M}}(w = w_t^{(i)})$ to $\gamma \cdot \max_w p_{\mathcal{M}}(w)$.
   • If the acceptance probability $\leq 0.5$, the token is rejected. All tokens from position $i$ to $n$ are discarded. A replacement token is sampled (or greedily selected) from $p_{\mathcal{M}}$ and the speculative decoding loop ends.
   • If the acceptance probability $> 0.5$, the token is accepted and evaluation continues to $i+1$.
   • If all $n$ tokens are accepted, the next $n$ tokens are fetched from the same corpus passage and evaluation continues (this can repeat indefinitely).

7. Append and mask: The accepted tokens (possibly a truncated span plus a resampled token) are appended to $y_{<t}$. The source passage for any accepted span is masked in the corpus to prevent re-retrieval.

8. Continue to the next step $t+1$, where a new $(\mathcal{K}', \mathcal{V}')$ is built (or updated) based on the first-stage retrieval results, and the process repeats from step 1.

---

Design Rationale and Hyperparameter Summary

The paper's design choices reflect a consistent philosophy: use the LM's own representations and distributions as the source of truth for quality control, while using the corpus as a source of candidate text. The LM is trusted to know when the retriever is useful (RRC), to select the best matching corpus context (max-pooling), and to reject corpus text that doesn't fit (relaxed speculative decoding). The corpus is treated as a proposal source, not an authority—every span it provides must pass the LM's quality test.

Key hyperparameters and their settings across all experiments:

The $\gamma$ values are the most task- and model-sensitive hyperparameter, increasing with model size. The paper explains that larger models can afford stricter acceptance (higher $\gamma$) because their own mixture distribution is more reliable, so fewer corpus spans pass the quality threshold. The Pile of Law setting uses a very low $\gamma = 5 \times 10^{-4}$ across all models because legal text requires more aggressive copying—the exact wording of statutes and clauses matters, and even a strong model's paraphrasing may be inaccurate.

The combination with in-context RA (RA-Nest) uses an even higher $\gamma = 5 \times 10^{-1}$ (i.e., 0.5), meaning it is much stricter about accepting corpus spans. This is because the in-context passages already provide strong factual guidance through the prompt; aggressive span copying on top of that could lead to over-reliance on the corpus at the expense of the LM's own reasoning.

Why these hyperparameters were tuned in stages: The paper first tunes $\alpha$ and $\tau$ on perplexity (a language modeling task that doesn't involve span selection or speculative decoding), then fixes those and tunes $\delta$ and $\gamma$ on generation metrics. This staged approach reduces the search space and avoids the confounding effect of evaluating all hyperparameters simultaneously—a practical necessity given that each evaluation involves generating hundreds of examples with an expensive retrieval pipeline.

4. Key Insights and Innovations

Innovation 1: Reframing Corpus Retrieval as a Proposal Mechanism with LM-as-Verifier

The dominant framing in retrieval-augmented generation treats the corpus as an authoritative source of truth that the LM must faithfully incorporate. In-context RA prepends passages and trusts the LM to attend appropriately; RETRO trains specialized cross-attention to hard-wire retrieval into the architecture; even kNN-LM, despite operating at the output level, blends the retrieval distribution as an equal partner in the mixture. In all cases, the retrieved content carries an implicit presumption of correctness—if the retriever found it, the LM should use it.

Nest fundamentally inverts this relationship. The corpus is demoted from authority to proposal source—it suggests candidate text spans, but only the LM (through the mixture distribution) has the authority to accept or reject them. This is not merely a change in implementation; it is a conceptual reframing of the LM-retriever relationship. The corpus becomes analogous to the draft model in speculative decoding: fast, approximate, and valid only to the extent that the target model (here, the mixture distribution p_ℳ) endorses its outputs.

This shift resolves a tension that prior work could not escape. Wang et al. (2023a) showed that kNN-LM degrades open-ended generation quality precisely because the corpus was treated as co-equal with the LM—tokens from unrelated documents were blended into the output with no quality filter. In-context RA, by making the corpus part of the input, gives the LM no mechanism to reject specific spans while accepting others; the model must attend to the full retrieved context or ignore it wholesale. Nest's proposal-verification architecture means that the corpus can be aggressively consulted (proposing spans whenever the retriever is confident, as controlled by the threshold δ) without risking generation quality, because every proposed token must pass the LM's own quality test via relaxed speculative decoding.

The evidence for this reframing's power is in the ablation study (Appendix C, Table 3). Adding relaxed speculative decoding to the system—the component that operationalizes "LM as verifier"—improves Biography FActScore from 41.6 to 46.8, a gain that comes entirely from rejecting corpus spans that don't fit the generation context. This is a quality improvement that no amount of better retrieval could achieve, because it addresses not whether the retrieved text is factual, but whether it is contextually appropriate—a judgment only the LM can make.

Innovation 2: Retrieval Confidence as a First-Class Signal for Adaptive Interpolation

The standard kNN-LM framework uses a fixed interpolation coefficient λ—a single global hyperparameter that weights the parametric vs. non-parametric distributions identically for all tokens, all contexts, and all tasks. This is manifestly wasteful, but the field's response has been surprisingly limited: follow-up work either trained λ predictors (He et al., 2021; Bhardwaj et al., 2023), which broke the plug-and-play property, or performed extensive per-task tuning (Drozdov et al., 2022), which broke zero-shot generality. No prior work asked the simpler question: can the retrieval scores themselves tell us how much to trust the retriever?

Nest's Relative Retrieval Confidence (RRC) answers this with a diagnostic insight: the distribution of retrieval similarities encodes the retriever's internal uncertainty. When all retrieved neighbors have similar similarity scores, the retriever cannot confidently distinguish between candidate tokens—its signal should carry less weight. When one neighbor dominates the similarity distribution, the retriever is confident—its signal should carry more weight. The min-max ratio of absolute similarities captures this in a single normalized statistic that requires no training, no additional forward passes, and no per-task calibration beyond two hyperparameters (α and τ) that transfer across tasks.

The significance of this insight extends beyond the specific formula. It establishes a general design principle for semi-parametric LMs: any retrieval mechanism that produces a distribution over candidates also produces, as a byproduct, information about its own confidence, and this meta-signal should be used to modulate the retrieval's influence. This principle is invisible in the fixed-λ framework but obvious in retrospect—it is essentially the same idea as using softmax entropy to measure model uncertainty, but applied at the retrieval interface rather than within the LM.

The ablation provides quantitative support: adding RRC to a two-stage kNN-LM baseline improves NQ answer-level recall from 40.8 to 44.4 and Biography FActScore from 34.8 to 41.6 (Table 3). These gains come purely from dynamic allocation of trust between the LM and retriever, before any span-copying or speculative decoding is applied. The sensitivity analysis (Figure 3b) further shows that the method is not brittle—perplexity is relatively flat across a range of α and τ near the chosen values—which is exactly what one hopes for in a heuristic that must work across diverse tasks without per-task tuning.

Innovation 3: Span-Level Copying Without Training via Approximate Speculative Decoding

The Copy Generator (CoG; Lan et al., 2023) established that copying multi-token phrases from a corpus—rather than retrieving tokens individually—improves both attribution quality and generation coherence, because consecutive tokens from the same source inherit the source's internal consistency. But CoG required joint training of a phrase encoder and the LM, making it incompatible with off-the-shelf models. This created an apparent dilemma: either accept the training cost and architectural modification of CoG-like approaches, or accept the degradation that token-level kNN-LM produces on open-ended text (Wang et al., 2023a).

Nest resolves this dilemma through a key conceptual move: treating the corpus as a draft source for speculative decoding. The speculative decoding framework (Leviathan et al., 2023) was developed for a completely different purpose—accelerating inference by using a small draft model to propose tokens that a large target model verifies—but Nest recognizes that exactly the same verification logic can validate corpus-extracted spans. The corpus proposes an n-gram; the mixture distribution (LM + retrieval) verifies it token by token; the first implausible token truncates the span.

This is not an obvious adaptation. Standard speculative decoding requires knowing the draft model's output distribution q(w) to compute correct acceptance probabilities. For corpus-extracted spans beyond the first token, no such distribution exists—the corpus is not a probabilistic model with known output probabilities. Nest's relaxed acceptance criterion—replacing the unknown q(w) with γ · max_w p_ℳ(w)—is a theoretical compromise that sacrifices the exact distribution-matching guarantee of standard speculative decoding in exchange for applicability to non-probabilistic draft sources. The paper does not prove that this relaxation preserves any particular property, but the empirical results (Table 1, across nine benchmarks) demonstrate that in practice, the relaxed procedure effectively filters inappropriate corpus text.

The conceptual significance is that Nest decouples span-copying from training. The speculative verification step replaces what CoG achieves through learned copy mechanisms—determining whether a phrase from the corpus should be incorporated into the output—with inference-time computation. This is philosophically aligned with recent trends toward test-time compute scaling: when more computation at inference can substitute for more training, the tradeoff favors methods that work with off-the-shelf models.

The attribution results (Table 2) demonstrate the practical value: depending on the model and task, 33.2% to 95.5% of generated tokens can be traced to a specific corpus passage, with average attributed spans of 3.0 to 27.9 consecutive tokens. This is span-level attribution—a complete factual claim like "It was released in August 1993 as the first single from Reba's compilation album Greatest Hits Volume 2" can be traced to a single source—which neither standard kNN-LM (token-level attribution, fragmented across sources) nor in-context RA (passage-level attribution, no guarantee of faithful copying) can provide.

Innovation 4: The Dual-Use Nature of Nest as Simultaneously an Attribution Method and an Acceleration Method

Most work on LLM inference treats quality improvement and latency reduction as independent optimization axes—you can make the model better (retrieval augmentation, chain-of-thought, self-consistency) or you can make it faster (speculative decoding, quantization, pruning), and these efforts proceed in parallel research communities with different methods and assumptions. Nest is unusual, perhaps unique among current methods, in achieving both simultaneously: the same span-copying mechanism that provides direct attribution also generates multiple tokens per inference step, reducing the number of sequential LM forward passes and thus accelerating generation.

Figure 2 makes this dual benefit concrete. The latency breakdown (Figure 2a) shows that even the slowest Nest configuration on the 70B model—with full retrieval overhead including passage search, token index building, and speculative verification—is faster than the base LM generating tokens one at a time. The mechanism is visible in the average proposed span length: as the relaxation factor γ decreases (making acceptance more lenient), the average span length increases from ~5 tokens to ~35 tokens per step, directly reducing the number of inference steps. And crucially, the accuracy (FActScore) does not monotonically degrade with leniency—it peaks at γ = 5×10⁻², where span length is moderate and generation speed is already substantially improved.

This is not a minor engineering optimization. It represents a fundamental insight about the relationship between retrieval and latency: when the retriever can locate coherent spans in the corpus, those spans amortize not only the retrieval cost (one search retrieves many tokens) but also the LM inference cost (one forward pass verifies many tokens). The 1.8× speedup on the 70B model is not an upper bound—it reflects a specific operating point (γ = 5×10⁻² on Wikipedia tasks), and more aggressive leniency could yield even larger speedups at the cost of quality (the tradeoff visible in Figure 2b).

The practical implication is that for practitioners deciding between retrieval augmentation and speculative decoding—two currently separate toolchains—Nest suggests they are not alternatives but can be synergistically combined. The corpus serves double duty: it provides factual grounding (reducing hallucination) and provides draft tokens (reducing latency). This unification is likely to influence how deployed LLM systems are architected, since it means a single infrastructure investment (corpus indexing, retrieval pipeline) yields improvements on both fronts.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The paper evaluates on nine benchmarks spanning four task categories (Section 4.1): (1) Text completion on WikiText-103 (Merity et al., 2017) and Pile of Law (Henderson et al., 2022), with test sets derived from the HuggingFace splits; (2) Question answering on 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); (3) Fact verification on a biography-generation task (Min et al., 2023) evaluated with FActScore, and TruthfulQA (Lin et al., 2022); (4) Closed-set tasks on MMLU (Hendrycks et al., 2021) across 57 subjects. For language modeling and text completion, 128-token prefixes are provided and 256 consecutive tokens are generated as targets; for question answering, the maximum generation length is 128 tokens; for fact verification, it is 512 tokens. Dev sets from WikiText-103, NQ, and Biography are used for hyperparameter tuning (Appendix B).

• Base model(s). All experiments use Llama-2-Chat models (Touvron et al., 2023b) at three scales: 7B, 13B, and 70B parameters. The instruction-tuned ("chat") variants are used because the paper focuses on evaluating instruction-following behavior under zero-shot conditions, with all few-shot demonstrations removed from prompts to simulate realistic deployment. The models span roughly an order of magnitude in parameter count, enabling assessment of how Nest's benefits scale with base model capability — a critical axis given that retrieval augmentation is known to benefit smaller models more (Borgeaud et al., 2022).

• Metrics. Six distinct metrics are reported depending on the task (Section 4.1, Table 1). For language modeling: perplexity (PPL). For text completion: ROUGE-1, ROUGE-2, ROUGE-L (Lin, 2004) measuring n-gram overlap with reference text; MAUVE (Pillutla et al., 2021) measuring distributional similarity to human text; and average generation length. For question answering: answer-level recall (Hit@1; Karpukhin et al., 2020), which checks whether the generated output contains any correct answer string rather than requiring exact match — appropriate for the zero-shot free-form generation setting where the model may produce verbose answers. For TruthfulQA: ΔBLEU and ΔROUGE-1, defined as the difference between the maximum similarity to true reference answers and the maximum similarity to false reference answers (following Lin et al., 2022). For biography: FActScore (Min et al., 2023) with length penalty, computed by decomposing generated text into atomic facts and verifying each against a knowledge source using a retrieval+llama+npm pipeline. For MMLU: macro accuracy averaged across 57 subjects, grouped into STEM, Social Sciences, Humanities, and Other. For attribution analysis (Section 4.6): the proportion of generated tokens traceable to the corpus and the average length of consecutive attributed spans.

• Baselines. Five baselines are compared (Section 4.3): (1) Base LM — the unmodified Llama-2-Chat model with no retrieval, representing the parametric-only lower bound; (2) In-Context Retrieval Augmentation (RA) — the most widely used approach, where the top-3 retrieved passages are prepended to the prompt as background context (Ram et al., 2023; Shi et al., 2024a); (3) Two-Stage kNN-LM — the standard kNN-LM (Khandelwal et al., 2020) with the same two-stage retrieval architecture as Nest (Section 3.1, top-40 passages, fixed interpolation coefficient λ, no span selection or speculative decoding), using λ = 0.7 for Wikipedia-based tasks and λ = 0.9 for Pile of Law; (4) Nest — the full method as described; (5) RA-Nest — Nest combined with in-context retrieval augmentation, where retrieved passages are added to the prompt and Nest operates on top. The first-stage retriever (Dragon+ plus BM25 with fusion) is shared by kNN-LM, Nest, and RA-Nest, ensuring that any performance differences are due to the inference-time mechanisms rather than retrieval quality.

• Generation budget / compute accounting. Generation uses greedy decoding throughout (no sampling, temperature 0) to eliminate randomness as a confounding factor in factuality evaluation. For latency measurements (Section 4.5), compute is measured in wall-clock time on 8× A100 GPUs for model parallelization and 32 CPU threads for search, with batch size 1. The base LM implementation is described as "internal, research-purpose" and "not optimized for latency," meaning absolute latency numbers are not directly comparable to production-optimized inference engines — the relevant comparison is relative speedup (1.8×) rather than absolute milliseconds per token. The cost of retrieval is included in latency measurements: the breakdown in Figure 2a shows passage search and token index building as separate components, which together account for roughly 25-30% of total latency. This is a strength of the evaluation — the speedup claim includes retrieval overhead, not just generation speed — though the fixed cost of maintaining the dense and sparse indexes (8.96 GB + 3.48 GB on disk, loaded into RAM) is not factored into latency.

• Cross-validation / statistical protocol. No formal cross-validation or statistical significance testing is reported. Hyperparameters are tuned on dev sets of WikiText-103, NQ, and Biography, then applied uniformly to all other tasks without per-task adjustment. The paper explicitly states in the NeurIPS checklist (item 7): "This paper does not include significant tests considering the performance gap between the proposed approach and the baselines." This is a meaningful limitation: with only 500 test questions for some datasets and no confidence intervals, the reliability of small performance differences (e.g., Nest-70B achieving 45.2 answer-level recall on the QA average vs. 44.9 for kNN-LM-70B, a 0.3-point gap on what is likely 100-500 questions) cannot be assessed. The ablation study (Table 3) uses validation sets, not test sets, which is appropriate for component analysis but means the reported improvements are measured on data that may have influenced hyperparameter selection.

Main Quantitative Results

Text Completion and Language Modeling Results

Table 1 (upper) reports results on WikiText-103 and Pile of Law. The headline findings:

Perplexity: RA-Nest achieves the lowest perplexity across all model sizes on both datasets. On WikiText-103 with the 70B model, RA-Nest achieves 4.8 PPL compared to 9.9 for the base LM (a 51.5% reduction), 5.3 for RA alone, and 7.1 for kNN-LM. The pattern holds at all model scales: RA-Nest consistently outperforms both RA alone and Nest alone, suggesting that input-level and output-level retrieval augmentation address complementary failure modes. On Pile of Law, the same ranking holds (70B: RA-Nest 4.7 vs. base 6.9), though the absolute perplexity numbers are lower across the board, likely reflecting the narrower domain of legal text.

ROUGE scores for text completion: Here the picture is more nuanced. On WikiText-103, RA alone achieves the best ROUGE scores (e.g., 70B: RA achieves ROUGE-1 of 40.5 vs. Nest's 32.6), while Nest underperforms RA but substantially outperforms kNN-LM and the base LM. However, on Pile of Law, RA-Nest achieves the best ROUGE scores for the 7B and 13B models, and competitive scores for the 70B model. The paper's interpretation (Section 4.4) is that "for legal documents, quoting the exact clauses from the source might be more favourable compared to Wikipedia" — legal text rewards verbatim copying, which Nest's span mechanism provides, while Wikipedia text rewards paraphrasing, which RA's prompt-based approach enables.

MAUVE scores: RA and RA-Nest consistently achieve the highest MAUVE scores, indicating that retrieval augmentation improves the distributional match to human-written text. Nest alone sometimes underperforms the base LM on MAUVE (e.g., 70B on WikiText-103: Nest 82.6 vs. base 88.6), which the paper does not explicitly discuss but likely reflects the tradeoff between factual accuracy and stylistic fluency — verbatim corpus spans may reduce MAUVE if they differ stylistically from typical Wikipedia prose.

Average generation length: Nest and kNN-LM tend to produce longer outputs than the base LM (e.g., 70B on WikiText-103: Nest 236.3 tokens vs. base 239.6, but on Pile of Law: Nest 251.3 vs. base 250.1). RA-Nest sometimes produces shorter outputs on WikiText-103 (70B: 233.1 vs. 235.9 for RA alone), which the paper attributes to the span-copying mechanism being more concise than parametric generation.

Question Answering Results

Table 1 (lower left) reports answer-level recall across NQ, TQA, HQA, and MQA. The headline findings:

RA and RA-Nest dominate, but Nest outperforms kNN-LM consistently. Across all four QA datasets and all model sizes, RA achieves the highest or tied-for-highest answer-level recall. For the 70B model, RA achieves 75.5 (TQA), 55.4 (NQ), 52.5 (HQA), and 16.0 (MQA), with a four-dataset average of 49.9. RA-Nest performs essentially identically on the 70B model (average 49.8), slightly better on the 13B model (45.7 vs. 45.3 for RA alone), and comparably on the 7B model (44.1 vs. 43.7). The gap between Nest and kNN-LM varies by dataset and model size: on the QA average, Nest holds a consistent advantage over kNN-LM (70B: 45.2 vs. 44.9; 13B: 38.4 vs. 38.4 tie; 7B: 37.1 vs. 37.2, essentially tied).

The benefit of retrieval diminishes with model scale. This matches prior findings (Borgeaud et al., 2022). On TQA, the base 7B model achieves 61.1, RA achieves 69.5 (+8.4), while the base 70B model achieves 74.0, RA achieves 75.5 (+1.5). The gap between base LMs and retrieval-augmented methods narrows substantially as model size increases, suggesting that larger models have already internalized more of the knowledge present in the corpus. This is an important boundary condition for Nest's practical value: the largest gains are on smaller models, which is where retrieval augmentation is most cost-effective anyway (since larger models are proportionally more expensive to serve).

MedMCQA shows the lowest absolute scores. On MedMCQA, even the best configuration (RA-Nest 70B) achieves only 16.3 answer-level recall. This dataset involves medical domain knowledge that may be poorly covered in the Wikipedia corpus used for retrieval, highlighting the dependence of all retrieval-augmented methods on corpus coverage.

Fact Verification Results

Table 1 (lower middle) reports TruthfulQA and Biography results. The headline findings:

On TruthfulQA, semi-parametric methods outperform RA, and RA can be harmful. On TruthfulQA, measured by ΔBLEU and ΔROUGE-1 (higher is better, indicating the model's output is more similar to true references than false references), Nest and kNN-LM consistently outperform base LMs and RA. For the 13B model: Nest achieves ΔBLEU of 0.29 and ΔROUGE of 0.98, compared to base LM at 0.13 and 0.81, and RA at -0.16 and 0.25. RA's negative ΔBLEU on TruthfulQA across all model sizes (7B: -0.34, 13B: -0.16, 70B: -0.13) is a striking finding: prepending retrieved passages increases the model's tendency to generate false answers on this adversarial dataset. The paper explains this as TruthfulQA containing questions "where in-context RA is more susceptible to the 'evidence' in the prompt (e.g., astrology and myths)." In contrast, Nest and kNN-LM, which operate at the output level rather than the prompt level, are less susceptible to misleading retrieved content. RA-Nest shows negative ΔBLEU as well, indicating that the in-context RA component's harmful influence dominates.

On Biography, RA achieves the highest FActScore, but Nest substantially outperforms kNN-LM. For the 70B model, RA achieves FActScore of 52.9, compared to 41.6 for Nest and 36.1 for kNN-LM. However, Nest's advantage over kNN-LM is substantial and consistent: 70B: 41.6 vs. 36.1 (+5.5); 13B: 35.7 vs. 31.1 (+4.6); 7B: 38.9 vs. 30.6 (+8.3). This is direct evidence that Nest's mechanisms (dynamic interpolation, span selection, relaxed speculative decoding) meaningfully improve over standard kNN-LM on factuality. RA-Nest achieves 59.2 on the 70B model, outperforming RA alone (52.9), suggesting that span-level copying can add factual precision on top of what prompt-level retrieval provides. The degradation for RA-70B relative to RA-13B (52.9 vs. 59.1) is attributed to the 70B model generating shorter claims, which are penalized by the FActScore length penalty.

The number of facts generated varies substantially. The "Facts" column shows that RA-Nest 70B generates 53.8 facts on average, while RA alone generates 42.1, and Nest alone generates 56.2. Nest's longer outputs (more facts) combined with moderate FActScore suggests it generates more claims but with imperfect factuality; RA generates fewer claims but with higher precision.

Closed-Set Task Results (MMLU)

Table 1 (lower right) reports MMLU macro accuracy. The headline findings:

RA, Nest, and RA-Nest all cluster within a few points of each other. For the 70B model, the average MMLU scores are: base LM 43.2, RA 45.1, kNN-LM 43.2, Nest 43.5, RA-Nest 45.1. The differences are small: RA and RA-Nest achieve identical average scores at 70B, Nest edges ahead of kNN-LM by 0.3 points. The breakdown by subject category shows RA and RA-Nest performing best across STEM (45.9/45.8 vs. 43.5 for base), Social Sciences (39.7 vs. 37.9), Humanities (46.2 vs. 44.4), and Other (48.6/48.9 vs. 47.0). Nest outperforms the base LM in all categories but by margins of 0.3-0.8 points — modest but consistent.

Why MMLU gains are small: MMLU is a multiple-choice task evaluated by comparing perplexity of each option concatenated with the question. Since Nest's span selection and speculative decoding mechanisms only apply during generation, and MMLU evaluation uses perplexity scoring rather than free-form generation, Nest's benefits are limited to the interpolation mechanism (RRC) for this task. The small gains over the base LM (43.5 vs. 43.2 for 70B) reflect the contribution of dynamic interpolation alone, without span-copying or speculative verification.

Latency Results

Figure 2 reports latency measurements on the Biography validation set using the Nest-70B model (specific configuration: α=0.3, τ=0.1, δ=0.5). The headline findings:

Latency breakdown (Figure 2a): The largest component is LM encoding (approximately 45-50% of total latency), followed by speculative decoding (variable, depending on γ), with passage search and token index building accounting for a relatively constant fraction. As γ decreases (more lenient acceptance), the speculative decoding time decreases because fewer tokens are rejected and resampled, while the LM encoding time remains relatively constant per step. The base LM latency (dashed line) is higher than all Nest configurations shown, confirming the speedup.

Span length vs. accuracy tradeoff (Figure 2b): As γ decreases from 10⁻¹ to 10⁻⁴, the average proposed span length increases from approximately 5 tokens to approximately 35 tokens per step. The FActScore on Biography shows a non-monotonic relationship: peak FActScore occurs at γ = 5×10⁻², with lower FActScore at both higher γ (too much rejection, not enough corpus copying) and lower γ (too little rejection, poor-quality spans accepted). This validates the existence of a "sweet spot" where both low latency and high accuracy are achieved simultaneously. The paper selects γ = 5×10⁻² for the 70B model in main experiments, which corresponds to the FActScore peak in Figure 2b.

The 1.8× speedup claim: The paper states "for Llama-2-Chat 70B, it achieves a 1.8× speedup in inference time." Based on Figure 2a, this appears to compare the Nest configuration at γ = 5×10⁻² against the base LM. The exact computation is not specified in the text, but the latency breakdown shows that total Nest time is roughly 55-60% of the base LM time, consistent with a ~1.8× speedup.

Attribution Results

Table 2 reports attribution analysis on NQ and Biography tasks. The headline findings:

High attribution ratios, with strong model-size dependence. On NQ, the proportion of generated tokens traced to the corpus ranges from 52.4% (Nest-13B) to 93.4% (Nest-7B). The 7B model shows substantially higher attribution ratios than the 13B and 70B models for Nest alone (93.4% vs. 52.4% vs. 58.8%), which is consistent with the finding that retrieval benefits smaller models more — the 7B model relies more heavily on the corpus, so more of its output is directly copied. RA-Nest configurations show lower attribution ratios (33.2-77.5%), likely because the in-context passages provide factual guidance that the model can paraphrase rather than copy verbatim.

Average attributed span length varies widely. The longest attributed spans occur with Nest-7B on Biography (27.9 tokens) and NQ (18.4 tokens), indicating that the system can copy entire factual claims as coherent units. The shortest attributed spans occur with RA-Nest configurations (3.0-5.9 tokens), where the combination of prompt-level and output-level retrieval may lead to more fragmented copying. The qualitative examples in Table 2 show green-highlighted spans corresponding to source passages, with the NQ example demonstrating a complete factual statement ("It was released in August 1993 as the first single from Reba's compilation album Greatest Hits Volume 2") attributed to a single source.

No baseline comparison for attribution: The paper claims that "neither of the baselines can achieve the same granularity and preciseness for the attribution as Nest" — kNN-LM can attribute individual tokens but not coherent spans, and in-context RA can attribute to passages but not to specific spans within passages. However, no quantitative attribution comparison against baselines is provided. This is a meaningful gap: without measuring attribution ratios for RA or kNN-LM, the claim that Nest provides superior attribution is qualitative rather than quantitative.

Ablation Studies and Robustness Checks

The paper includes a progressive ablation study (Appendix C, Table 3) and sensitivity analyses for key hyperparameters (Appendix C, Figures 3a-c).

Starting point: The ablation begins with "kNN-LM (two-stage)" — the standard kNN-LM framework applied with the same two-stage retrieval architecture as Nest (top-40 passages, top-1024 tokens), using a fixed interpolation coefficient. On the validation sets of WikiText-103, NQ, and Biography, this baseline achieves ROUGE-1 of 20.1, answer-level recall of 40.8, and FActScore of 34.8.

Adding Relative Retrieval Confidence: This replaces the fixed λ with the dynamic λt from the RRC formula. WikiText-103 ROUGE-1 improves from 20.1 to 24.7 (+23%), NQ ALR from 40.8 to 44.4 (+8.8%), and Biography FActScore from 34.8 to 41.6 (+19.5%). This is the single largest component-level improvement, confirming that dynamic interpolation is substantially more effective than a fixed λ. The gain on Biography (+6.8 FActScore) is particularly notable — the dynamic coefficient is doing meaningful work in determining when to trust the retriever for factual claims.

Adding Dynamic Span Selection: This adds the n-gram extraction mechanism (δ=0.5, n=64). The metrics remain essentially flat: WikiText-103 ROUGE-1 goes from 24.7 to 24.5, NQ ALR from 44.4 to 44.6, Biography FActScore from 41.6 to 41.6. The paper acknowledges this: "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." This is an important interpretive point — span selection does not improve aggregate metrics, but it changes how the output is generated, enabling coherent multi-token attribution that per-token retrieval cannot provide. The metrics used (ROUGE, ALR, FActScore) may not capture attribution quality directly.

Adding Relaxed Speculative Decoding: This adds the full speculative verification procedure with relaxation factor γ (values unspecified for this ablation, but presumably the same as main experiments). WikiText-103 ROUGE-1 improves from 24.5 to 26.8 (+9.4%), NQ ALR from 44.6 to 45.4 (+1.8%), and Biography FActScore from 41.6 to 46.8 (+12.5%). The large gain on FActScore is the key result: speculative decoding is filtering out corpus spans that would degrade factual accuracy, and this filtering is particularly important on factuality-focused tasks. The smaller gain on NQ ALR suggests that for question answering, the answer is often captured by the first token of the span, and subsequent tokens matter less for the recall metric.

Sensitivity to number of retrieved passages and tokens (Figure 3a): On WikiText-103 validation perplexity using the Nest-7B model, increasing the number of passages from 10 to 80 at fixed token count reduces perplexity by approximately 0.5-1.0 per doubling, with diminishing returns beyond 40. Increasing the number of tokens from 256 to 2048 at fixed passage count produces a similar pattern. The chosen operating point (40 passages, 1024 tokens) is near the knee of both curves, balancing accuracy and latency. The paper notes that "as Nest needs to encode the retrieved passages on the fly, the latency also increases linearly w.r.t. the number of passages" — hence 40 is chosen as the practical tradeoff.

Sensitivity to interpolation hyperparameters α and τ (Figure 3b): On WikiText-103, perplexity is relatively flat across a range of τ (0.01 to 1.0) and α (0.0 to 0.8). At high τ (1.0), the choice of α has minimal impact because the sigmoid is near-uniform. At low τ (0.01), the sigmoid is near-binary, and the sweet spot is around α=0.4. The chosen values (α=0.3, τ=0.1) sit in a region where perplexity is near-optimal and the transition is sharp but not binary.

Sensitivity to span selection threshold δ (Figure 3c): On NQ answer-level recall, δ = 0.0 (never copy spans) achieves the lowest recall, δ = 0.5 achieves the peak, and δ = 1.0 (always copy spans) achieves intermediate recall. The peak at δ=0.5 means that copying spans only when the retriever is confident enough to dominate the mixture (λt ≤ 0.5) is better than either never copying or always copying. This is consistent with the intuition that span-copying should be deployed selectively.

Relaxation factor γ and model size interaction: The paper reports using larger γ for larger models on Wikipedia tasks (7B: γ=5×10⁻⁴; 13B: γ=5×10⁻³; 70B: γ=5×10⁻²) and observes that "as the model gets stronger, using larger γ which leads to more rejection, is more beneficial to generation quality." This is a non-obvious finding: stronger models are more selective about accepting corpus spans, not less. The interpretation is that larger models have more reliable parametric knowledge, so the quality threshold for accepting external text should be higher. For Pile of Law, all models use γ=5×10⁻⁴, suggesting that legal text specifically benefits from more aggressive copying regardless of model size, likely because exact statutory language matters.

Critical Assessment

Claim: Nest improves generation quality over base LMs and kNN-LM. The evidence supports this, but with important nuance about which quality metrics improve and by how much. On perplexity, Nest consistently outperforms kNN-LM and the base LM (Table 1, upper). On ROUGE scores for text completion, Nest outperforms kNN-LM but underperforms in-context RA on WikiText-103, while being competitive on Pile of Law. On QA answer-level recall, Nest slightly outperforms kNN-LM (by 0-2 points depending on model size) but substantially underperforms RA. On FActScore, Nest substantially outperforms kNN-LM (by 4-8 points) but underperforms RA. On MMLU, the gains are minimal (0.3-0.8 points). The paper's abstract claims Nest "significantly enhances the generation quality and attribution rate," which is true for kNN-LM but overstated relative to RA — on most generation quality metrics, RA outperforms Nest by wider margins than Nest outperforms the base LM. The paper is honest about this in the results text ("Nest is able to outperform base LMs and kNN-LM's on most tasks while being on par with RA"), but the abstract's framing is stronger than the data supports. The real strength of Nest relative to RA is not quality but attribution granularity and latency, which are not measured as quality metrics in the main table.

Claim: Nest achieves 1.8× speedup in inference time on Llama-2-Chat 70B. This claim is supported by Figure 2a, but with important caveats. First, the speedup is measured on an "internal, research-purpose implementation of the base Llama-2-Chat model which did not optimize for latency." Optimized inference engines (vLLM, TensorRT-LLM, etc.) might show different relative speedups. Second, the speedup includes retrieval overhead but does not account for the fixed cost of maintaining the retrieval indexes. Third, the speedup varies with γ: at the operating point chosen for accuracy (γ=5×10⁻²), the speedup is approximately 1.8×, but at lower γ (faster but less accurate), speedups could be substantially larger. The 1.8× figure should be understood as the speedup at the quality-optimized operating point, not the maximum possible speedup.

Claim: Nest provides span-level attribution that baselines cannot match. This claim is qualitatively supported by Table 2 and the examples therein, but not quantitatively compared against baselines. The paper reports attribution ratios for Nest and RA-Nest (33.2-95.5%), but does not report any attribution metric for kNN-LM or RA. Without this comparison, the reader cannot assess how much better Nest's attribution is. kNN-LM can in principle attribute individual tokens to their retrieved sources; the claim that Nest is superior rests on the coherence of attributed spans (multiple consecutive tokens from the same source), which is demonstrated qualitatively but not measured quantitatively. A metric like "average number of source switches per 100 tokens" or "average length of contiguous spans from the same source" would strengthen this claim considerably.

Missing experiments that would strengthen the paper:

• Sampling-based generation: All experiments use greedy decoding. Nest's speculative decoding procedure has a natural interaction with sampling (the acceptance probability becomes a random variable), but this is not explored. Since many deployments use temperature sampling for diversity, understanding how Nest behaves under stochastic decoding is practically important.

• Comparison with REST (He et al., 2024): The concurrent work REST also uses a datastore for speculative decoding but does not modify the output distribution. A direct comparison on latency and quality would clarify the value of Nest's interpolation mechanism over pure acceleration. The paper discusses REST in related work but provides no empirical comparison.

• Attribution quality metrics compared across baselines: As noted above, quantitative attribution metrics for RA and kNN-LM would substantiate the attribution claim.

• Confidence intervals or statistical tests: With test sets of 500 questions and differences as small as 0.3 points (e.g., Nest vs. kNN-LM on MMLU 70B), statistical significance matters. The paper's explicit decision not to report significance tests (NeurIPS checklist item 7) is a weakness.

• Varying the corpus: All Wikipedia-based experiments use the same Wikipedia 2021 dump. Testing with different corpora (news, scientific papers, books) would probe how Nest's performance depends on corpus domain match to the task.

• Larger batch sizes for latency: The latency experiments use batch size 1. In production, batching is common, and the relative overhead of retrieval (which can be amortized across a batch) vs. speculative decoding (which processes per-sequence) would change.

Conditions under which claims hold:

• The quality improvements over kNN-LM are robust across model sizes and tasks, with the largest absolute gains on smaller models (7B) and factuality-focused tasks (Biography). The improvements over the base LM are more modest, particularly on MMLU and QA tasks where in-context RA is stronger.

• The speedup claim is conditional on the specific hardware configuration (8× A100 GPUs, 32 CPU threads) and the unoptimized baseline implementation. Relative speedups may differ on different hardware or with optimized inference stacks.

• The attribution claim is qualitatively demonstrated but not quantitatively benchmarked against alternatives. The paper convincingly shows that Nest can provide span-level attribution, but does not quantify how much more attribution it provides than kNN-LM or how much more precise that attribution is than RA.

Strengths of the evaluation:

• Nine benchmarks across four task categories provide broad coverage and reveal the task-dependence of Nest's benefits (strong on factuality, weak on MMLU, competitive on QA).
• Three model sizes reveal the scaling behavior — Nest helps smaller models more, which is consistent with prior retrieval-augmentation findings and practically useful.
• Ablation study separates the contributions of RRC, span selection, and speculative decoding, revealing that RRC and speculative decoding drive quality improvements while span selection enables attribution without harming quality.
• Latency breakdown includes retrieval overhead, making the speedup claim honest about total system cost rather than cherry-picking generation-only speed.
• TruthfulQA results show RA can be harmful, a finding that is important for practitioners and not obvious a priori.

6. Limitations and Trade-offs

Two-Stage Retrieval Is a Recall Filter That Can Drop Relevant Evidence

The assumption or constraint. Nest depends critically on the first-stage passage retriever to select the top-$b=40$ passages from which all token-level search is conducted. If a passage containing the correct factual answer is not ranked among these 40, the token-level retriever never sees it, and Nest cannot possibly copy the relevant text. The paper acknowledges this implicitly in Section 6:

"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"

but frames it as a general caveat rather than exploring when this failure occurs or how often.

The consequence. The two-stage architecture creates a hard recall ceiling: even if the token-level retrieval and speculative decoding work perfectly, the system's factual accuracy is upper-bounded by the passage retriever's recall@40. On tasks where the correct evidence passage is lexically or semantically distant from the query—complex multi-hop reasoning, questions requiring inference across documents, or queries with vocabulary mismatch—the dense+sparse hybrid retriever may fail to surface the right passage, and Nest will copy from whatever topically adjacent but factually incorrect passages happen to rank highly. This is particularly acute because Nest copies verbatim spans: when it copies from the wrong passage, it not only answers incorrectly but does so with the confidence of verbatim text, potentially making errors more convincing to users than the base LM's parametric hallucinations.

What evidence exists in the paper. The paper does not directly measure recall@40 of the passage retriever on any task. The MedMCQA results (Table 1, lower left) provide indirect evidence: even the best configuration (RA-Nest 70B) achieves only 16.3 answer-level recall on this medical QA dataset, substantially lower than other QA tasks. This is consistent with the Wikipedia corpus having poor coverage of medical knowledge, but it may also reflect the retriever's difficulty in matching medical queries to relevant Wikipedia passages—the failure mode (corpus coverage vs. retrieval failure) is not disentangled. The HotpotQA results (multi-hop QA, where the evidence is distributed across multiple passages) show a similar pattern: RA-Nest 70B achieves 52.4 vs. 41.4 for Nest alone, suggesting the passage retriever captures some but not all of the needed information. The paper provides no analysis of whether errors are due to retrieval failure (wrong passages in the top-40) or downstream processing (right passages retrieved but spans not copied correctly).

Mitigation status. Partially addressed through the hybrid dense+sparse retriever with aggressive first-stage recall (4,000 passages retrieved before pruning to 40). This "cast a wide net" strategy improves recall@40 relative to a dense-only or sparse-only system, but it is a heuristic proxy, not a guarantee. The paper does not experiment with varying $b$ (the number of passages passed to the second stage) beyond the sensitivity analysis in Appendix C (Figure 3a), which shows diminishing perplexity returns beyond 40 passages but does not measure recall of task-critical facts. Future work on better passage retrieval, query reformulation, or adaptive passage count per query could raise the recall ceiling, but the fundamental limitation—that token-level search is restricted to a small window of the corpus—is architectural, not parametric.

---

Attributed Spans Are Factual IFF the Corpus Is Factual

The assumption or constraint. Nest copies spans from a corpus and presents them as attributed evidence. The truth value of those spans is entirely determined by the corpus content. The paper acknowledges this in the Broader Impact statement (Section 8):

"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"

but frames it as a deployment concern rather than an intrinsic limitation of the method. Nest has no internal mechanism to verify the factual correctness of the corpus—the LM's speculative decoding evaluates whether a span is likely under the mixture distribution, which measures contextual fit and retrieval confidence, not truth.

The consequence. When the corpus contains false information, Nest will confidently copy and attribute it, giving users a false sense of verification. The TruthfulQA results (Table 1, lower middle) partially illustrate this: on an adversarial dataset containing common misconceptions, Nest achieves positive ΔBLEU and ΔROUGE-1 (meaning it's better than the base LM at distinguishing true from false statements), but the absolute scores remain far from perfect. More concerning is a failure mode the paper does not address: if the retriever surfaces a passage from a biased or outdated source, Nest may copy spans that are technically present in the corpus but factually wrong or misleading. Since the attribution points to a real source, users may be less likely to question the claim than if the base LM generated it without citation—the attribution creates an illusion of authority that can amplify corpus errors.

What evidence exists in the paper. The paper provides no direct measurement of this effect. The TruthfulQA results show that Nest is less susceptible to retrieval-induced falsehoods than in-context RA (which has negative ΔBLEU scores across all model sizes), but this only addresses the case where the corpus contains both true and false information on a topic and the retriever surfaces the wrong one. It does not address the case where the corpus is systematically biased or contains no correct information on a topic. The ablation study (Table 3) shows that relaxed speculative decoding improves FActScore—this means the LM is rejecting some corpus spans, which is good, but it does not measure whether the accepted spans are factually reliable. The paper's use of Wikipedia as the corpus is a best-case scenario for factual reliability; results would likely degrade with a noisier corpus (e.g., web crawl data, social media, or domain-specific corpora with known biases).

Mitigation status. Not addressed. The paper suggests filtering as a preprocessing step (Section 8), but this is external to the Nest framework and shifts responsibility to the corpus curator rather than the method. A more integrated approach—such as having the LM's own factual knowledge vote on whether a span is likely true (as opposed to likely appropriate, which is what speculative decoding currently assesses)—is not explored. The paper also does not discuss how attribution interacts with corpus errors: a user who follows the attribution link and finds the source passage contains an error has no recourse; Nest has no mechanism for flagging or qualifying uncertain attributions.

---

Latency Gains Depend on Unoptimized Baselines and Specific Hardware

The assumption or constraint. The paper's headline speedup claim—"a 1.8× speedup in inference time when applied to Llama-2-Chat 70B"—is measured against an "internal, research-purpose implementation of the base Llama-2-chat model which did not optimize for latency" (Section 4.5). The latency breakdown (Figure 2a) is collected on 8× A100 GPUs with 32 CPU threads for search, with batch size 1.

The consequence. The 1.8× figure is relative to a specific, potentially slow baseline. Production inference engines (vLLM, TensorRT-LLM, DeepSpeed-Inference) can achieve 2-5× speedups over naive implementations through kernel fusion, continuous batching, KV-cache optimization, and quantization. If the optimized baseline is, say, 3× faster than the research implementation while Nest's retrieval overhead remains constant (since passage search and index building are CPU-bound and do not benefit from GPU kernel optimization in the same way), the relative speedup of Nest over an optimized baseline could shrink substantially or even reverse. The paper does not compare against any optimized inference framework.

Additionally, the latency experiments use batch size 1. In batched inference (common in production), the relative overhead of Nest's components changes: LM encoding can be parallelized across the batch (GPU utilization increases), but the CPU-bound passage search and token index building must be performed per-sequence and may not benefit from batching to the same degree. The per-token speculative decoding, which processes each sequence's span independently, also scales linearly with batch size. Without batch-size scaling experiments, the latency advantage at realistic production batch sizes (8-64) is unknown.

What evidence exists in the paper. The paper provides only Figure 2, which measures a single hardware configuration and batch size 1. There is no comparison with optimized inference engines, no batch-size scaling analysis, and no measurement of throughput (tokens per second per GPU) vs. latency (time to complete one query). The breakdown in Figure 2a shows that LM encoding accounts for roughly 45-50% of total latency; if an optimized inference engine cuts LM encoding time by 3×, the retrieval overhead (which remains unchanged) would become a much larger fraction of total time, potentially erasing the speedup entirely. The paper's explicit caveat about the unoptimized baseline is honest, but the 1.8× figure is presented in the abstract without qualification, and its sensitivity to the baseline is never quantified.

Mitigation status. The authors are transparent about the hardware and the unoptimized baseline in Section 4.5, which is a genuine strength of the paper's reporting. However, no sensitivity analysis or optimized-baseline comparison is provided. A practitioner evaluating Nest for deployment would need to benchmark against their own optimized inference stack; the 1.8× figure should be treated as an existence proof that speedup is possible, not as a calibrated estimate for production systems.

---

No Guarantee That the LM Will Not Ignore or Override Retrieved Evidence

The assumption or constraint. Nest's mixture distribution blends the LM's parametric predictions with the retrieval distribution via the dynamic coefficient $\lambda_t$. The speculative decoding procedure then validates corpus spans against this mixture. However, when $\lambda_t$ is large (the retriever is uncertain), the LM's own distribution dominates, and the system can generate tokens that are factually inconsistent with the retrieved passages even when those passages are in the top-$b$ set. The paper acknowledges this possibility only indirectly, noting that validation loss was used for early stopping during fine-tuning.

The consequence. Nest can suffer from a failure mode that in-context RA is explicitly designed to avoid: the LM overriding retrieved evidence with its own parametric knowledge. In in-context RA, the retrieved evidence is prepended to the prompt, forcing the LM to attend to it (though, as the TruthfulQA results show, attention does not guarantee compliance). In Nest, when RRC determines that the retriever is uncertain—which can happen for many reasons, including genuinely ambiguous queries, poor retrieval quality, or the LM's own hidden state being a poor match to any corpus context—the system defaults to the LM's parametric distribution. If the LM has internalized incorrect facts, those will be generated despite the corpus containing correct information. This is the mirror image of the previous limitation: Nest can fail either by copying from a factually wrong corpus (Limitation 2) or by ignoring a factually correct corpus (this limitation), and the RRC mechanism decides which regime applies at each token without any direct truth assessment.

What evidence exists in the paper. The RRC sensitivity analysis (Appendix C, Figure 3b) shows that the chosen $\alpha=0.3$, $\tau=0.1$ settings produce a relatively sharp transition: at most steps, $\lambda_t$ is either close to 0 (retriever dominates) or close to 1 (LM dominates), with relatively few steps in the ambiguous middle ground. This means that on a substantial fraction of tokens, the LM is operating nearly independently of the corpus. The paper does not report the distribution of $\lambda_t$ values across tasks, so the reader cannot assess how often the LM overrides the retriever on fact-critical tokens. The attribution ratios in Table 2 (33.2% to 95.5% of tokens traced to the corpus) provide an indirect measure: the complement (4.5% to 66.8% of tokens) are generated parametrically and have no corpus grounding, but this metric does not distinguish between tokens where the corpus had no relevant information (acceptable) vs. tokens where the corpus had correct information that was overridden (problematic).

Mitigation status. Not addressed. The RRC mechanism is designed to modulate trust based on retriever confidence, not on factual correctness of the LM vs. the corpus. A more principled approach would incorporate the LM's own uncertainty (entropy of $p_{\text{LM}}$) into the interpolation decision, but this is not explored. The paper also does not investigate whether accepted spans (those passing speculative decoding) sometimes contradict the LM's parametric knowledge—this would be a direct measure of the override problem.

---

Hard Problems and Adversarial Settings Expose Fundamental Limits

The assumption or constraint. Nest's mechanisms assume that when the retriever is confident and the corpus contains relevant text, incorporating that text as verbatim spans will improve generation quality. This assumption breaks down in two important regimes: (1) tasks where the correct answer requires reasoning across multiple passages rather than extracting a span from a single passage (multi-hop reasoning, synthesis), and (2) adversarial settings where retrieved passages are designed to mislead.

The consequence. For multi-hop reasoning tasks (exemplified by HotpotQA in the paper's benchmark suite), Nest must assemble facts from multiple retrieved passages into a coherent answer. But the span-copying mechanism copies consecutive text from single passages—it has no mechanism for synthesizing information across passages. The LM's parametric component can perform cross-passage reasoning in principle (since $p_{\text{LM}}$ is conditioned on the full generation prefix), but the span-copying component, which provides the attribution, is inherently single-passage. This creates a tension: the attributed parts of the output are restricted to what single passages contain, while any cross-passage reasoning must be performed parametrically and is therefore unattributed.

For adversarial settings (exemplified by TruthfulQA), the paper's own results show that in-context RA produces negative ΔBLEU scores because the LM is "more susceptible to the 'evidence' in the prompt (e.g., astrology and myths)" (Section 4.4). Nest performs better than RA on TruthfulQA (positive ΔBLEU), suggesting that output-level interpolation is less susceptible to adversarial retrieval than prompt-level augmentation. However, the paper does not test adversarial retrieval explicitly—TruthfulQA is adversarial at the question level (questions are designed to elicit false beliefs), not at the retrieval level (the corpus is still Wikipedia, which is generally factual). A genuinely adversarial retrieval setting—where the corpus has been poisoned or the retriever has been manipulated to return misleading passages—is untested.

What evidence exists in the paper. HotpotQA results (Table 1) show Nest-70B achieving 41.4 answer-level recall vs. 52.5 for RA—Nest substantially underperforms RA on this multi-hop task. This gap is consistent with the limitation described above: RA can attend to multiple passages in parallel through the prompt, while Nest must choose specific spans from specific passages. The TruthfulQA results are discussed in Section 5's critical assessment and show Nest outperforming RA but still achieving imperfect scores—the positive ΔBLEU means Nest is better than the base LM, but not that it is robust to adversarial information.

Mitigation status. Not addressed for multi-hop reasoning. The paper does not discuss how Nest could be extended to synthesize information across passages, and the architectural assumption that spans are copied from a single corpus location makes cross-passage attribution a fundamental challenge for the current design. For adversarial robustness, the paper notes that TruthfulQA results demonstrate some resilience, but does not propose or test explicit defenses against retrieval manipulation. The Broader Impact statement (Section 8) mentions filtering as a mitigation for corpus quality, which is a weak defense against adversarial attacks specifically designed to circumvent filters.

---

Difficulty Estimation Is Implicit and May Fail Silent

The assumption or constraint. Nest's dynamic interpolation via RRC estimates retriever confidence from the distribution of retrieval similarity scores. This is an implicit difficulty estimation mechanism: when the retriever is uncertain (all scores similar), Nest defaults to the LM; when the retriever is confident (one score dominates), Nest trusts the retriever. However, this conflates two distinct conditions: the query is difficult for the retriever (genuine uncertainty about which token is correct), and the corpus contains no relevant information (all retrieved tokens are equally irrelevant). In both cases, the similarity scores will be concentrated (all neighbors are similarly bad), leading to high $\lambda_t$ and reliance on the LM.

The consequence. When the corpus contains no relevant information for a query, Nest silently falls back to the base LM. This is the correct behavior (the corpus has nothing to contribute), but it is indistinguishable, from the user's perspective, from cases where the corpus does contain relevant information and the retriever simply failed to identify it. The system provides no signal that a query is outside the corpus's knowledge coverage. This matters for trust calibration: a user who sees attributed spans (green text with citations) in some parts of the output and unattributed text in others cannot tell whether the unattributed text reflects the LM's parametric knowledge filling corpus gaps, or the LM overriding corpus information the retriever failed to surface. Both cases produce the same observable behavior, but they carry very different implications for how much the user should trust the unattributed portions.

What evidence exists in the paper. The paper provides no direct measurement of this phenomenon. The attribution ratios in Table 2 show that some tokens are always unattributed (in the best case, Nest-7B on NQ attributes 93.4% of tokens, leaving 6.6% unattributed), but there is no analysis of whether unattributed tokens correspond to corpus gaps (acceptable) or retrieval failures (problematic). The sensitivity analysis in Appendix C (Figure 3c) shows that the threshold $\delta$ controls the aggressiveness of span copying, but this is a tunable parameter, not a diagnostic of when copying is appropriate. The paper does not measure retrieval precision or recall at the fact level, only perplexity and generation quality metrics, which confound retrieval quality with the LM's own capabilities.

Mitigation status. Not addressed. A natural mitigation would be to provide a confidence score for each attributed span (e.g., the RRC value $\lambda_t$ at the time the span was selected, or the acceptance probabilities from speculative decoding) so that users can calibrate their trust. Alternatively, the system could flag queries where the retriever's confidence is uniformly low as "no supporting evidence found" rather than silently falling back to the LM. The paper does not explore either option.

7. Implications and Future Directions

How This Work Changes the Landscape

Nest does not introduce a new model architecture, a new training objective, or a new retrieval algorithm. It introduces a new inference-time protocol for how LMs and retrieval corpora should interact, and that protocol carries a conceptual shift that is larger than the sum of its components.

The shift is this: the corpus is demoted from authority to proposal source, and the LM is promoted from consumer to verifier. Prior work in retrieval-augmented generation—whether in-context RA, RETRO, or standard kNN-LM—treated the retrieved content as something the LM should incorporate. The design question was how to incorporate it: through the prompt, through cross-attention, through output interpolation. In all cases, the retrieved text carried an implicit presumption of correctness. If the retriever surfaced a passage, the LM's job was to use it.

Nest inverts this. The corpus proposes candidate text, but every token of that proposal must pass the LM's own quality test (relaxed speculative decoding) before it enters the output. The corpus is not an authority; it is a draft model. The LM is not a consumer of retrieved content; it is a verifier that accepts or rejects proposals. This is not semantics—it changes where error responsibility lies. In in-context RA, when the model hallucinates despite having correct passages in the prompt, we blame the model's attention mechanism. In Nest, when the model copies an incorrect span, we can trace the error either to the corpus (the span was factually wrong) or to the verification step (the LM should have rejected it but didn't). The error is attributable, which is exactly what Nest provides at the output level: traceability of every claim to its source, including traceability of errors.

This reframing reconciles a tension that has run through the retrieval-augmentation literature since its inception. On one side, kNN-LM (Khandelwal et al., 2020) and its variants provide direct token-level attribution but degrade fluency (Wang et al., 2023a). On the other side, in-context RA (Ram et al., 2023; Shi et al., 2024a) preserves fluency but provides no guarantee that the retrieved content is actually used, and can even be harmful on adversarial datasets (TruthfulQA, where RA achieves negative ΔBLEU in Table 1). These were treated as competing approaches with an inherent fluency-attribution tradeoff. Nest's results suggest the tradeoff is not inherent—it is an artifact of treating the corpus as authority rather than as a proposal. When the LM gets veto power over corpus spans (through relaxed speculative decoding), fluency is preserved because implausible spans are rejected, and attribution is provided because accepted spans are verbatim. The ablation (Table 3) makes this concrete: adding relaxed speculative decoding to the system improves Biography FActScore from 41.6 to 46.8 while maintaining or improving fluency metrics, breaking the apparent tradeoff.

The paper also reconciles the contradictory findings around whether retrieval helps or hurts on adversarial tasks. The TruthfulQA results (Table 1) show that in-context RA produces negative ΔBLEU across all model sizes (7B: -0.34, 13B: -0.16, 70B: -0.13), meaning prepending retrieved passages makes the model's output more similar to false reference answers. This is consistent with prior findings that LLMs can be misled by retrieved content. But Nest achieves positive ΔBLEU (13B: 0.29, 70B: 0.17), meaning output-level interpolation is less susceptible to misleading retrieval than prompt-level augmentation. The mechanism is clear in retrospect: in-context RA bakes the retrieved content into the model's context, making it impossible for the model to selectively ignore misleading passages while attending to helpful ones. Nest's per-token interpolation and per-span verification allow the model to reject corpus text that doesn't fit the generation—a capability that prompt-level methods structurally cannot provide. This doesn't mean Nest is immune to adversarial retrieval (the corpus can still contain misleading information, and the verification step can still fail), but it establishes that the LM's ability to reject retrieved content is a critical design dimension that prior RA taxonomies (Asai et al., 2024) did not capture.

A less obvious shift is in how the field should think about the relationship between retrieval augmentation and inference acceleration. These have been separate research communities: retrieval-augmentation people work on making models more factual, and inference-acceleration people work on making models faster. Nest shows that the same mechanism—copying spans from a corpus and verifying them with the LM—simultaneously provides attribution (by linking spans to sources) and reduces latency (by generating multiple tokens per inference step). The 1.8× speedup in Figure 2a is not an incidental benefit; it is a direct consequence of the same span-copying that provides attribution. This unification matters because it changes the cost-benefit calculation for practitioners. A deployment that adds retrieval for factuality reasons gets acceleration for free (or rather, the acceleration helps offset the retrieval overhead). Conversely, a deployment that uses speculative decoding for acceleration can add factuality benefits by sourcing draft tokens from a trusted corpus rather than a draft model. The paper does not fully explore this dual-use property—the RA-Nest combination (which adds prompt-level retrieval on top of output-level copying) is tested but not analyzed through the lens of how the two mechanisms' costs and benefits compose—but it opens the door to thinking about retrieval and acceleration as jointly optimizable rather than competing for a latency budget.

The paper also provides a concrete diagnostic that should influence future work: the relaxation factor γ is a tunable knob that controls the fluency-factuality-speed tradeoff. Figure 2b shows that FActScore peaks at an intermediate γ (5×10⁻² for the 70B model), with lower FActScore at both stricter and more lenient settings. This means there is no universally optimal γ—the best setting depends on the task, the model, and the user's priorities. This is not presented as a limitation but as a feature: practitioners can dial γ up (stricter verification, higher quality, lower speedup) or down (more lenient verification, faster generation, more corpus copying) depending on their needs. For fact-critical applications like medical QA, higher γ provides a quality safeguard; for latency-critical applications like real-time chatbots, lower γ maximizes speed. The paper's finding that larger models benefit from larger γ (7B: 5×10⁻⁴, 13B: 5×10⁻³, 70B: 5×10⁻²) provides a practical heuristic: as the base model gets stronger, you should trust the corpus less, because the model's own parametric knowledge is more reliable.

Follow-Up Research This Work Enables

Direct comparison with REST on latency-quality Pareto frontier. REST (He et al., 2024) also retrieves drafts from a datastore for speculative decoding, but does not modify the output distribution—the datastore provides draft tokens, but the target distribution remains the base LM's. Nest modifies the distribution through kNN interpolation. The critical empirical question is: does interpolating the retrieval distribution into the target (as Nest does) produce better factuality than using the base LM's distribution and relying on the datastore only for draft proposals (as REST does)? A controlled experiment would fix the datastore, the retriever, and the base model, then compare Nest and REST at matched latency budgets (not matched generation budgets, since the two methods have different per-step costs). The key metric would be FActScore vs. tokens-per-second, producing a Pareto frontier for each method. If Nest's frontier dominates REST's, the interpolation mechanism is providing value beyond pure acceleration. If REST's frontier dominates, then modifying the output distribution is not worth the added complexity for factuality—the datastore's value is purely as a draft source, and factuality should be addressed through other means (e.g., better pretraining, prompt-level RA). This experiment would clarify whether Nest's contribution is primarily an acceleration method (with REST as the baseline) or a factuality method (with kNN-LM as the baseline).

Attribution quality metrics that enable cross-method comparison. The paper claims Nest provides better attribution than baselines but does not quantify this claim (Section 5 critical assessment). A follow-up study should define and measure attribution quality across Nest, kNN-LM, and in-context RA on the same tasks. Concrete metrics could include: (a) Attribution precision: for each factual claim in the output, what fraction of the attributed source text is actually relevant to the claim? (b) Attribution recall: what fraction of factual claims have any attributable source? (c) Source coherence: average length of consecutive output tokens attributed to the same source passage—high for Nest (span-level copying), low for kNN-LM (per-token independent retrieval). (d) Attribution accuracy: when a claim is attributed to a passage, does that passage actually support the claim? This last metric requires human annotation or a strong NLI model and would distinguish between Nest (which copies verbatim, so attribution is accurate by construction if the span is correct) and in-context RA (which may paraphrase and drift from the source). A study reporting these metrics on the Biography dataset (which has existing FActScore annotations) would directly test the paper's central attribution claim.

Training the RRC mechanism: learned interpolation with retrieval features. The RRC formula is a heuristic—the min-max ratio of retrieval similarities passed through a sigmoid—that requires manual tuning of α and τ. This is effective (the ablation shows it substantially improves over fixed λ), but it is almost certainly suboptimal. A natural follow-up is to train a lightweight predictor that takes the full set of retrieval scores (not just the min and max) plus features from the LM's own distribution (entropy, max probability, disagreement between top LM prediction and top retrieval prediction) and predicts the optimal interpolation coefficient. The training signal could come from a language modeling objective (minimizing perplexity on a held-out corpus, similar to how α and τ were tuned) or from a downstream task reward (maximizing FActScore). The key questions are: (a) how much does a learned λ predictor improve over the heuristic RRC? (b) do LM-derived features (uncertainty estimates) add value beyond retrieval-derived features? (c) does the learned predictor transfer across tasks and corpora, or does it overfit to the training distribution? Even a negative result—that the heuristic RRC is near-optimal and learning adds little—would be valuable, as it would validate the paper's design choice and establish RRC as a strong baseline for future work on adaptive interpolation.

Nest with on-policy retrieval and iterative refinement. The current Nest implementation uses the same set of top-b passages throughout generation (or rebuilds the token-level store periodically). But as generation proceeds, the context grows, and the relevant passages may change. A more sophisticated system would perform passage retrieval at regular intervals during generation (e.g., every 16 or 32 tokens), using the partial generation as the query, to update the set of candidate passages. This is particularly important for long-form generation (the Biography task generates up to 512 tokens) where early-retrieved passages may be irrelevant to later parts of the output. The experiment would measure: (a) does periodic re-retrieval improve FActScore and attribution coverage? (b) at what interval is the cost-benefit tradeoff optimal? (c) does the verification step (relaxed speculative decoding) naturally reject spans from stale passages, making re-retrieval unnecessary? This connects to work on adaptive retrieval (Asai et al., 2023b; Self-RAG) and would test whether Nest's LM-as-verifier architecture makes the system robust to retrieval staleness.

Adversarial evaluation with a poisoned corpus. Section 6 identifies that Nest's factual accuracy depends on corpus correctness, but provides no empirical measurement of this dependence. A controlled stress test would take a corpus known to be factually correct (e.g., Wikipedia), inject a known set of factual errors (e.g., swap birth years of famous people, change numerical values in scientific claims), and measure how often Nest copies the poisoned spans vs. rejecting them through speculative decoding. The key questions: (a) At what error rate does Nest's FActScore drop below the base LM's? This establishes the corpus quality threshold at which retrieval becomes harmful rather than helpful. (b) Does the RRC mechanism (which controls interpolation weight) effectively downweight retrieval when the corpus is noisy? The min-max ratio of similarities might not change if the poisoned passages are topically similar to the query—the retriever could be confidently wrong. (c) Can the relaxation factor γ be tuned to make the system robust to corpus noise, and what is the resulting speed-factuality tradeoff? This experiment would directly inform deployment decisions for practitioners using Nest with web-crawled or user-generated corpora where factual reliability is uncertain.

Combining Nest with structured knowledge sources beyond free text. The paper uses Wikipedia as a flat text corpus. But many knowledge-intensive applications have structured or semi-structured data available: medical guidelines with entity relationships, legal codes with cross-references, product catalogs with attribute-value pairs. Nest's span-copying mechanism is designed for free text and would not naturally handle structured data where the "continuation" in the source doesn't correspond to a coherent linguistic span. A follow-up could extend Nest to structured knowledge by defining span selection differently: instead of following the n-gram continuation in text, follow a structured "continuation" in the knowledge graph (e.g., from an entity to its properties, from a legal clause to its sub-clauses). The speculative decoding step would need to verify that the structured content can be expressed as fluent text under the LM's distribution. This is non-trivial—structured data verbalized as text often reads unnaturally—but if successful, it would extend Nest's attribution benefits to domains (medicine, law, e-commerce) where structured knowledge sources are common and factual accuracy is critical.

Practical Applications and Downstream Use Cases

Verifiable AI-generated content for high-stakes domains. In journalism, legal document drafting, and medical summarization, the ability to trace every factual claim to a specific source passage is not a nice-to-have—it is a requirement for adoption. Nest's span-level attribution (Table 2: 33.2-95.5% of tokens traceable to the corpus, with average spans of 3.0-27.9 tokens) provides this capability without requiring architecture changes or fine-tuning. A news organization could deploy a Llama-2-Chat 70B model with Nest pointing to its own article archive, generating draft summaries where every factual statement is color-coded with a citation to the source article. The 1.8× speedup means the system would be faster than the base LM alone, despite the additional verification. The key practical consideration is corpus quality control: the organization must maintain a curated, fact-checked corpus (their own published articles) rather than relying on a general corpus like Wikipedia, since Nest copies verbatim and will reproduce any errors in the source. The attribution ratios in Table 2 suggest that even with a general corpus, most tokens are attributable; with a domain-specific curated corpus, the ratio would likely be higher.

Cost-efficient retrieval-augmented serving with latency guarantees. Current production retrieval-augmented systems (e.g., ChatGPT with browsing, Perplexity.ai) face a tension: adding retrieved passages to the prompt improves factuality but increases latency because the LM must process more input tokens and generate output tokens autoregressively. Nest offers a way to break this tension: by generating multiple tokens per inference step through span-copying, it can offset the retrieval overhead. A deployment serving QA queries at scale could use Nest with a moderate γ (quality-optimized) for fact-critical queries and a lower γ (speed-optimized) for latency-critical queries, adjusting the tradeoff per-query based on user needs or SLA requirements. The latency breakdown in Figure 2a shows that passage search and token index building account for roughly 25-30% of total latency—this is the cost of adding factuality, and the speculative decoding provides the speedup that makes it palatable. With the 70B model, the net effect is still a 1.8× speedup over the base LM with no retrieval, meaning the system is both more factual and faster than the status quo. This changes the deployment economics: rather than choosing between a fast-but-inaccurate base model and a slow-but-accurate RA model, practitioners can deploy Nest and get both properties simultaneously.

Offline data generation for training more factual models. The paper's method is training-free, but its outputs could be used to train better models. The idea is simple: use Nest to generate attributed, factually-grounded text on a large corpus of prompts, then fine-tune a base LM on these outputs (with or without the attribution annotations). The resulting model would ideally internalize the factual precision of the retrieval-augmented system while maintaining the low latency of a pure parametric model at inference time. This is a form of knowledge distillation from a semi-parametric teacher to a parametric student. The key advantage over standard distillation from in-context RA is that Nest's outputs have explicit span-level provenance: the training data includes not just the generated text but also the source passage for each attributed span. This enables training objectives that explicitly reward the student for reproducing facts from the cited sources, potentially leading to better factual grounding than distillation from an RA teacher where the link between output and source is implicit. The paper's finding that smaller models benefit more from Nest (7B shows larger relative gains than 70B in Table 1) suggests that distilling Nest-70B's outputs into a 7B or 13B model could be particularly effective—the smaller model gets the factual benefits of the larger model's verification capability without the inference cost.

When to Prefer This Method

The paper explicitly positions Nest against two alternatives: standard kNN-LM (output integration with fixed interpolation) and in-context retrieval augmentation (input augmentation with retrieved passages in the prompt). The experiments in Table 1 provide the empirical basis for decision rules, and the latency analysis in Figure 2 adds another dimension. The following decision rules emerge:

• Prefer Nest over standard kNN-LM when attribution coherence matters (you need spans traceable to single sources, not a patchwork of per-token attributions), or when generation quality is important (the ablation in Table 3 shows Nest substantially outperforms the two-stage kNN-LM baseline on all metrics, with RRC and speculative decoding providing the gains). The fixed interpolation coefficient of standard kNN-LM is strictly dominated by Nest's dynamic λt—there is no task in Table 1 where kNN-LM outperforms Nest by a meaningful margin.

• Prefer Nest over in-context RA when (a) attribution granularity is required (Nest provides span-level citations, RA provides passage-level context with no guarantee of faithful use), (b) latency is a constraint (Nest achieves 1.8× speedup on 70B while RA adds latency from processing extra prompt tokens), or (c) the task is adversarial or the corpus may contain misleading content (TruthfulQA results show RA can be harmful while Nest maintains positive ΔBLEU). The Biography results (Table 1) show RA achieving higher FActScore than Nest alone (70B: 52.9 vs. 41.6), but RA-Nest achieves 59.2, suggesting the two methods are complementary for factuality. Nest alone may not match RA's factuality ceiling, but it provides attribution and latency benefits that RA cannot.

• Prefer in-context RA over Nest when (a) raw factuality is the sole metric and attribution is not required (RA achieves higher FActScore and QA recall than Nest alone across most settings), (b) the task requires synthesizing information across multiple passages (HotpotQA: RA achieves 52.5 vs. Nest's 41.4 for 70B, consistent with RA's ability to attend to multiple passages in parallel while Nest copies from single passages), or (c) the deployment cannot support the retrieval infrastructure Nest requires (dense + sparse indexes, on-the-fly token encoding, speculative verification)—in-context RA requires only a passage retriever and a prompt template.

• Prefer RA-Nest (combined) when the application demands both maximum factuality and span-level attribution. The combination achieves the best of both worlds in Table 1: highest FActScore on Biography (59.2 for 70B), lowest perplexity on language modeling, and competitive QA performance. The cost is increased complexity (two retrieval pipelines: one for prompt augmentation, one for output interpolation) and potentially higher latency (not measured for RA-Nest separately), but for offline or asynchronous applications where factuality and attribution are paramount, this is likely the optimal configuration among the tested alternatives.

• Prefer the base LM (no retrieval) over Nest when the problem is well within the model's parametric knowledge (MMLU results show Nest provides only 0.3-0.8 point gains, not worth the infrastructure cost) or when the inference infrastructure cannot support the additional components (GPU cluster for the LM plus CPU-based dense/sparse search with 32 threads for retrieval). The latency breakdown (Figure 2a) shows Nest is faster than the base LM on the paper's hardware, but this advantage may not hold on all hardware configurations.