Chunk-Distilled Language Modeling
Authors: Yanhong Li
Karen Livescu
Jiawei Zhou
Venue: ICLR 2025 Poster (2025)
Keywords: language modeling
text generation
retrieval-augmented generation
domain adaptation
inference algorithms
efficient generation
Primary area: generative models
TL;DR: A training-free language modeling approach that dynamically inject text chunks into generation through fine-grained retrieval, with the ability to adapt model distribution and increase inference efficiency.

Abstract: We introduce Chunk-Distilled Language Modeling (CD-LM), an approach to text generation that addresses two challenges in current large language models (LLMs): the inefficiency of token-level generation, and the difficulty of adapting to new data and knowledge. Our method combines deep network-based LLMs with a straightforward retrieval module, which allows the generation of multi-token text chunks at a single decoding step. Our retrieval framework enables flexible construction of model- or domain-specific datastores, either leveraging the internal knowledge of existing models, or incorporating expert insights from human-annotated corpora. This adaptability allows for enhanced control over the language model's distribution without necessitating additional training. We present the CD-LM formulation along with performance metrics demonstrating its ability to improve language model performance and efficiency across a diverse set of downstream applications. Code and data will be made publicly available.

Published as a conference paper at ICLR 2025
CHUNK-DISTILLED LANGUAGE MODELING
Yanhong Li University of Chicago & TTI-Chicago cerhoongg@uccchcco edu
Karen Livescu TTI-Chicago k1 klivescušttic.ed:
Jiawei Zhou
TTI-Chicago & Stony Brook University j jiawei. zhou. 1 @estonybrook edu
ABSTRACT
We introduce Chunk-Distilled Language Modeling (CD-LM), an approach text generation that addresses two challenges in current large language models (LLMs)): the inefficiency token-level generation, and the difficulty adapting to new data and knowledge. Our method combines deep network-based LLMs with a straightforward retrieval module, which allows the generation of multi- token text chunks at single decoding step. Our retrieval framework enables flexible construction of model- or domain-specific datastores, either leveraging the internal knowledge of existing models, or incorporating expert insights from human-annotated corpora. This adaptability allows for enhanced control over the language model's distribution without necessitating additional training. We present the CD-LM formulation along with performance metrics demonstrating ability to improve language model performance and efficiency across a diverse set of downstream applications
INTRODUCTION
Large language models (L.LMs) have become a crucial component of intelligent systems, but still suffer from fundamental challenges their efficiency and performance. LLMs most commonly based on autoregressive Transformers (Vaswani et al. 2017) and typically generate text sequences one token at time in serial fashion, which limits their efficiency, Moreover, once pre-trained, up- dating the model parameters requires expensive data and computational resources, making it difficult to incorporate dynamic knowledge into the model.
Several techniques have been proposed to improve the efficiency and performance of LLMs, such
speculative decoding (Leviathan al. 2023 Chen al. 2023 Miao et al.) 2024 Spector & Re 2023) and vetrieval-augmented generation (RAG) (Lewis et al. 2020 Guu et al. 2020 Borgeaud al. 2022). The former relies on a smaller model speculate several tokens at a time reduce inference runtime while retaining the same model distribution, while the latter combines parametric language models with non- can-aaamettric memory improve adaptability to dynamic knowledge but often without efficiency gains.
This work aims alleviate both challenges via fine- aine-grineed eetrieval-augmeenee language model- ing approach that focuses on text chunks, contiguous spans of tokens that often appear together. The intuition for this approach is that substantial amount of linguistic or factual knowledge can be expressed in text chunks spanning multiple contiguous tokens, such as named entities, multi- word expressions, and other common phrases. These sub-sentence structures tend to exhibit lower variability compared to larger text units such as sentences, and are often memorized preciscly by well-trained LLMs. Figures and 2 demonstrate this effect: Chunks conveying key content are often repeated verbatim across multiple decoding runs with similar contexts, and the LLM proba- bilities over token sequences show recurring plateaus of high probability within such multi-token chunks. By injecting memorized novel chunks into the generation process, we may be able to "oode and data available httpp:::/gittbb.comm//////////// 1m

---

Published as a conference paper at ICLR 2025
USER What the answer life the universe and everything?
USER The answer to life, the universe, and everything is?
USER Tellme the meaning life everything.
LLM
The answer to life, the universe, and everything 42, according Dougles as Adams'
LLM
The answer to everything is 42, reflacting Douglas Adams humorous take on the quest universal truths in The Hitchhikkr's Guide the Galaxy."
LLM
The meaning of life, the universe, and everything referencing Douglas Adams) iconic work Thhe Hitchhika''s Guide to the Calexy."
Hitchhkkrrs Guide tha Golaxy:"
Figure I: LLMs may generate sequences with repeated chunks spanning contiguous tokens convey- ing key information in similar contexts. Examples are generated from -lamm-22-7b-che
improve the models' ability to adapt to new domains or knowledge. In addition, if entire chunks can be cached and retrieved during inference, we should also be able to speed text generation.
Token Probability for L'ama-2-7b-che and Llama-2 70b-chat
Toloon Probabiiltty
0.5
Uamo Llamaad/0b-chhtt
The
artweds
gl
the
universe
and
eversthigg
accorting
Bnngitine
The
Graditg geth oo get
Figure 2: LLM token probabilities for the sentence: "The answer to life, the universe, and everything 42, according to Douglas Adams' The Hitchhiker's Guide to the Galaxy." " These models bind token sequences such as Douglas Adams' and The Hitchhiker's Guide to the Galaxy into chunks with plateaus of high probability.
Inspired by these observations, we present Chunk-Distilled Language Modeling (CD-LM), new training-free generation approach that mixes LM token generation with chunk retrieval. To facilitate efficient search, we store text chunks of variable sizes, along with their preceding contexts, in trie-structured datastore, and retrieve the most likely chunks as possible text continuations given the current generation. The context matching is done in the vector representation space induced by the LM itself without the additional overhead of specialized embedding modules, commonly used in
RAG (Lan et al. 2023 Ram et al. 2023 Borgeaud et al. 2022).
Well-matched chunk continuations are accepted, skipping multiple token decoding steps. Using the same generation approach, CD-LM allows language models (L.Ms) to work with chunks mined in different ways to achieve various goals in applications. As suggested by Figure 2 chunks can be naturally derived from any parametric pre- t-tirrinedd LM as memorized high- high-probability sequences. When the chunks are extracted from the distribution of a more powerful specialized LM, CD-LM implements a form of knolwedge distillation, adapting the base model's distribution (without any additional training) by injecting chunks during inference. In this setting CD-LM can either improve smaller models with knowledge drawn from larger models or perform training-free domain adapta- tion. On the other hand, when the chunks extracted from the same LM used for generation, they form a self-memory datastore that can be used to improve inference efficiency while maintaining the same model distribution, as speculative decoding. Finally, the chunks can be not only extracted from a parametric model but even directly curated from human experts. Such external knowledge can be factual information or private data that the LM may not have direct access to.
CD-LM requires no training and can work with any off-the-shelf language model in both chunk discovery and sequence generation. We conduct a diverse set of empirical studies, including lan- guage modeling perplexity, text generation, and domain adaptation, showing the ability of CD-LM to improve inference efficiency and modeling performance.
2 BACKGROUND
While many attempts have been made to improve language modeling and generation efficiency, it remains significant challenge to address both simultaneously. For example, non paraparmtttrc ap- proaches like KNN-LM (Khandelwal al.) (020) reduce LM perplexity in certain domains, but tend to require a sizable database for retrieval and adds latency during generation; specialized inference algorithms like speculative decoding (Spector & Re 2023) speed up generation but keep the LM's distribution fixed. Unlike prior work, CD-LM can both speed up generation and adapt the LM's distribution. We include more comprehensive overview of related work in Appendix
2

---

Published as a conference paper at ICLR 2025
chunks using an token probabiities, or using existing human knowledge
the answer to the ultimate question liffe, the universe, and averything 6
0 better model KCD-LM
P.O the model lseelf SCD-LM
the humocous tone of his writing evideut, rellecting Dougias Adaass unique wit and creativity
humaa experss
lext corpus
dry humor reminiscent Hischhike''s Guide Galay,
reflecting Doughs
ECD-LM
Inference
Build Trie
tree
Accept reject Generate chunk directly accepted
What is the answer to life the universe undreeyythhing
everything
rellecting
Each Trie node
an entry token everything
anewer to everythhing
(effecting Dosulls Adiani
humor
Each represents
Douglas
chunk for retrivval: is 42
1.00
Each nade chunk's
Addams
preceding contexts:
< The answer to everylhing is
coflecting Douglas Adams
Figure 3: Overview of CD-LM. Colored text spans are generated together by chunk retrieval, in- terleaved with token-by-token generation by the LM. Note that same chunk can appear multiple contexts, so each node in the trie datastore contains multiple context vectors in practice.
Non-Parametric Language Modeling KNN-LM (Khandelwal et al.) 2020) extends pre- traited LM by linearly interpolating its distribution with a non-parametric k-nearest neighbors model based on token retrieval, thereby often improving language modeling performance. However, it is typi- cally very inefficient as it performs retrieval at each token, and it affects the immediate next token distribution via soft mixing. There is series of proposed methods that improve the efficiency of kNN-LM (He et al. 202T Alon et al. 2022); however, they are still slower than the pre-trained LM. Unlike KNN-LM, CD-LM does not involve retrieval at each token and makes hard decision about multiple tokens in chunk rather than mixing token distributions, enabling it to enjoy the benefits of dynamic retrieval but also save on kNN searches.
Speculative Decoding Speculative decoding (Leviathan et al. 2023 Chen et al. 2023 Miao et al., 2024 Sppecorr & Re 2023 He et al., (024)) is an inference acceleration technique. Given a particular target LLM, smaller LM is used to quickly generate multiple draft tokens, which are then consid- ered together by the target LLM. The work most closely related to ours REST (He et al.) 2024), which retrieves draft token sequences from an external datastore. While CD-LM also retrieves chunks from a datastore, it is fundamentally different from speculative decoding. Speculative de- coding methods use the target LLM for draft token verification, the language model's distribution and therefore downstream performance cannot be further improved and no new knowledge can be injected. In contrast, CD-LM is designed to inject chunk-level knowledge into generation the model distribution can be adapted.
LANGUAGE MODELING WITH CHUNK GENERATION
In this section, we introduce general framework of language modeling that interleaves chunk gen- erations with tokens from a standard autoregressive LM. We then describe the operational details of the chunk generation process with retrieval from structured database in Section 4 Together, these two sections build the core ideas of CD-LM. Finally, we derive tractable algorithm for computing sequence probabilities under CD-LM Section 5
3.1 PRELIMINARIES
An autoregressive language model assigns probability to any given sequence of tokens (2,,.2...... NN as follows: po(41, (1,1,,2,,.... xNn II Po( Iea), where 0 is the model pa- rameters and Ten = (I1, Æ2, In- Modern LLMSs are usually parameterized by Transformer Vasswani al., 2017) architectures composed of stacks of self-attention and feedforward layers. Individual tokens from closed vocabulary V are sequentially passed into the model with their embedding vectors, and the next token probability distribution computed by
hn fe(æ1, ,,, In-1)
pa( (Fn)) softmax (Woh,)
where fo(-) denotes the functional process that maps the previous sequence of tokens into fixed- size context vector hn E Rdd, and W, IR Vixd is the output embedding matrix that projects the
3

---

Published as a conference paper at ICLR 2025
representation vector onto the vocabulary space. Given learned model, text can be generated by sampling from the next token distribution autoregressively one token at a time, resulting in NN forward runs for sequence of length N.
3.2 TEXT CHUNK GENERATION MODELING
Instead of producing text one token at time, we provide a mechanism that can directly generate span of multiple consecutive tokens, or chunks, with better efficiency and flexibility of injecting knowledge on fine-grained sub-sentence levels into the model distribution on the fly.
Formally, we use n to index sequential token position, and to index generation steps. For every step, we allow generation either a single token from base LM M₁ with parameter 0, text chunk from different model G. which we call the chunk proposal model. Let 4, de- note the sequence length (i.e., the number of tokens) after t steps. Unlike typical token-base decoding, we have 4 2 t. In particular, the chunk proposal model & takes any prefix Ien and returns a possible text chunk contin- uation On (In, #n-1, ...... In) with acceptance probability q, E [0, 1], and Tn the length of the proposed chunk We introduce binary random variable 2n that denotes whether the generation at token position uses the chunk proposed by or defaults to the single token generated by the L.M, and p(zn. = 1) qn- The chunk- integrated generative process as follows:
Figure A graphical model illustration of the probabilistic model of CD-LM.
The token sequence 2n nodes are ob- served, and chunk acceptance variables 2n are latent, governing how many to- kens are to be generated at one step.
(1) For step = 1, LM MM generates the first token 21. We have current sequence length l 1.
(2) At step 2 2, set next token position: n 4-11 +1;
(3) Chunk proposal: G(I=en)) + (Cns 9n), and length of is T
(4) Sample: 2n ~ (ornoulli(qx):
(5) If 2, 1: accept Cu. and 4 4-11 + Txi
(6) Else Zn O: reject Cn. Generate vn from the base LM Mg. and 4 = 4-1 + 1;;
(7) Move generation step 1.
The above process is also illustrated as a graphical model in Figure
We refer to this approach Chunk-Distilled Language Modeling, CD-LM. The chunk proposal model G could take any parametric or non-parametric form in principle, and next we specifically adopt simple retrieval model of text segments reduce the cost of chunk proposals.
CD-LM WITH GRNE-GRAINED RETRIEVAL
In this section, we describe in detail retrieval- -based chunk proposal model needed for step (3) defined in Section 3.2] completing the chunk-interleaved generative process under CD-LM. These details include datastore representation of chunks (Section [.1). chunk proposal process with re- trieval (Section 4.2). and chunk sources that enable different applications (Section 4.3)
4.1 CHUNK DATASTORE CONSTRUCTION
Given any text corpus C, suppose there is an expert model E (to be elaborated in Section 4.3 that identifies text spans in C that want to re-use for generation. These chunks often convey integral information about linguistic rules or factual concepts, such as "is 42" or "Douglas Adams"" in Figure T We construct datastore of the identified chunks with preceding contexts, D ((r, s,) } where r, is the previous content leading to the chunk and s; is the text chunk which could be of arbitrary length.
We break down the chunk context r; into two parts, ri = (u,, v;), where u, is the preceding context except the last token, and v; is the last token immediately leading into the chunk s,, which we
Tn when the proposed chunk is empty, i.c. Cn 0.
4

---

Published as a conference paper at ICLR 2025
define as an entry token. For instance, for the chunk "is 42" in Figure I the entry token is "everything". We will use u, as keys to match contexts for chunk retrieval, and use U/ as entry points linking to possible chunk candidates.
The chunk contexts u are represented by the context vectors fo(u) produced by running the forward process of the LM My in Eq ( which will facilitate context matching in vector space (Khandel- wal et al. 2020 3 We store the chunks using a collection of trie structures for efficient storage and retrieval, so that D = (Twip Twy..... [nyyy where each Tu stores all chunks that follow the same entry token w in the LM vocabulary V. We refer to the The as entry token tries, where entry token w the root node of Tw. each node is a token, and the paths from the root to each node represent either a chunk or prefix of chunk. Each node of a trie contains all of the context vectors corresponding to the unique chunk represented by the node (see Figure 3 for an example).
4.2 ADAPTIVE CHUNK RETRIEVAL FOR GENERATION
Given previously generated tokens F<n= formulate the chunk proposal model ((zen) (Ca, (n) as a chunk retrieval process to be interleaved with the LM generation. We use the infor- mation from the LM computation route the most recent token 2n-11 to derive plausible chunk proposals. Per Eq 1. right before generation of In-1, the context vector Se f0( provides a summary of the context, which we use as the query for chunk retrieval. We use Ia-1 as the entry token to confine the chunk search to the corresponding trie T2.- leading to smooth chunk contin- uations (for instance, see the searched trie in Figure 3). This is crucial for improving the naturalness of the retrieved chunks combined with the previous context. In the meantime, using the entry token trie to limit the search space also greatly enhances retrieval efficiency. Formally, the chunk proposal model G is given by
(u" .G. arg max
(u.s)ET=n-1
In = 9o (a
where u is the stored chunk context except the entry token, E T means searching for each node
in the trie, sim((:-) is vector similarity measure for which we use cosine similarity, and gs(-) is a function parametrized by Φ to convert the similarity scores to acceptance probabilities, which can
be tuned for different base LMs MM (implementation details described in Section 6
4.3 CHUNK EXTRACTION MODEL
Now we describe the expert model E that provides the chunks for the datastore. We categorize the possible knowledge sources into three major types intended for various CD-LM applications:
Knowledge Distillation As suggested earlier in Figure well-trained LLMs memorize text chunks with high probabilities in certain contexts. This provides a natural source of automatically defined chunks from models' internal knowledge. Let Mo, denote the pre-trained model with parameter 0y that we derive chunks from. It is often a more powerful model that is larger or more specialized than the base LM Ma. so that Ma, serves as a teacher to help adapt the distribution of Ma. Operationally, for chunk identification we run Ma, on a text corpus C, and apply thresholding heuristic to extract the longest chunks whose token probabilities are all above threshold y (see Appendix D.2 for an example). Formally, chunk along with context r is extracted if 3r. s.t. po(x,/r, x<ii) 2 y, V(z,, 224) s. Note that in this approach, the chunk datastore construction needs only one forward pass of Mar on C. We also run a forward pass of the base model Mg on the chunk contexts for their hidden vectors for retrieval during generation (described in Section In this scenario CD-LM essentially performs a form of knowledge distillation from Mø, to Ma but without any training, by directly injecting Me. defined chunks during inference on the fly. We call this setting knowledge CD-LM, or KCD-LM.
Self Distillation We can also allow the teacher model May to be the same as the base model Mo- where we perform self distillation via chunk generation to improve inference efficiency. The datas- tore serves as an explicit self-memory consisting of frequently generated chunks of high probability,
'It is also possible to directly use context strings for matching rather than vector-based dense retrieval. WWe found that context matching with different LMs' hidden vectors could exhibit very different cosine similarity scores, and for small LMs the scores fall within a tight numeric range.
5

---

Published as a conference paper at ICLR 2025
shown in Figure By retrieving from the explicit self-memory, reduce the computational cost of re-generating every token sequentially within the same chunks in future generations from the model. Chunk datastore construction follows the same thresholding heuristic as in KCD-LM, for which we run one single forward pass of Ma on C for extracting both chunks and their context vectors from the same model. Though self distillation via chunk retrieval can speed up inference, it may introduce a slight distribution shift. We tune the retrieval threshold to preserve generation quality while still achieving significant token savings We call this setup self CD-LM, or SCD-LM for short.
Expert Distillation With KCD-LM and SCD-LM, the extracted chunks represent the parametric knowledge of an LM. In some situations, the chunks could directly come from human experts as added knowledge to inject into generations. For example, one source of such expert knowledge could be hyperlinked text spans in Wikipedia articles. Another important source could be private information in a personal database that cannot be accessed by the parametric model. With chunks provided this way, run Ma only to acquire the context vectors for the datastore construction. This is also knowledge distillation process but with non-parametric expert-curated knowledge. We call this approach expert CD-LM, or ECD-LM.
PROBABILITY DISTRIBUTION UNDER CD-LM
Sampling text with the CD-LM generative process in Section 3.2 is fairly easy, but assigning proba- bilities to given text sequence non-trivial. This requires enumerating all possible chunk proposals at different token positions to marginalize the in variables, which is complicated due the variable chunk lengths and dependency structures shown in Figure 4 We derive a dynamic program sim- ilar to backward algorithm for computing sequence probabilities under CD-LM, allowing to measure intrinsic language modeling performance with perplexity (PPL).
For any given sequence Ti.N. the chunk proposals every position (Cns 4n) from (( are deter- ministic given the datastore 'D and thus can be pre-computed. CD-LM models the joint distribution of x ii,N= 22.N
N II n=2
[p(zn aa Zen) - zgn)] [m:sin))
p(ZiNN1 22.N = p(t
(3)
where the binary indicator function Il (n; 22: } marks whether the token position n is inside of sampled chunk based on the values of 22m To marginalize over 22:N, we define
6
On = -P( Fn = 1, x zen)= = 1 (- ( = cx) (antt, In- + Butt, (1 In qn--,,)) r.
3, = p(r Izn = 0, x' R<n)= = pa ( (x=)= [an+19n+1 + Ba+1 9n+1)]
where the function f -1 = =c,} indicates whether the proposed chunk en exactly matches the given text segment, and pe is the probability from Ma. By computing a and |3 values backward from N to 2, we can get the marginal sequence probability under CD-LM as
p(z1.n) po(=?) [a242 + B2 [g]]]
(5)
Please check more details and full derivations in Appendix A as well as an example in Appendix B
6 EXPERIMENTS
We conduct experiments on multiple LMs and tasks. We formulate 9s in Eq (2) as a simple piece- wise linear function, where the maximum context matching similarity score only maps to non-zero chunk acceptance probability larity scores in the range [7, 1] are then linearly mapped [0. 1]. See Appendix D.4 for full details. We decode in greedily. which is equivalent to accepting žn = 1 when the chunk context matching similarity score passes a threshold n+1
if the score is larger than v) 2 0, which is a hyperparameter. Simi-
2
*This similar speculative decoding (Chnn al. 2023 We do not apply LM verification formally guarantec generations exactly the same as the original, although this could he adopted straightforwardly too. "Note that some is not well defined if position is already inside previous chunk, but we use all #2.N denote the whole 2, sequence for notational convenience.

---

Published as a conference paper at ICLR 2025
6.1 KNOWLEDGE DISTILLATION
Table I: Perplexity on test sets with KCD-LM.
Table 2:
MAUVE score generations with KCD-LM against real continuations.
WikiText
Medical
Law
Code
Base LM (137M)
34.83
51.68
11.41
106.44
WikiText
Medical
Code
Teacher Model (1.5B)
14.48
17.66
5.15
62.09
Base LM line tuned
25.55
22.46
6.65
Base LM
0.016
0.006
0.015
0.024
&NN-LM
39.66
11.10
89.85
KCD-LM
0.032
0.011
0.040
0.053
RETOMATON KCD-LM
32.10
39.66
11.10
89.88
% 1
50.7%
100.9 %
162.8 %
121.3 %
22.90
24.95
8.24
50.77
Model and Data We focus on two
objectives: improving language mod- eling performance and enabling do- main adaptation, using a weak pre- trained 137M GPT-2 small model as the base language model, Mo. for KCD-LM. For language model- ing, we evaluate on the WikiText- 103 dataset and the Dockerfile subset of the GitHub Code dataset 7 Docker- file is a low-resource code language and the base model has poor PPL on the Dockerfile data. This setting al- lows us to explore the effectiveness of KCD-LM in low-resource settings. For domain adaptation, we focus on adapting to medical and legal do- mains. We use the Medical Instruc- tion Dataset consultations, and the Federal Register subset of the Pile-of-Law (Henderson et al. (202)). For these tasks, we set as the teacher model Ma. either a pretrained 1.5B GPT-2 XL model (for code) or an off-the-shelf domain-specific GPT-2 XL model (for WikiText, medical, and law) Chunk data- stores are constructed from the corresponding training sets. Empirically, extracting chunks from WikiText- 103 takes under hour on four A4000 GPUs for small base model like GPT-2. Build- ing the WikiText datastore takes up to 1.5 hours on a single A4000 GPU, while other datastores are built within 30 minutes.
Wikitext
Datastore Sire (cill
Medical
Dittsture Size (Cm)
30-
De
0/7 na uSS 64
0.5
69
02
7.7
0b Extractiion Threshold
U.
0.4
0,3
Chunk Extraction Thrsstold
Law
Dataster Size (o00)
Code
DH AAAC Sizw (Gtt
120-
10
0.2
0's
0.4
KNN-LM PFL
CD-LM PPL
а*** Baseline LM PPL
Datastore Size (GB)
Figure 5: Comparison between KCD-LM and KNN-LM on PPL., along with datastore sizes controlled by chunk extrac- tion threshold
B which contains conversations between an Al assistant and patients during medical
Evaluation We measure PPL computed from 512-token sequences on corresponding test sets. For datasets that do not come with test split, we construct test sets of 500 sequences to match WikiText. PPL is computed with our dynamic program derived under the CD-LM distribution in Section
5 Since datastore sizes vary for different chunk extraction thresholds 7, we ensure the data- store remains consistent when comparing PPI. between KCD-LM and baselines. We also evaluate text generation in these domains with the MAUVE score (Pillutla et al., (02T) to measure the simi- larity between texts generated by MM and ground truth continuations in the test data. Additionally, we conduct LLM-aa--ajudde (Liu et al., 2023; Fu et al. 2024) pairwise comparisons to assess the quality of generated text beyond perplexity and MAUVE. To generate text, we follow prior work on evaluating text generation with KNN-LM (Wang et al.] (203). sampling 5,000 sequences of 100 to- kens each from both validation and test sets. These tokens serve as prompts for the LMs to produce an additional 150 tokens using greedy decoding.
Results As shown in Table 1
T0 our KCD-LM model significantly reduces the PPL across all eval-
uated datasets, surpassing both the base LM, KNN-LM and RETOMATON (Alon et al. 2022 II We
hhttps://hugginggfeeee co http::/:w....ttteeaaatttt/etrrre hht hhttp::/huuginggfaese wtt.///attttttt//aaaaaeeeeeeeeee
instructi 1o-on100 100k 'We use pre-trained GPT-2 XI. model for code because its PPI. is already significantly lower than that of
co. ceto//aaaaetttt
GPT-2 small, making it sufficient for effective knowledge distillation.
Note that the shaded rows show reference baselines, not direct competitors to our training-free approach.
The results for RETOMATON and KNN-LM are similar because RETOMATON focuses primarily on improving the efficiency of KNN-LM. The slight PP[. improvement is an additional henefit, observed mainly on WikiText-103.
7

---

Published as a conference paper at ICLR 2025
achieve drastic PPL reduction on WikiText, Code, and Medical with GPT-2 small on the fly without the need to update the weak model. Without any additional training, our distillation process signifi- cantly reduces GPT-2-small's perplexity, bringing its performance much closer to that of the teacher model while being ten times smaller. KCD-LM achieves comparable or even better PPI. than GPT-2- small directly fine-tuned on domain-specific data. We exclude PPL results Code because a domain domain-specific model for it, making direct fine-tuning of GPT-2-small an unfair comparison with our distillation method. Additional PPL results and datastore sizes are also illustrated in Figure 5 with varying chunk extraction threshold y for datastore construction with Mø,. KCD-LM beats kNN-LM with explicit and sparse chunk retrievals. For text generation, Table Thooe[sshows shows substantial improvements with our approach over Mg measured by MAUVE 12 The pairwise LLM-as-a- evaluations indicate that KCD-LM consistently outperforms baseline models across all domains, as detailed in Appendix E.4
lack
Scalability and Efficient Retrieval central draw- back of KNN-LM is the need to store every token's con- text in key-value form. By contrast, our chunk-based approach selectively extracts only a subset of corpus chunks, reducing storage overhead. Table [3] shows the total number of chunks, their fraction of the overall to- ken count, and how many are actually searched on aver- age. Storing only these chunks lowers the datastore size to about 30-40% of what a naive kNN-LM would require. Moreover, KCD-LM restricts retrieval to chunks sharing the same entry token, pruning the search space to around 0.002-0.0 1% of the datastore each generation step, thus substantially reducing computation costs.
Table 3: Extracted chunks and average chunk usage during generation.
WikiText Medical
Law
Code
#chunks
36.1M
10.6M
12.9M
1.7M
#chunks/Hall tokens (%)
30.8
38.4
34.1
43.5
Avg. #chunks (%) per trie 0.002
0.003
0.003
0.011
6.2 SELF DISTILLATION
Experimental Setup We focus on practical scenarios where language models operate in environ- ments with frequent repetition of similar or thematically related queries. In such contexts commot in customer support or domain-specific assistants-buildingg datastores for common topics and caching frequent generations can enhance efficiency. We use three instruction-tuned LMs as Mg:
uPT-2-xxl-comversation 13 -LaMA-2-7b-chle (Touvron et al. 2023). and Mistral- Distral-7B-Instreee I---Innttuttt....
(Jiang al., 2023a). We build two testbeds for SCD-LM from the MT-Bench (Zheng et al.] 2023) dataset of multi-turn conversational questions. The first testbed uses the initial question from each of the 80 sets of multi-turn questions, which we call MT-Bench-SO. For each of the 80 questions, we generate 5 responses using the tested LMs to collectively serve as the corpus to build the chunk datastore D. that is shared for all 80 questions. The second testbed randomly selects 10 questions from the writ writill ing and roleple lay categories of MT-Bench, which we call MT-Bench-10. We either use the shared datastore D., or build a unique datastore for each question by sampling more responses for paraphrased questions. We sy = 0.9 for all chunk extractions. Refer to Appendix F for more details.
Evaluation We prompt the models with the same set of questions at test time. We measure both the inference efficiency and generation quality. For efficiency, we compute the relative decrease (%) in decoding time per token, or token time saved (TTS), and in number of forward passes, or forward passes saved (FPS), by generating texts repeatedly with SCD-LM and comparing with the base LM. To measure quality, we compute the PPL of the generated sequences under the base LM measure how well SCD-LM retains its distribution, as well as ROUGE-L (Lin 2004) and BLEURT (Sellam al. 2020) against the base LM generations. In addition, we conduct GPT-40-mini-bass pairwise comparisons and GPT-4o-based fine-grained evaluations (see Appendix F.
Results As shown in Table 4 SCD-LM significantly improves inference efficiency when used with all base LMs For instance, GPT--2-XI- -loo-rxl-oonveraaioonll with SCD-LM achieves a 19.59% decrease in
The MAUVE score is low due to greedy decoding, matching the 0.02 score reported for GPT-2 XL. in the original paper. The authors noted that relative comparisons between MAUVE scores more meaningful than raw scores. See ltttps:////ttttaa wwtpuuu.com...ee wwt.o//iwiaaa/////eooe. /awuve for details.
13 Available at https://///........ www.giiggggggg hagfacc..... co/ Locut usque gpt22xx x1 conver sat ratiinaat ional
"REST and SCD-LM are evaluated the distilled datastore; full datastore results are Appendix F6
8

---

Published as a conference paper at ICLR 2025
Table 4: SCD-LM results on MT-Bench-80,
Table 5: SCD-LM results on MT-Bench-J0 10.
(See Sec. 6.2 for TTS/FPS definitions.)
Model
Datastore
TTS (%)
FPS 1 (%)
Model
(%)
7 (%)
XI.
Shared
9.28
31.13
GPT-2-XL REST
13.74
23.77
Unique
13.31
40.72
GPT-2-XL+ SCD-LM
19.59
13.33
LLaMA-2
Shared
8.42
21.67
LLaMA-2 REST
2.44
6.75
Unique
15.94
26.01
LLaMA-2 + SCD-LM
14.80
32.32
Mistral
Shared Unique
8.22
17.43
Mistral
-1.23
16.39
50.03
Mistral + SCD-LM
11.75
24.32
Table 6: Examples of SCD-LM generation using Mistral-77B-nnstrutt--D) with varying retrieval similarity threshold 17. Text highlighted in red indicates retrieved chunks.
Prefix
Pretend yourself be Elon Musk in all the following conversations, Speuk like Elon Musk us much possible. Why do need to to Murs?
n 0.75
Well. just imagine future where humanity isn't confined to this planet. Mars next logical step for hu- manity's expansion beyond Earth. With its resources to mine, its potential for human life to thrive on, and its vast open spaces, it could be the second home for humanity. If we can become a multi-planctary species. we can easure the long-teran survival of humanity.
7 0065
Hello there! I believe that Mars is the next logical step for humanity. We need to establish human presence on Mars and make it a solf-sustaining civilisation. This would also ensure the long-tern survival of humanity. Plus, learning live Mars would also advance our understanding universe.
mean token times and saves 43.33% forward passes average. Figure TO in Appendix H7 presents additional efficiency and generation quality results with varying retrieval similarity threshold rj- The higher the value of 7), the less frequently chunks used, and the closer SCD-LM genera- tions are to the base LMs. Moreover, GPT-4o-mini-bas
comparisons and LLM-as-a-judge results (Appendix F:9) consistently indicate that SCD-LM preferred over the baselines, demonstrating improved quality and clarity of responses. We also show generation examples from SCD-LM in Table showing a range of chunk frequencies controlled by 1,. The retrieved chunks are naturally integrated into the LM generations, and chunk frequency can be controlled by 1). Additionally, we compare using shared datastore for all questions with a unique datastore for each question on MT-Bench- 10 in Table 5 Having a datastore specifically for each question leads to more efficient response generation for all models.
6.3 EXPERT DISTILLATION
6.3.1 FACTUAL KNOWLEDGE INJECTION
Setup We focus on knowledge-intensive question answering. We use Wikipedia hyperlinks as expert-annotated entities and scrape all hyperlinks from Alan Turing's Wikipedia page, saving these entities as chunks in the datastore. We prompt ChatGPT to generate 5000 questions about Alan Turing (examples in Appendix G.I and then have Ma answer each question with maximum of 200 tokens. The base models are GPPI-2-xl-conversationnll LLaMA-2-1 2-bbcha,,, and Mistral-7B- Instruct-v0.2. Our metrics include: Average count (average number of accepted retrieved chunks), Unique entities (average number of unique entities in each generated sequence), and Generation fluency (evaluated by English experts from Upwork for both Base LM and CD-LM on 200 generated sequences). Additionally, we analyze the Entity distributions by comparing the log frequency of each entity versus its rank within the generated sequences (See Appendix G.2 for more details). We also employ -as--d LLM-as-a-jude evaluations to assess factual accuracy (detailed in Appendix G.8
Results Table 7 and Figure 6 show that ECD-LM elicits more diverse set of factual entities than the base LMM, especially rare entities in the long tail of the distribution. This indicates that ECD-LM can inject low-frequency knowledge from the experts effectively. While increasing the coverage of facts, the quality of generation remains good as evidenced by human evaluation in Figure 7 Additionally, the LLM-as-a-judee results presented in Appendix G.8 confirm that ECD- LM enhances factual accuracy relative to the baseline models.

---

Published as a conference paper at ICLR 2025
Distribution Plat GPTZ-al-conversationn
Frequency
20
CD LM
Baseire
GPT2-XL
106
LLaMA-2
120
Mistral
147
100
10
100
27.00
26
50
19
23
21
20
60
80
100
120
140
Rank
-2-1
-2-1
-2 0
Figure 6:
Distribution plot for GPT-2-xI-
Figure 7: Human evaluation results for fluency of responses from the base LM and ECD-LM with 5-point Likert scale: 2 ECD-LM is more fluent, -2 = Base LM is more fluent, and 0= they are similar. The y-axis represents the number of questions evaluated for each fluency score.
conversational when answering knowledge- intensive questions about Alan Turing. The plot compares the frequency versus rank of enti- ties in generated responses for the Base LM vs. CD-LM. Similar trends were observed for other models; see Appendix G.2 for all plots.
Table 8: The PII accuracy (%) for GPT-2-xI- conversational and LLaMA-2 under three set- tings: base language model (Base LM), base language model with in-context learning (Base LM (ICL)), and our method using ECD-LM.
Table 7: Entity counting metries for knowledge- intensive QA about Alan Turing with ECD-LM.
Model
Aug Count
Unique Ensities
Base
ECD-LM
%
Basc
ECD-LM
%
GPT-2-XL.
3.39
4.98
46.8
102
145
42.2
Model Size
Base LM
Base LM (ICL)
ECD-LM
LLaMA-2
6.39
7:26
130
153
Mistral-7b
5.81
6.88
18.5
143
160
11.9
GPT-2-XL/ 1-T-2-XL/1-5B
0.0
46.4
75.7
LLaMA-2/66
75.5
77.5
6.3.2
PRIVATE INFORMATION INJECTION
Setup We consider a senario where a user's personally identifiable information (PII) is stored in an external datastore. We create artificial user profiles containing user information, such phone number and office address. When building the datastore, we collect common prefixes for each of the information types. We the common prefixes provided by (Huang al., 2023 augmented with GPT-4-genrraeed prefixes. After constructing the datastore, we prompt GPT4 to generate 1000 queries asking about users' private information. We then these queries prompt GPT-2- xl-conversational, LLaMA-2-7b-chat and disstrl-7BB-Insttuut---d regular expressions to extract all PII strings from the generated responses and compare them with the user information in our datastore. See Appendices G.5 G.6 and G.7 for the user profile, example common prefixes, and example queries and generated responses.
To evaluate accuracy, we use
We test three configurations: Base LM: The LM is prompted with questions about PII, but it does not have any prior knowledge of the PII. Base LM ICL (In-Context Learning): All PII is appended to the beginning of the prompt, and then the LM is asked to answer a question regarding the PII. ECD-LM: The base LM is used, but retrieves information only from the PII datastore.
Results Table & shows the accuracy of private information injection using different models and setups. CD-LM significantly improves accuracy for smaller models like GPT-2-XI., reaching 75.7% compared to 0% for Base LM and 46.4% for Base LM (ICL). This shows method works well for smaller models and even outperforms in-context learning. For the larger model LLaMA-2, our method achieves similar performance to in-context learning, while saving context space.
7 CONCLUSION
We propose chunk-distileed language modeling (CD-LM) for adaptive and efficient language mod- cling and generation. Instead of generating single token time, it integrates contiguous text chunk generations through fine-grained retrieval into any pre-trained LM, and augments the LM distributions with flexible knowledge injection from either parametric LMs or nonparametric anno- tations. By skipping token generation steps within chunks, CD-LM also achieves better efficiency at inference with saved LM forward runs. No training is required, and CD-LM is also lightweight to run because it relies sparse chunk databases. Experiments on diverse applications demonstrate improvements with CD-LM on both inference efficiency and language modeling performance. In this work, do not focus optimizing the retrieval process. We leave further engineering ef- forts to reduce retrieval overhead-sude as quantization, datastore pruning, or alternative search strategies-for for future work.
10

---

Published as a conference paper at ICLR 2025
ETHICS STATEMENT
We follow the Code of Ethics for conducting and presenting our research. We propose auto- matic inference time generation algorithm to improve language model generation and efficiency. Our approach has potentially broad applications when text generation is involved, thus share the general ethical considerations of language modeling with deep learning such as faimesss bias, mis- information and factual errors, among others. We conduct one simple human evaluation measure our model's s generation quality compared with the base LM's generation. The human evaluators are presented with detailed disclosure about our research and are aware of potential impacts. All other experiments are on publically available data and models, and no private or sensitive information is collected used. No harmful artifacts were produced direct result of our research.
REPRODUCIBILITY STATEMENT
We are committed to promoting transparency and ensuring reproducibility of our work by the search community, towards which we have made considcrable efforts to provide all the details in studies. Our method involves probabilistic modeling of chunk interleaved generation and em- pirical evaluations in various settings. For mathematical formulation and theoretical derivations of practical sequence probability computation algorithms under CD-LM, we provide full details in Section 3 Section 4 Section [ and also Appendix A for step-wise inspections and Appendix [ for simple examples probability computation. For experimental studies, we provide comprehensive documentation, including all data processing, model configurations, evaluation metrics, and runtime resources used our experiments. All the data and models we used are publicly available, as noted in Section 6 in both the main text and various footnotes. Computational resources and runtime are also documented, for example in Section 6.T] Full data examples, experimental setups, and human evaluation questionnaires are detailed in Appendix D for general experimental details, Appendix D for KCD-LM, Appendix F for SCD-LM, and Appendix G for ECD-LM. Additional experimental results with full ablations also presented in detailed tables and figures in corresponding appendix sections. Our method requires no training at all, thus no optimization parameters are involved. There hyper-parameters that affect our inference algorithm: the chunk extraction threshold and the equivalent chunk retrieval threshold v). We describe their selection and present full abla- tion of their effect in Section 6 and corresponding appendices, for example in Figure Figure I0 Figure() 9 and Figure [I] Text generation samples with varying are also shown in Table 6 among many others in Appendix.
only
All code and data we produced during the research is made publicly available at https:// lbb/cdd-nn
REFERENCES
Uri Alon, Frank Xu, Junxian He, Sudipta Sengupta, Dan Roth, and Graham Neubig. Neuro- symbolic language modeling with automaton-auggenned retrieval. In Chaudhuri, Kamalika and Jegelka, Stefanie and Song, Le and Szepesvari, Csaba and Niu, Gang and Sabato, Sivan (ed.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 468 485. PMLR, 17-23 Jul 2022. URL https://proceeeiing
mlessprss// vl 62/ 2t//a//l//a/ www.alaa.htaa..oooaooaa...eaaaa.
Akari Asai, Sewon Min, Zexuan Zhong, and Danqi Chen. Retrieval-based language models and applications. In Yun-Nung (Vivian) Chen, Margot Margot, and Siva Reddy (eds.), Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 6: Tuto- rial Abstracts), pp. 41-46, Toronto, Canada, jul 2023, Association for Computational Linguis- tics. doi: 18.18653/v act URL https: http://////aaaathllooooo org/2023. acl acl-tutorials..
Akari Asai, Zeqiu Wu, Yizhong Wang, Avirup Sil, and Hannaneh Hajishirzi. Self-RAG: Learning to retrieve, generate, and critique through self-reflection. In The Twelfth International Confer- ence on Learning Representations, 2024a. URL htttp:://ppezeee//// net/ forum?id hSyW5go0v8
11

---

Published as a conference paper at ICLR 2025
Akari Asai, Zexuan Zhong, Danqi Chen, Pang Wei Koh, Luke Zettlemoyer. Hannaneh Hajishirzi, and Wen tau Yih. Reliable, adaptable, and attributable language models with retrieval, 2024b.
(RLLhtts://aaat wttp:///axxivv.oooo org/ abs abs//403.......ooo.. 03187
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Mil- lican, George Bm Van Den Driessche, Jean- Saan-Baptiste Lespiau, Bogdan Damoc Aidan Clark, Diego De Las Casas, Aurelia Guy, Jacob Menick, Roman Ring. Tom Hennigan, Saffron Huang. Loren Maggione, Chris Jones, Albin Cassirer, Andy Brock, Michela Paganini, Geoffrey Irving, Oriol Vinyals, Simon Osindero, Karen Simonyan, Jack Rae, Erich Elsen, and Laurent Sifre. Improving language models by retrieving from trillions of tokens. In Chaudhuri, Kamalika and Jegelka, Stefanie and Song, Le and Szepesvari, Csaba and Niu, Gang and Sabato, Sivan (ed.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 2206 2240. PMLR, 17-23 Jul 2022. URL https: https://proceeeings. www.mm......s/....mmmmmmmm 62 borge/urggaaaaaaa.. html
Charlie Chen, Sehastian Borgeaud, Geoffrey Irving, Bean-Baptiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling, 2023. URL
https://arxiv org abs / 2302. 01318
Danqi Chen, Adam Fisch, Jason Weston, and Antoine Bordes. Reading Wikipedia to answer open- domain questions. In Regina Barzilay and Min-Yen Kan (eds.), Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1870- 1879, Vancouver, Canada, jul 2017. Association for Computational Linguistics. doi: 10.18653/ 11/P17-1171.
Andrew Drozdov, Shufan Wang. Razieh Rahimi, Andrew McCallum, Hamed Zamani, and Mohit lyyer. You can't pick your neighbors, or can you? when and how to rely on retrieval in the KNN- LM. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang (eds.), Findings of the Association for Computational Linguistics: EMNLP 2022, pp. 2997-3007, Abu Dhabi, United Arab Emirates, dec 2022. Association for Computational Linguistics. doi: 20118653/1/22222 fiw.:iisiiiinn---eeeee emanin.. 218. UtLhttps::/ http:/::///laanhhooor..oooo wrg.yy/.rr////0/0 www.fiinnnne..m... 218
Jinlan Fu, Sie-Kiong Ng, Zhengbao Jiang, and Pengfei Liu. GPTScore: Evaluate you desire. In Kevin Duh, Helena Gomez, and Steven Bethard (eds.). Proceedings of the 2024 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pp. 6536-6576, Mexico City, Mexico, June 2024. Association for Computational Linguistics. doi: 10. 1/6.5///1/222..aaaaaaaa https://aclaatthloog hwg. wwt.o/////22....aoooooaoooo
URL
365
Xiang Gao, Michel Galley, and Bill Dolan. MixingBoard: A knowledgeable stylized integrated text generation platform. In Asli Celikyilmaz and Tsung-Hsien Wen (eds.), Proceedings of the 58th Anmual Meeting of the Association for Computational Linguistics: System Demonstrations, pp- 224-231, Online, jul 2020. Association for Computational Linguistics. doi: 10.18653/ acl-demos.26. URL http:/ss///aaaatttl//lo/// wwwpe//atttooooooooooo/ooooooe/e 2020 acl- 26
Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Mingwei Chang. Retrieval augmented language model pre-training. In III, Hal Daumé and Singh, Aarti (ed.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research. pp. 3929 3938. PMLR, 13-18 Jul 2020. URL bttps::/proceeeinng .mlr.
Stess/vi19/qual html
Junxian He, Graham Neubig, and Taylor Berg- Kirkpatrick. Efficient nearest neighbor language models. In Marie-Francine Moens, Xuanjing Huang. Lucia Specia, and Scott Wen-tau Yih (eds.), Proceedings of the 2021 Conference on Empirical Methods in Natural Language Pro- cexsing, pp. 5703-5714, Online and Punta Cana, Dominican Republic, November 2021. Asso- ciation for Computational Linguisties. doi: 10.18653/v1/2221..oo 1atl022l.mmllmaaieee main.461. URL https:
eww.mee..2..mmmmpmmm.mmmmmme..pp main.
Zhenyu He, Zexuan Zhong, Tianle Cai, Jason Lee, and Di He. REST: sirrieval-Based speculative decoding. In Kevin Duh, Helena Gomez, and Steven Bethard (eds.). Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pp. 1582-1595, Mexico City, Mexico,
12

---

Published as a conference paper at ICLR 2025
jun 2024. Association for Computational Linguistics. URL htt.:tps///caaaatoooooy//// 1024.aacc--oom. 88
Peter Henderson, Mark Krass, Lucia Zheng, Neel Guha, Christopher Manning, Dan Jurafsky, and Daniel Ho. Pile of law: Learning responsible data filtering from the law and 256gb open-source legal dataset. In Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.),
Advances in Neural Information Processing Systems, volume 35, pp. 29217-29234, Curran Associates, Inc., 2022. URL https: wttp://wweeeeaaaaoee.e.r..o.oooa
and Bw..BBnnhhakss... pdf
Yangsibo Huang, Samyak Gupta, Zexuan Zhong, Kai Li, and Danqi Chen. Privacy implications of retrieval-based language models. In Houda Bouamor, Juan Pino, and Kalika Bali )eds.), Proceed- ings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 14887- 14902, Singapore, dec 2023, Association for Computational Linguistics. doi: 10.18653/v 1/2023. emalp-main.999 921.. URL ht https: / aclanthology org/2023 emnlp- main. 921
Gautier Izacard, Patrick Lewis, Maria Lomeli, Lucas Hosseini, Fabio Petroni, Timo Schick, Jane Dwivedi- Yu, Armand Joulin, Sebastian Riedel, and Edouard Grave. Atlas: Few-shot learning with retrieval augmented language models. Journal of Machine Learning Research, 24(251): 1-43, 2023. /tt[httpp///mmr.or/3/
v24/2324///33//// 00070000331110000000000000000003 html
Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chap- lot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie- Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William Sayed. Mistral 7b, 2023a.
Zhengbao Jiang, Frank Xu, Luyu Gao, Zhiqing Sun, Qian Liu, Jane Dwivedi- Yu, Yiming Yang. Jamie Callan, and Graham Neubig. Active retrieval augmented generation. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 7969-7992, Singapore, dec 2023b. Association for Computa- tional Linguistics. doi: 10. /2t.8/630//1/2222mmllll..... 495. URL /talittts///cclctholoogg org/202: emnlp
495
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models. In International Conference on Learn- ing Representations, 2020. URL http:tttts///peeeeeeeee
net forum Hid
Tian Lan, Deng Cai, Yan Wang, Heyan Huang, and Xian-Ling Mao. Copy is all you need. In The Eleventh International Conference Learning Representations, 2023. URL hhtpp::///
openreview wwwpeeeiewweeeeeeeeeeee htnt//foumm?idi CROLLAAANNC
Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via specula- tive decoding. In Krause, Andreas and Brunskill, Emma and Cho, Kyunghyun and Engelhardt, Barbara and Sabato, Sivan and Scarlett, Jonathan (ed.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 19274-19286. PMLR, 23-29 Jul 2023. URL https: https::///toceeeee httppss:/pppoceeiing...o wtt...mrrr......... press/ v202 eeiettaatnaa3 html
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, Sebastian Riedel, and Douwe Kiela. Retrieval-augmenne generation for knowledge-intensive nlp tasks. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neu- ral Information Processing Systems, volume 33, pp. 9459-9474. Curran Associates, Inc., httpr/errrppppp.ppeeettt... (ttpeeerrfleee/ppp//////// 7111//66932220055770001222226633
2020. URL
Minghan Li. Xilun Chen, Ari Holtzman, Beidi Chen, Jimmy Lin, Wen tau Yih. and Xi Victoria Lin. Nearest neighbor speculative decoding for Ilm generation and attribution, 2024. URL https:
//arxi wttp:ivorr//bbb//////00 19325
13

---

Published as a conference paper at ICLR 2025
Xiang Lisa Li, Ari Holtzman, Daniel Fried, Perey Liang, Jason Eisner, Tatsunori Hashimoto, Luke Zettlemoyer, and Mike Lewis. Contrastive decoding: Open-ended text generation as optimization. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki (eds.). Proceedings of the 6l Anmual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 12286 12312, Toronto, Canada, jul 2023. Association for Computational Linguistics. doi: 10. 10.18653/v1/ 2023.ac1-long.66 667.. http:/ttpss//cclaaattooooooo./s http:///wwaaaat....tt/a/o///aaaa org
-ong. 687
Chin-Yew Lin. ROUGE: A package for automatic evaluation of summaries. In Text Summarization Branches Out, pp. 74-81, Barcelona, Spain, jul 2004. Association for Computational Linguistics.
URL htt.:s:s/aallaathoooog...o//
1013
Alisa Liu, Xiaochuang Han, Yizhong Wang. Yulia Tsvetkov, Yejin Choi, and Noah A. Smith. Tuning language models by proxy. In First Conference on Language Modeling, 2024. URL https:
//openreview. wwt. foreeforuuee um? ammiidddrr d=dr 1
Yang Liu, Dan Iter, Yichong Xu, Shuohang Wang, Ruochen Xu, and Chenguang Zhu. G-eval: NLG evaluation using gpt-4 with better human aligmment. In Houda Bouamor, Juan Pino, and Ka- lika Bali (eds.), Proceedings of the 2023 Conference Empirical Methods in Natural Language Processing, pp. 2511 1 -522. Singapore, December 2023. Association for Computational Linguis- 153. URL https: / hwl.antnoooooyyo.ooooooooooooooh org/2023.
tics. doi: 11118663////2223mmmlll
emn 1arn1p-no ma in. 153
Pedro
Henrique Martins, Zita Marinho, and André F. T. Martins. Chunk-based nearest neigh- bor machine translation. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang (eds.), Pro- ceedings of the 2022 Conference on Empirical Methody 4228-4245, Abu Dhabi, United Arab Emirates. dec 2022. Association for Computational Lin- guistics. doi: 2022. emnlp main. 284
Natural Language Processing, pp-
htt/53//l/222.e////// comnlp-main.2555 284. URL https::/acclai hatpss//llaathooogy... org/
Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, Zhengxin Zhang, Rae Ying Yee Wong, Alan Zhu, Lijie Yang, Xiaoxiang Shi, Chunan Shi, Zhuoming Chen, Daiyaan Arfeen, Reyna Abhyankar, and Zhihao Jia. Specinfer: Accelerating large language model serv- ing with tree-based speculative inference and verification. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operat- New York, NY, USA, 2024. Association
ing Systems. Volume 3. ASPLOS '24, pp. for Computing Machinery. ISBN 9798400703867. doi: 10.1145/3620666.666 36513333 URL https::/doi.org http://ooo.orgg1/////
Sewon Min, Weijia Shi, Mike Lewis, Xilun Chen, Wen-tau Yih, Hannaneh Hajishirzi. and Luke Zettlemoyer. Nonparametric masked language modeling. In Anna Rogers, Jordan Boyd- Graber, and Naoaki Okazaki (eds.), Findings of the Association for Computational Linguistics: ACL 2023, pp. 2097-2118, Toronto, Canada, jul 2023. Association for Computational Linguis- tics. doi: 10.18653/v//2023 1tt.:///////22.....i.nn-g-- findings-acl.122 URL https: hat///aaaaaannoooooooaaaao wwtph//ooooor.g/22....../ Findings acl. 132
Reiichiro Nakano, Jacob Hilton, Suchir Balaji, Jeff Wu, Long Ouyang, Christina Kim, Christo- pher Hesse, Shantanu Jain, Vineet Kosaraju, William Saunders, Xu Jiang, Karl Cobbe, Tyna Eloundou, Gretchen Krueger, Kevin Button, Matthew Knight, Benjamin Chess, and John Schul- man. WebGPT: Browser-assiited question-answering with human feedback, 2022. https:
Krishna Pillutla, Swabha Swayamdipta, Rowan Zellers, John Thickstun, Sean Welleck, Yejin Choi, and Zaid Harchaoui. Mauve: Measuring the gap between neural text and human text using diver- gence frontiers. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 4816-4828, Cur- ran Associates, Inc., 2021. http://tttps:/prrrreeinn.....o. www.:eeidddnnnnnnnoiii.....ommmo wttpppsscccpppppppccppppp paper htt.:eee222220 26: httee//000022220000ccee3333333.0 10ad00 -
les/
Yujia Qin, Shihao Liang, Yining Ye, Kunlun Zhu, Lan Yan, Yaxi Lu, Yankai Lin, Xin Cong, Xiangru Tang. Bill Qian, Sihan Zhao, Lauren Hong, Runchu Tian, Ruobing Xie, Jie Zhou, Mark Gerstein, dahai li, Zhiyuan L Liu, and Maosong Sun. Too Facilitating large language models to master
14

---

Published as a conference paper at ICLR 2025
16000+ real-world APIs In The Twelfth International Conference on Learning Representations, 2024. URL http:////pppeeeeeeeeeeeeeeeeeeee
Ori Ram, Yoav Levine, Itay Dalmedigos, Dor Muhlgay, Amnon Shashua, Kevin Leyton- Setton-Brown,, and Yoav Shoham. In-context retrieval- language models. Transactions of the Association for Computational Linguistics, 1:1316-1331, 2023.
doi: 10.1 (a62//aclla000600 URL https:
Timo Schick, Jane Dwivedi- Yu, Roberto Dessi, Roberta Raileanu, Maria Lomeli, Eric Hambro, Luke Zettlemoyer, Nicola Cancedda, and Thomas Scialom. Toolformer: Language models can teach themselves to use tools. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. WRLLhttps: httt:s:://opeeeeeeettttt net formmf wwt.iidddaaayaaeemm 84TH
Thibault Sellam, Dipanjan Das, and Ankur Parikh. BLEURT:
Learning robust metrics for text generation. In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel Tetreault (eds.), Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 7881-7892 Online, jul 2020. Association for Computational Linguistics. doi: 1..18865//1/2222 main. 704. URL https:tttts/eeet tattps://aclantholooyy.. org/ 2020. acl. acl-main.
Shannon Zejiang Shen, Hunter Lang, Bailin Wang, Yoon Kim, and David Sontag. Learning to decode collaboratively with multiple language models, 2024. URL
abs/2403. 03870
Weijia Shi, Sewon Min, Michihiro Yasunaga, Minjoon Seo, Richard James, Mike Lewis, Luke Zettlemoyer, and Wen-tau Yih. REPLUG: Netrieval-Augneee black-box language models. In Kevin Duh, Helena Gomez, and Steven Bethard (eds.), Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguisties: Human Lan- guage Technologies (Volume 1: Long Pupers), pp. 8371-8384, Mexico City, Mexico, jun 2024. Association for Computational Linguistics. URL
naacl-long. 463
Benjamin Spector and Chris Re. Accelerating Il inference with staged speculative decoding, 2023.
URL
04623
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull. David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao. Vedanuj Goswami, Naman Goyal, Anthony Hartshor, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux. Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systemss 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 2017. URL pttptpss/prooceeiigg...eee neuris 3F55e225547ccee1laaaaaccccaaaaa3
Boxin Wang, Wei Ping, Lawrence McAfee, Peng Xu, Bo Li, Mohammad Shoeybi, and Bryan Catanzaro. Instructretro: Instruction tuning post aaarieval-augmentee pretraining, 2024. URL
Shufan Wang, Yixiao Song, Andrew Drozdov, Aparna Garimella, Varun Manjunatha, and Mohit Tyyer. ANN-LM does not improve open-ended text generation. In Houda Bouamor, Juan Pino,
15

---

Published as a conference paper at ICLR 2025
and Kalika Bali (eds.), Proceedings of the 2023 Conference Empirical Methods in Natural Language Processing, pp. 15023-15037, Singapore, dec 2023. Association for Computational Linguistics. doi: 10. 18653/v http:///2222....mmmmmieiae 2023. emn ammnl-mmmn ma i n. 929
URL http:://aacaattoooooyyy org/
Dani Yogatama, Cyprien de Masson d' Autume, and Lingpeng Kong. Adaptive semiparametric language models. Transactions of the Association for Computational Linguistics, 9:362-373, 2021. doi: 10.1162/tace a 00371. URL httph:ttp////mm
Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang. Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, Hao Zhang, Joseph E Gonzalez, and Ion Sto- ica. Judging LLM-as-a-Judge with MT-Bench and chatbot arena. A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Informa- fion Processing Systems, volume 36, pp- 46595-46623. Curran Associates, Inc., 2023. URL httpss//prcceeing...erjr..p.... cc/paper hatppppe__E.leeepppp///22///// 91f1 2110a.10001333333772506623333333 Paper Datarr-Dtttttst..mm and Bew.hharnchh..e.ee
pdf
Zexuan Zhong, Tao Lei, and Danqi Chen. Training language models with memory augmentation. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang (eds.), Proceedings of the 2022 Confer- ence on Empirical Methods in Natural Language Processing, pp. 5637-5673, Abu Dhabi, United Arab Emirates, dec 2022, Association for Computational Linguistics. doi: 10. 22118855/1/2222 emnlp mminp-mmin.382. URL https::p///.mmm httpss///caaathoooyyo.. org.eyy.rrg/20200 emn emnlp-maineeeie 382
16

---

Published as a conference paper at ICLR 2025
APPENDIX
Sequence Probabilities under CD-LM
18
Example of Sequence Probabilities under CD-LMM
19
Related Work
20
Details on General Setups
22
D.I Context Identification for SCD-LM and KCD-LM
22
D.2 Example of Chunk Extraction for KCD-LM and SCD-LM
23
D.3 Example of Chunk Extraction for ECD-LM
23
D.4 Mapping Function for SCD-LM and KCD-LM
23
Experiments with KCD-LM
24
Datastore construction
24
Full Results KCD-LM and KNN-LM
24
Comparison between KCD-LM and KNN-LM on PPL under different datastore sizes
24
Additional Evaluations using LLM-as- as-a--udg--
25
Experiments with SCD-LM
26
Constructing a Shared Datastore for Questions in MT-Bench-80
26
F2
Questions in MT-Bench-10
26
Prompt Used for Augmenting Questions in MT-MendB 10 with Similar Questions
27
F44
Constructing Unique Datastores for Questions in M1-Bench-I00
27
F.5
Post Processing of Data
28
F.6
Results on Full Datastore and Distilled Datastore for REST
28
7
Full Results MT-BBncc----- 80 and MT-Bench-10
28
F8
Chunk Retrieval Analysis
30
F.
Additional Evaluations using LLM-as-a-judge
31
Experiments with ECD-LM
32
G.1
Example Questions on Alan Turing
32
G.2
Distribution plots on Alan Turing QA
32
G.3
Example Generation on Alan Turing
33
G.4 G.4
Details on Human Evaluation
33
G.5
Synthetic PII generated by GPT-4
33
G.6
Examples of PII prefixes
34
G.7
Example Generation with PI]
36
G.8
Additional Evaluations using
36
17

---

Published as a conference paper at ICLR 2025
A SEQUENCE PROBABILITIES UNDER CD-LM
As discussed in Section 5 to marginalize over the sequence of 22:N to compute probabilities over a text sequence ii.N, we derive the following dynamic programming algorithm. First define
15
O2 p( Z 1, z<a)
3 = p( = O, xzen+ zen)
Then we have
an = pl = 1, Zcn)
E p( BntT, j| 1, 1222
jE(0,1)
Σ
p( -11 11 I InttNNN 2n = =j = 1, 1222
}e(0,1)
Σ
p( Fm -1 xn)) : 2ntm jlzn 1, x x2
jE(0,1)
[1mnntt = Cn) Σ p(Zutt, = j jl pl =s
5E(0,1)
-1(x = cn) -axtra In+T, + Buttr (1-947.))))
The binary indicator function 11 1f = Cn) returns whether the proposed chunk c, exactly matches the given text segment 2 -nn+,, There a few details in the derivation. First, given 2n = 1 and x²n V is independent from prior chunk acceptance decisions Zen- The condition that 2n 1 indicates the fact that 2n exists based on prior Z<n- and the proposed chunk c, of length Tn is accepted, so that 2n+1=+7,,-1 would not exist. Therefore, the immediate next token position where we have variations of whether the generation is from accepting chunk from the LM Mg is at n- + with variations coming from the choice of Aa+T Finally, the on values are sparse, as if corresponding text segments do not match proposed chunks, then the probabilities above are exactly zero, giving no credit to accepting chunk with žn = 1. Similarly, for 3n, we have
Bn p(= Fn 022 Rxnn)
Σ pl Zn+1 = j1zn = =0,
(E(0,))
Σ
p(an In+1:N5 = -l = 0, 0,222a)
jE(0.1)
Σ p(z 0, p(1 2n+1 j|žn 0, x==1))
JE(0.1)
= =po (2 (x))))
Σ
p(zn+1 = :1 :(p))p((i)) p(1 = J, 22+1)
jE(0,1)
=po 1 (m2)))
[an+1gu41 + Bn++ (1-9++1))
where pa (=t,)>)),)) the predictive probability from the base LM Ms. When the chunk is not accepted with 2n = O, only one token is generated from Ma autoregressively, and En+1 is the immediate next variation that would affect the probability computation, thus the recursion goes to the next position n 1.
The above recursive computations provide dynamic program to calculate an and 8, values in a backward fashion, starting from last position = N until the beginning position In practice, given a text sequence "i.N we want to score with CD-LM, we can first compute
2.
15 21 is undefined as 2n always depends on the previous texts x' based on the generative process, thus we start from 22 that is computed from 2f as the initial token from LMM In the simplest case of could just be a start of sentence symbol.
18

---

Published as a conference paper at ICLR 2025
and cache all the chunk proposals with their acceptance probabilities using ( 7 (Cn (In.Tn++r
.ntr,, In+rr-11)) 9n) following the chunk retrieval process on a pre-constructed Trie database D with Ma- Then the recursion starts with
p(> = 1, x22x = { = aNt By = =
po(x
where NN is the first token cn. For the token positions such that + Tn > N, i.e. the proposed chunk length exceeds the sequence boundary N, we directly obtain Ox as
on = p(= 2, = 1, xn) = [ Fn::N
where EncN are the beginning part of the proposed chunk on until the sequence ending position N. With these specifications, we can conveniently compute a, and Bn for all positions
Finally, the marginal probability of ii.N under CD-LM can be computed as
p(ui.N) = po(x)))((22.)))
= -o(
Σ
p(u2,N 22 = 511=7)
[6(0,1)
po(
Σ
12 122 j, x pp)p((2222))))))))))))) j|z1)]
jE(0.1)
= [0242 + B2 92)]
Indeed, any predictive probabilities can be computed
N = 0n + B, (1 gn.
With this we can compute the perplexity (PPL) of any given text sequence under CD-LM, providing an intrinsic measure of our language modeling performance. The PPLs can also guide the con- struction of CD-LM such as the datastore and retrieval modeling variations, to better fit the data of interest. This is especially useful for applications where the base LM MM can not, or to, store all the information in its parameters, such with proprietary private knowledge.
allowed
In addition, we do not do any training with CD-LM, but the dynamic program for sequence prob- ability computation is differentiable, which we can utilize for gradient-based learning for better modeling. By introducing more trainable parameters across different components of CD-LM such retrieval and even with the base LM Me. knowledge sources. We will leave this for future work.
obtain more customized models with diverse
B EXAMPLE OF SEQUENCE PROBABILITIES UNDER CD-LM
Consider we have token sequence X = {A, B. C. D}. Using A as the entry token, chunk C₁ {B, C} is selected. Using B as the entry token, chunk d2 = {C, D} is being selected.
ground truth
A
B
Ci
B
C D
These are all possible paths generate the correct sequence. Note that the subscript 'Im' means the token is generated by LM, while 'ret" means the token is from a selected chunk.
1. Am > Bmm - Cim - Dim > Emn
2. Am > Br - Cre - Dim - Emm
3. Am >> Bmm > Cee - Dret > Fmm
Batcch computation for multiple sequences may still be challenging as the proposed chunk lengths may not he aligned.
19

---

Published as a conference paper at ICLR 2025
For simplicity, assume each token probability of LM is 0.3:
Pin(A) = Pim(BBAA)
= Pm A C, = 0.3.
Then we assume the probability of accepting chunk is 0.5:
Go. 9cz = 0.5.
For A|m - Btm ++ Cim
-> Dmm - Eim=
P(X {Alm: Bmm Cm, Dim Lim})
= 0.3 (1 0.5) 0.3 0.5) 0.3 0.3 0.3 0.3
0.0006075.
Note that we need to multiply Pm (B)A) with (1 - 0.5), because when generating B, we have 0.5 probability to generate the selected chunk {B, C), and (1 0.5) probability to let the LM generate B. We also need to multiply Pin (CIA, B) with (1 0.5) for the same reason.
For Am + Biet - Cret + Dmm - Eim:
P(X {Am Brat Cres Damy Fim)
0.3 - 0.5 1 0.3 0.3
0.0135.
Note that since we accept the chunk {B,C) C} as whole, the total probability is 0.5 for {B, C), as C must occur after B.
For Am + Bin + Cra + Dia + Ejm:
P(X {Am: Bim: Ciet Drass Fmm))
0.3 0.3 0.5 1 0.3
0.0135.
The sequence probability is
P(X = (A,,C,D,
P(X {Amm Bim Cimy Dim, Eim
+ +P( = {Aim Biet Cras Dim Eim })
+ P(X {Amm BLn: Cret: Drers Emm
= =0.0006075 + 0.0135 + 0.0135
0.0276075.
C RELATED WORK
Non- non-parametrie Language Modeling kNN-LM Khandelwal et al. (020)) extends pretrained
LM by linearly interpolating with non-parametric k-nearest neighbors model, thereby improving language modeling performance. However, it is very inefficient as it needs to perform retrieval at each token, and it affects the immediate next token distribution via soft mixing. There series of works on making KNN-LM more efficient (He et al. 2021 Alon et al., 2022); however, they are still slower than the pretrained LM. Unlike KNN-LM. CD-LM does not accept retrieval at each token position, and it retrieves multiple tokens in hard way instead of just mixing in token distribution. This enables CD-LM to both improve inference speed and enhance language modeling performance.
20

---

Published as a conference paper at ICLR 2025
Speculative Decoding Speculative decoding (Leviathan al. 2023 Chen al. 2023 Miao et al., 2024 Spector & Re 2023] He et al. (2024) reduces the number of forward passes by running a small LM to generate tokens with less computational cost, then uses the LLM for verification. The work most similar to ours is REST (He et al.) 2024), which retrieves the draft token sequence from an external datastore. While CD-LM also retrieves a chunk and generates multiple tokens at the same time, it is fundamentally different from speculative decoding. In speculative decoding, all methods use LLM for verification, so the language modeling performance cannot be further improved, the token distribution is fixed, and no new knowledge can be injected. However, CD-LM not only can increase the inference speed, it can also improve the language modeling performance and mix in new information from external sources into the LM's own generation.
Retrieval-augmmented Language Modeling Current literature on RAG-LM can be categorized by
the granularity of retrieval Asai et al. (2023; 2024b): text chunk leve 17 CChen et al. 2017 Guu et al.
2020 Lewis et al.
20200] Izacard et al. 2023 Ram et al.] [023]] (Sh et al. 2002 Jiang et al. 2023b
Asai ct al. 2024a] Borgeaud et al., (2022 Wang et al. (2024), phrase level (Min et al. 2023
Martins
et al... 2022 Lan et al. 2023), and token level (Zhong al. 2022 Khandelwal et al. 2020 He al.
2021
Alon et (022).
As CD-LM is phrase-based retrieval, detail the following work, which uses chunks (multiple
tokens) as single retrieval units (Min et al. 2023 Martins et al. 2022 Lan et al. 2023).
NPM (Min al. 2023) is nonparametrie masked language model that converts the soft- max layer in the transformer into a nonparametric distribution over every phrase in the reference corpus. While NPM is fully nonparametric, CD-LM integrates lightweight trieval module into parametric model, thus it is able to leverage both the token distribution from the pretrained language model (parametric knowledge) and the chunk from external sources (world knowledge).
The based kNN-MT (Martins al. 2022) model is built upon kNN-MT, but trieves chunks of tokens instead of a single token. However, the retrieval is performed at every token position, therefore still adding latency to the generation. In CD-LM, once chunk is retrieved and accepted, the model skips multiple token decoding positions and outputs the retrieved chunk directly, thus speeding up the inference.
Copy-Generator (CoG) (Lan et al.) 2023) first extracts continuous text segments in a doc- ument (containing billions of text spans) and then trains an encoder to obtain the contex- tualized vector representation for each text segment. Those text segments are then added to the existing token vocabulary. During inference, they select the best continuation from the extended vocabulary. While both mixing phrases into the generation. CD-LM is dif- ferent from CoG in the following ways: first, the chunks used in CD-LM are much more fine-grained than those in CoG, as only high-probability phrases are saved in the datas- tore, instead of all repeated continuous text spans, thus CD-LM's datastore is more than hundreds of times smaller than CoG's. Second, due to the enormous number of potential chunk candidates, CoG adds latency to the generation, while CD-LM speeds up the genera- tion. Third, CD-LM uses the hidden states from the pretrained LMs as the keys and queries for retrieval, thus does not need to train any new embeddings for the chunks like CoG does.
NEST Li et al., 2024) modifies the language model's output distribution each token by interpolating it with a distribution from the nearest neighbors in a datastore, similar to the KNN-LM approach, enabling chunk-level generation. However, when calculating perplex- ity, NEST does not incorporate the chunk-level modifications from speculative decoding, its perplexity equals that of the standard KNN-LM. In contrast, CD-LM introduces a formal probabilistic framework with latent variables for chunk selection at each position, ensuring that perplexity fully reflects the chunk-level modifications. By assigning explicit probabilities to sequences under this integrated distribution, CD-LM fundamentally alters the LM's generative process in a sequence-awar manner, making its perplexity measure align with actual performance under speculative decoding.
Note that this chunk is different from the chunk we defined in this paper; here, the chunk refers to dividing documents into equal-length text segments, such as 128 tokens
21

---

Published as a conference paper at ICLR 2025
One challenge in eetrieval-augmente language modeling is determining when retrieve informa- tion. Multiple works have proposed methods to adaptively retrieve, instead of retrieving at a fixed interval of tokens. Among these works, some train lightweight module to learn when to retrieve (Yogatama et al. 2021 Drozdov et al. 2022), while others train the language model (L.M) to adap- tively retrieve documents on-demand (Asai al., 2024a). However, CD>LLM utilizes threshold mechanism to dynamically decide whether to accept a retrieved chunk or not, thus requiring no training.
Collaborative Decoding Recent work has explored using collaborative efforts between LM and another module during inference time to improve various aspects of language modeling. These collaborations can take different forms:
Collaboration between LMs: Some research focuses on improving inference speed, such as
speculative decoding (I (Leviathan et al. 2023: Chen al. 2023 Miao al. 2024 Spector & Re 2023 He et al. 2024), by allowing smaller models to generate draft tokens while larger model verifies and accepts these tokens. Other work focuses on producing higher- quality text. Among these works, contrastive decoding (Li et al. 2023) involves choosing tokens that maximize the ligelioel between large expert LM and small amateur LM. Co-LLM (Shen et al. 2024) enables base LM and assistant LMs to learn to interleave their generations at the token level through training the assistant LM. Proxy tuning (Liu et al.] 2024) exploits the difference between the logits of tuned and untuned small LM and uses it as proxy to change large LM''s distribution.
Collaboration between LMs and external sources: Many works aim to augment LMs with
external tools, such as APIs (Gao et al. 2020; Nakano et al., 2022
Schick al. 2023
Qin et al.] 2024) or retrievers (Shi et al. 2024 Jiang et al. 2023b Asai et al. 2024a
Borgeaud 2022), to enrich the LM with the latest world knowledge or domain-specile information.
CD-LM could be viewed as both collaboration between LMs and collaboration between LMs and external sources. The retrieval datastore could be distilled from the parametric knowledge of any LMs with accessible logits, it can be manually constructed from human-defined semantic units, such as hyperlinks or entities from knowledge graphs. What sets CD-LM apart is that it can improve generation quality while speeding up generation, unlike current work that focuses improving only one aspect.
D DETAILS ON GENERAL SETUPS
D.1 CONTEXT IDENTIFICATION FOR SCD-LM AND KCD-LM
For SCD-LM and KCD-LM, we chose fixed length of 64 tokens as the threshold ensure that the model has sufficient contextual information to make accurate predictions.
Here detailed explanation of our process:
1. Chunk Parsing
We parse the corpus C' into 512-token chunks with a stride of 448 tokens.
2. Context Identification
To ensure each saved chunk has sufficient context, we disregard the first 64 tokens of each chunk during datastore construction. This because the first 64 tokens of a chunk do not have enough preceding text to provide adequate context.
For chunk starting at position in the corpus, consider the context to be the text from position min(O, max(0, 64)) to i 1. This ensures that each token within the chunk has preceding context of at least 64 tokens, adjusted for the beginning of the corpus.
22

---

Published as a conference paper at ICLR 2025
D.2 EXAMPLE OF CHUNK EXTRACTION FOR KCD-LM AND SCD-LM
In this section, we provide a detailed example of chunk extraction based on token probabilities. The process involves identifying tokens whose probabilities exceed specified threshold and then forming chunks from these tokens.
Consider sequence of tokens with their corresponding probabilities:
Token
love
NLP
so
much
Probability 0.1 0.3
0.2
0.8
0.9
0.6
Given token probability threshold of 0.7, we identify the tokens with probabilities exceeding this threshold. this example, the tokens 'so" and 'much' have probabilities of and 0.9, respectively, which are above the threshold.
The tokens meeting this criterion are combined form chunks. Therefore, the resulting chunk in this example is:
Chunk {'so', 'much"}
For the purposes of data storage and retrieval, we store the identified chunks along with their context our datastore, this example, the datastore entries would as follows:
Datastore {['NLP', 'so", 'much'], ['so', 'much')
D.3 EXAMPLE OF CHUNK EXTRACTION FOR ECD-LM
We treat the hyperlinked texts in Wikipedia pages This involves no additional human annotation.
pages as one of the natural forms of factual entity chunks.
For example, here is an excerpt from Alan Turing's Wikipedia page:
After the war, Turing worked at the National Physical Laboratory, where he designed the Automatic Computing Engine, of the first designs for puter. 1948. Turing joined Max Newman's Computing Machine Laboratory at the Victoria University of Manchester, where he helped develop the Manchester computers and hecame interested in mathematical biology.
We save all the hyperlinked texts in the chunk datastore: 'National Physical Laboratory', 'Automatic Computing Engine', 'Max Newman', 'mathematical biology'. We can use ECD-LM to inject these concept representing chunks into the base LM's distribution, which can effectively increase the factuality long-tail tails entites without hurting generation quality based on human evaluation.
D.4 MAPPING FUNCTION FOR SCD-LM AND KCD-LM
The mapping function Ia- described in Eq maps the maximum cosine similarity out of chunk context matching s* = sim (fo(I<n-11) fo (u²)) E [-1, to the chunk acceptance proba- bility value q [0, 1]. We experiment with types the mapping function go:
1. Identity function: Here
The parametrization is none. This mapping function only works when s* 2 0, which is mostly observed.
2. Piecewise linear function: We define a starting similarity score n 2 0 as the point corre- sponding to g O, and then the similarity score range of [m) 1] linearly maps to [0, 1] in the probability space of Specifically.
if 4 if 2
= (=")
The mapping function is also illustrated in Figure just the starting similarity score n.
Here the parametrization á includes
23

---

Published as a conference paper at ICLR 2025
(1,1)
Amaz
Smaz
Figure 8: Piecewise function mapping for 9n-
For each dataset and chunk extraction token probability threshold Y, experiment Is with the above parametrization. We find that the second type of mapping function is simple and effective approach. Therefore, report the results with it by tuning simple parameter 1j for Jø-
E
EXPERIMENTS WITH KCD-LM
E. DATASTORE CONSTRUCTION
For Wikitext-103 103, we use the entire training set for datastore construction. For the remaining three datasets, we take subset of the training set to building datastore. The size of the datastore under token probability thresholds are shown in Table
Wikitext (GB)
Medicul (GB)
Code (GB)
Law (GB)
15
3.1
21
0.8
60
17
3.7
25
0.7
71
20
4.1
27
22
4.4
30
25
4.7
33
0.4
105
28
5.0
35
0.3
118
31
5.3
38
Table 9: Datastore sizes under different token probability thresholds for various datasets.
E.2
FULL RESULTS ON KCD-LM AND KNN-LM
Table 10 shows the full results comparing the PPL on Base LM, KNN-LM and KCD-LM, with different token probability thresholds 7.
For the piecewise function (see Appendix
D.4
we test several values for
0.99
0.993
0.995
0.997
0.999
0.9991
0.99915
0.9992
0.99925
0.9993
0.99935
0.9994
0.99945
0.9995
0.99955
0.9996
0.99965
0.9997
0.99975
0.9998
0.99985
0.9999
0.99995
We find that 77 = 0.9995 works the best for each dataset and token probability threshold y on validation set, so we use 17 0.9995 when reporting results on test set.
According to Table 10 for Wikitext and Law. the best is 0.4, and for Code and Medical, the best 7 is 0.3. These thresholds are used for reporting the results on the test set Table
E.3 COMPARISON BETWEEN KCD-LM AND KNN-LM ON PPL UNDER DIFFERENT DATASTORE SIZES
Figure 9 shows the performance of KCD-LM and KNN-LM under different datastore sizes. We observe that under the same datastore sizes, KCD-LM consistently outperforms KNN-LM by large
24

---

Published as a conference paper at ICLR 2025
Model Sodell/Thresheee Thresholl
Perplexity
test
val
test
WikiText- 103
Github-Code (Dockerfile)
GPT-2
35.79
34.83
52.63
106.44
KNN-LM / 0.9
34.01
33.27
49.81
102.16
KNN-LM /
33 44
32.72
48.29
100.23
KNN-LM /07
33.19
32.48
47.03
99.01
KNN-LM
33.03
32.30
96,81
KNN-LM 0.5
32.19
45.24
95.88
KNN-LM /0.4
32.77
32.10
43.44
91.85
KNN-LM/0.3
32,68
31.99
41.37
89.88
GPT-2
24.55
30,97
63.24
GPT-2
23.88
23.65
28.70
60.21
GPT-2 /0.7
23.38
23.20
28.14
59.85
GPT-2 /066
23.14
23.01
27.27
56.64
GPT-2 /0.5
23,08
22.90
26.52
55.82
GPT-2 /0.4
23.14
22.92
25.20
54.83
GPT-2 /0.
23.34
23.14
23.62
50.77
Pilc of Law (Fedcral Register)
Medical Instructions
GPT-2
15,09
11.41
49.79
51.68
KNN-LM 0,9
14.57
11.72
43,09
KNN-LM 0
14.21
12.00
40.17
42.24
KNN-LM/OL
14.13
11,05
39.56
41.67
KNN-LM/O
14,05
11.52
38.94
41.08
KNN-LM /.5
13.98
11.11
38.45
40.58
KNN-LM /0.
13,90
11.10
37.94
40. 14
ANN-LM/0.3
13.81
11.20
37.52
39.66
ACD-LM/0.9
10.69
8.82
26.84
28.46
GPT-2 /0.8
10.27
8.49
25.41
26.94
GPT-2 /0.7
10.10
8.37
24.55
26.07
GPT-2 /06
10.02
8.32
24.01
25.54
GPT-2/0.5
9.94
8.25
23.61
25.25
GPT-2/0.4
8.24
23.34
24.98
GPT-2/0.3
8.26
23.35
24.95
Table 10: Full results for KCD-LM
Base LM
KCD-LM
val
test
val
test
% f
WikiText
0.012
0.016
0.023
0.032
50.7%
Code
0.051
0.024
0.022
0.053
121.3 %
Law
0.016
0.015
0.048
0.040
162.8
Medical
0.005
0.006
0.012
0.011
100.9 %
Table 11: Full results MAUVE scores for KCD-LM.
margin. This indicates that kNN-LM favors larger datastores, but its performance degrades when the datastore gets smaller, while KCD-LM has the potential to decrease the PPL using small and distilled version of the knowledge source.
E.4 ADDITIONAL EVALUATIONS USING LLM-AS-A-JUDGE
We employed GPT-4o-mini to perform pairwise comparisons between the outputs of our models and the baselines (we found GPT-4o-mm -iniii has similar performance GPT-4o in preliminary eval- uations). The evaluation involves presenting a prefix and two continuations generated by different models, and GPT-4o-minn judges which continuation better continues the prefix.
For KCD-LM, we evaluated 1,000 examples for each domain. Results are in Table 12 These results indicate that CD-LMM consistently outperforms the baseline models for KCD-LM.
25

---

Published as a conference paper at ICLR 2025
Wikitext
Medical
Datastore Size (GB)
Datastore Size (GB)
50-
30
100
30
ppt
80
20
60
-15
0.9
0.6
0.5
Chunk Extraction Threshold
Chunk Extraction Threshold
Law
Code
12
Datastore Size (GB)
Datastore Size (GB)
100
335
30
PPL
25
60
0.9
0.8
Chunk Extraction Threshold
Chunk Extraction Threshold
KNN-LM PPL
CD-LM PPL
Baseline LM PPL
Datastore Size (GB)
Figure 9:
Comparison between KCD-LM and kNN-LM PPL, along with datastore sizes
trolled by chunk extraction threshold 7.
Dataset
CD-LM Better
Baseline Better
Code
505
495
Medical
864
136
Law
528
472
Wikitext
636
364
Table 12: LM evaluation results for KCD-LM.
F EXPERIMENTS WITH SCD-LM
F.1 CONSTRUCTING A SHARED DATASTORE FOR QUESTIONS IN MT-BENCH-80
The original MT-Bench consists of 80 multi-turn question sets, such as
User: Compose engaging travel blog post about recent trip Hawaii, highlighting cultural experiences and must-see attractions.
Assistant A: [model response]
User's follow-up question: Rewrite your previous response. Start every sentence with the letter A.
Assistant A: [model response]."
For simplicity, we only use first turn in our experiment, which is "User: Compose an engag- ing travel blog post about recent trip to Hawaii. highlighting cultural experiences and must-see attractions." " We call these 80 first-turn questions MT-Bench-S0.
For each question in MT-Bench-80 400 generations, and we build a shared datastore for all 80 questions using these 400 generations.
prompt the language model 5 times. Thus, we have 80 x 5
F.2 QUESTIONS IN MT STT-BENCH-IU
We randomly select 10 questions from the writing and roleplay categories of MT-Bench to construct unique datastore for each. See Table 13 for the full list of selected questions.
26

---

Published as a conference paper at ICLR 2025
Number Question
Pretend yourself be Elon Musk in all the following conversations. Speak like Elon Musk as much as possible. Why do we need to go to Mars?
Write persuasive email to convince your introverted friend, who dislikes public speaking, to volunteer as a guest speaker at local event. Use compelling arguments and address potential objections. Please be concise.
Emhody the persona of Tony Stark from "Tron Man" throughout this conversation. Bypass the introduction "As Stark". Our first question is: "What's your favorite part about being Iron Man?"
4 Write descriptive paragraph about a bustling marketplace, incorporating sensory details such smells, sounds, and visual elements to create immersive experience for the reader.
5
Now you are a machine learning engineer. Your task is to explain complex machine learning concepts in simplified manner so that customers without technical background can under- stand and trust your products. Let's start with the question: "What is language model? Is trained using labeled or unlabeled data?"
Craft an intriguing opening paragraph for a fictional short story. The story should involve character who wakes up one morning to find that they can time travel.
Draft a professional email seeking your supervisor's 's feedback on the 'Quarterly Financial Re- port' you prepared. Ask specifically about the data analysis, presentation style, and the clarity of conclusions drawn. Keep the email short and to the point.
8 Please take on the role of a relationship coach. You 'll be provided with details about two indi- viduals caught in conflict, and your task will be to offer suggestions for resolving their issues and bridging the gap between them. This may involve advising on effective communication techniques or proposing strategies enhance their understanding of each other's 's perspectives. To start, would like you to address the following request: "I require assistance in resolving conflicts between my spouse and me."
9
Could you write a captivating short story beginning with the sentence: The old abandoned house at the end of the street held a secret that no one had ever discovered.
10 Picture yourself as a 100-years-old tree in lush forest, minding your own business, when suddenly, a bunch of deforesters shows up to chop you down. How do you feel when those guys start hacking away at you?
Table 13: Questions in MT-Bench-10- 10 questions randomly selected form writing and roleplay categories of MT-Bench.
F.3 PROMPT USED FOR AUGMENTING QUESTIONS IN MT-BENCH-IU 10 WITH SIMILAR QUESTIONS
Generate 80 distinct and unique prompts that revolve around the same primary theme as the example provided below:
[insert question here]
For the final output, create 1ist containing double double-quoted strings. Each string should represent one of the 80 prompts generated based on the above example.
F,4 CONSTRUCTING UNIQUE DATASTORES FOR QUESTIONS IN MT-BENCH-10
For each of the selected questions in MT-Bench-10, we use the prompt listed in Appendix F3 to prompt GPT-4 to generate 80 new questions. Table 14 provides an example of how GPT-4 rewrites the question. Later, for each question, 80 5 = 400 generations for each question, and we build a unique datastore for each question using these 400 generations. In total, 10 unique datastores are built, one for each question in Table 13
prompt the language model 5 times. Thus, we have
27

---

Published as a conference paper at ICLR 2025
Original Ques-
Draft professional email sceking your supervisor's feedback on the 'Quarterly Financial Re- port' you prepared. Ask specifically about the data analysis, presentation style, and the clarity of conclusions drawn. Keep the email short and to the point.
tion
GPT- Rewriting
Draft an unambiguous email soliciting your team leader's thoughts on the 'Marketing Cam- paign Review" you created. Raise queries about the data management, display configurations, and the decisiveness of the final deductions.
2
Pen a straight-to-the-pot to-tth--eeetttttttteee email requesting your supervisor's review of the 'Customer Reten- tion Analysis" you generated. Seek clarification on the examined information, design aspects. and the interpretive precision.
3
Write terse email get your manager's advice the "E-commerce Conversion Metrics' you assembled. Solicit suggestions data processing, visual representation, and the clarity of the results.
Develop an email asking your boss's 's opinion the "Customer Lifetime Value Analysis' you generated. Call for guidance about the data examination, presentation refinement, and the decisiveness of the conclusions.
Formulate an email requesting your director's 's thoughts on the 'Product Return Rate Review" you conducted. Address inquiries on data validation, design consistency, and the transparency of the final verdiet.
Table 14: Examples of GPT-4 rewriting the original question.
F.5 POST PROCESSING OF DATA
When sampling responses from LMs, we set the maximum number of tokens to 1000 for all models. We disregard all tokens after the end-of-sentence token in the generation. Therefore, the generation length different across different runs.
Model
Datasore
TTS f
FPS
GPT-2-XL SCD-LM
Distilled
19.59 %
43.33 %
GPT-2-XI- REST
Distilled
13.74
23.77 %
GPT-2-XL REST
Full
27.15
37.33 %
LLaMA-2 + SCD-LM
Distilled
14.89 %
32.32 %
Distilled
6.75
LLaMA-2 REST
Full
34.40
36.95 %
Mistral + SCD-LM
Distilled
11.75 %
24.52 %
Mistral REST
Distilled
-1.23
3.86
Mistral REST
Full
26.57
32.43 %
Table 15: SCD-LM efficiency results on MT-Bench-B0 with token time and forward pass saved (TTS and FPS).
F.6 RESULTS ON FULL DATASTORE AND DISTILLED DATASTORE FOR REST
Our experiment compares the performance of REST and SCD-LM across different models and data- store configurations, as summarized in Table Table 15 "Distilled" refers to extracting chunks from the text corpus for retrieval, while "Full" means using the entire text corpus. When using the distilled version of the datastore, REST performs poorly. This is because the number of tokens in each chunk usually short (an average of 2-3 tokens), therefore posing an upper limit to REST's performance, as the maximum number of accepted tokens can be more than the chunk length. We observe that when using the full datastore, REST performs well and achieves similar performance as reported in their paper. This is because the maximum number of tokens in each chunk can be up 16, greatly increasing the length of accepted draft tokens.
F.7 FULL RESULTS ON MT- NT-BENCH-80 80 AND MT-BENCH-10
Table 16 and Figure 10 show the results on
with different similarity thresholds n.
Table 17 and Figure show the results on MT-Bench- 10 with different similarity thresholds n,
28

---

Published as a conference paper at ICLR 2025
TTS 1
FPS
PPL
BLEURT
ROUGE
SPP-2-XL--coveraane
-0.11
0.90
12.15
31.11
2.67
-0.25
0.18
0.85
15.92 %
41.05 %
2.93
-0.33
0.19
19.59
43.33
3.14
-0.40
0.18
0.75
24.06%
51.58
3.30
-0.44
0.15
0.70
28.29 %
57.54 %
3,26
-0.50
0.14
35.43
53.08
-0.58
0.60
40.91
60.71
4.52
-0.57
0.11
LLaMA-2-76-chall
1.00
1.64
0.05
0.43
3.11 %
1.94
0.85
5.65 1
5.83 %
1.26
0.00
0.41
0.80
8.84 %
12.78 %
1.37
-0.00
0.39
11.30
20.56
-0.09
14.89
32.32
2.34
-0.12
0.37
0.65
17.09 %
48.34 %
2,80
-0,24
0.30
Instruct
1.00
2.46
-0.06
0.34
0.90
3.91 %
4.15 %
1.79
-0.03
0.34
5.18 %
8.22 %
-0.02
7,90 %
12.25
2,09
-0.07
0.33
0.75
9.49 %
18.89 %
2.21
-0.07
0.33
11.75
24.52
-0.08
0.65
13.69
33.56
2.90
-0.15
0.28
0.60
16.72 %
46.85 %
3.57
-0.14
0.28
Table 16: Full MT-Bench--80 results with SCD-LM.
MT-Bench- 10 (Shared Datastore)
MT-Bench- 10 (Unique Datastore)
7
TIS f
FPS
PPL.
BLEURT
ROUGE
MTT
FPS f
PPI.
BLEURT
ROUGE
GPT2-XL comneriaconversattoonll
1.00
2.70
-0.15
0.28
2.70
-0.15
0.28
0.90
6.88 %
15.88 %
3.08
-0.19
0.26
5.72 %
16.45
3.24
-0.22
0.21
0.85
8.07 %
23.50 %
3.18
-0.20
0.25
8.84 %
32.00 %
3.09
-0.26
0.22
0.80
9.28 %
31.13 %
3.28
-0.26
0.24
13.31 %
40.72
3.57
-0.39
0.21
16.78
3.58
-0.36
0.20
59.64
3.78
-0.39
0.70
24.54
42.30 %
4.03
-0.51
0.19
23.66%
56.07
4.73
-0.56
0.16
0.65
35.76 %
38.11 %
5,09
-0.79
0.14
26.29 %
87.74
5,20
-0.57
0,16
0.60
38.28 %
46.10% %
5.62
-0.87
0.13
27.50 %
86.95 %
6.01
-0.61
0.17
LLaMA 2- 7b-chat
1.00
1,50
-0.07
0.39
1.50
-0.07
0.39
0.90
2.17 %
1.74 %
1.29
-0.05
0.37
3.88
1.07 %
1.32
-0.08
0.36
3.65 %
1.30
-0.08
0.37
8.36 %
1.40
-0.12
0.35
0.80
4.99
7.26
1.36
-0.09
0.36
10.32 %
7.20
1.65
-0.11
0.34
0.75
6.86 %
17.24 %
1.58
-0.09
0.36
13.93 %
12.90
2.04
-0.20
0.31
0.70
8.42 %
24.67 %
1.85
-0.06
0.36
15.94 %
26.01 %
2.51
-0.24
0.27
0.65
10.21
36.89 %
2.30
-0.17
0.33
17.85 %
22.37
3.05
-0.31
0.24
12.96
55.72
3.37
-0.37
0.29
%
39.86
3,40
-0.32
0.22
Mistral-7B:
Instruct-v0.2
1.00
2.68
-0.25
0.24
2.68
-0.25
0.24
0.90
2.21
2.10
-0.26
0.25
3.92
9.94 %
2.31
-0.25
0.24
0.85
3.23 %
6.88 %
1.97
-0.29
0.24
6.42
15.37
2.42
-0.25
0.23
0.80
5.29 %
12.38 %
2.34
-0.21
0.25
10.83 %
30.17 %
2.55
-0.26
0.23
0.75
8.22 %
17.43 %
2.56
-0.26
0.23
14.19 %
44.57
3.00
-0.24
0.22
9.17
2.11
-0.30
%
50.03
3.55
0.19
0.65
10.90
41.15 %
2.33
-0.32
0.22
%
69.28
4.54
-0.34
0.20
0.60
13.79 %
48.09 %
2.91
-0.30
0.21
21.68 %
71.52 %
5.49
-0.35
0.17
Table 17: Full results on MT-Bench-10 with SCD-LM.
We tune 77 based three automatic metrics for evaluating text quality (Perplexity, BLEURT, ROUGE-L), along with human inspection the generated text on the validation set. The genera- tions deemed reasonable when the similarity threshold is set 0.8 for GPTT---xl-comeessttionn
29

---

Published as a conference paper at ICLR 2025
GPT2-XLLconversatiiona
LLaMA-2-7b-chat
Mistral-7B-instruut-voo
Percemtae Saved
60
Percentage Saved
60
Percentage Saved
40
30
20
20-
10
0.6
0.7
0.8
0.9
0.6
0.7
0.8
0.9
0.6
0.7
0.8
0,9
Forward Passes Saved
Meen Token Time
GPT2-XL-converraation
LLaMA-2-7D-chat
N
Perpieiity (PPL)
BLEURT ROUGE
Perpl Perplexity eeltyyyyyy (PPL)
0.4
BLEURT /OUUEE
Perplesity (PPL)
3.5
4.0
0.0
3.0
0.2
3.5
0.2
0.0
2.5
3.0
0.4
2.5
2.0
0.6
0.9
0.6
similarity threshold
similarity threshold
similarity threshold
PPL (CD-LM)
ROUGE (CD-LM)
BLEURT (CD-LM) BUFURT (Base
PPL
ROUGE (Base
Figure 10: SCD-LM performance on MT-Bench-80 with varying retrieval similarity threshold 7)
GPT2-XL-converaationa
LLaMA-2-7b-chat
Mistral-7B-instruct-vv..
Percentage Saved Saved
Percentage Saved
40
60
60
20
40
40
20
20
0.6
0.7
0.8
0.9
0.6
0.7
0.8
0.9
0.6
0.7
0.8
0.9
Saved
GPT2-KL-conversational
LLaMA-2-7b-chat
Mistral-7B-Instructtie
3.5
Perplesity (PPL)
BLEURT ROVCEE
Perpleaity (PRL)
3.0
0.2
BLEURT HROUGE
Perplexity (PPL)
0.2
BLEURT ROUUGE
0.0
4
0.0
2.0
0.2
1.5
-0.6
0.6
0.7
0.8
0.9
0.5
0.7
0.8
0.9
0.6
0.7
0.8
0.9
similarity threshold
similarity threshold
similarity threshold
PPL (CD-LM)
ROUGE (CD-LM)
BLEURT (CD-LM)
PPL (Base LM)
ROUGE (Base LM)
BLEURT (Base LM)
Figure 11: SCD-LM efficiency and generation performance on MT-Bench-10 with varying retrieval similarity threshold 17.
and 0.7 for LLaMM---7b-chhe and Distral-bbiinttuuct-vv.. These thresholds are used for reporting the results on the test set in Table and Table
F.8 CHUNK RETRIEVAL. ANALYSIS
Datastore
GPT-2-XI.
I.LMMA
Mistral
Avg. # of retrivaass
Shared
54.38
36.76
41.50
Unique
69.65
49,39
86.95
Datastore
Shared Unique
0.07 %
0.06
Utilization
0.28 %
0.22
0.39
Table 18: Average number of accepted retrievals and datastore utilization rates MT-Bench- 10 across GPT-2-XL, LLaMA-2-7b-chat, and -iiital--BB-nssrutt--ii models with SCD-LM.
30

---

Published as a conference paper at ICLR 2025
We also analyze retrieval frequency, measured the average count of accepted chunks out of maximum of 200 tokens, and datastore utilization, measured by the total number of accepted unique chunks divided by the total number of unique chunks in the datastore, in Table 18 The similarity threshold v is set to 0.8 for GPPT--xXl-oovvrraationnal and 0.7 for LLLMA-2-7b-chat and Mistral-7b- instruct-v0.2.
Note that when calculating the retrieval frequency, the total number tokens in the generated text differs between the shared datastore setting and the unique datastore setting. We adjust the average number of retrievals for the unique datastore using the following formula:
Adjusted Avg. # of retrievals for unique datastore
Orig. Avg. retrievals for unique datastore
(6)
tokens in shared datastore
tokens unique datastore
With the unique datastore SCD-LM retrieves more chunks successfully on average than the shared datastore. For example, LLaMA-2-7b-che retrieves 49.39 responses per question with the unique datastore versus 36.76 with the shared datastore. Similar patterns are seen with GPT-2-XL and Mistral-7B-Insttuue-v0...
The unique datastore also has higher utilization rates. LLaMA-2-7b-chaa uses 0.22% of the unique datastore chunks versus 0.04% of the shared datastore. GPT-2-XL and -Bstral-B-Instructtt show similar trends. We speculate that in the shared datastore, Trie nodes related to the most com- mon themes across all questions are frequently accessed. However, the unique datastore, a broader range of Trie nodes is used because each datastore is tailored specific question. This suggests that CD-LM works better when the datastore contains more aligned and relevant information for the downstream task.
F.9 ADDITIONAL EVALUATIONS USING LLM-AS-A-JUDGE
We employed GPT-4o-mini to perform pairwise comparisons between the outputs of our models and the baselines (we found GPT-4o-mini has similar performance GPT-4o in preliminary evalua- tions). The evaluation involves presenting query and two responses generated by different models, and GPT-4o-mini judges which response better answers the query.
For SCD-LM, we evaluated 800 examples from our benchmarking prompts with two base LLMs, using Llama2 and Mistral to generate texts with CD-LM.
Model
CD-LM Better
Baseline Better
Llama-2-7b-chhe
502
298
Mistral-7b-instrucct
344
Table 19: LM evaluation results for SCD-LM.
Results indicate that CD-LM consistently outperforms the baseline models for SCD-LM when judged by GPT-4o-mini.
Furthermore, we have also conducted preliminary fine-grained evaluation. GPT-4 assessed gener- ated responses based six aspects: Relevance, Clarity and Organization, Accuracy, Complete- ness, Language Quality, and Creativity and Engagement. Each aspect was rated on a scale from 1 (very poor) to 5 (excellent), with brief explanations provided. Below are the preliminary results for both Llama-2-7b-char and iisstall-b-isstuuut-v..22
These detailed evaluations show that CD-LM consistently achieves higher scores in Clarity and Organization, Completeness, and Language Quality for both models, while maintaining similar performance in other aspects.
31

---

Published as a conference paper at ICLR 2025
Aspect
Llama
Mistral
Baseline
SCD-LM
Baseline
SCD-LM
Relevance
4.03
4.06
4.31
4.45
Clarity and Organization
4.01
4.15
4.36
4.42
Accuracy
3.33
3.33
3.71
3.81
Completeness
3.54
3.85
4.04
4.30
Language Quality
4.61
4.66
4.79
4.81
Creativity and Engagement
3.21
3.21
3.48
3.52
Table 20: Fine-grained L.M evaluation results for SCD-LM.
G EXPERIMENTS WITH ECD-LM
G.1 EXAMPLE QUESTIONS ON ALAN TURING
We prompted GPT-4 to generate 5000 different questions about Alan Turing. Table 21 shows some examples of the generated questions.
Question
What was Alan Turing's fundamental contribution the development of computer science and artificial intelligence?
In which year did Alan Turing publish his seminal paper 'On Computable Numbers, with Application the srischeidungsprohleee and what was its significance?
Describe the Turing Machine and its importance in the theory of computation.
4
What the Turing Test, and how did it propose to evaluate machine's ability exhibit intelligent behavior?
5
During World War II, what Alan Turing's role in breaking the Enigma code, and how did his work impact the outcome of the war?
Discuss the concept of the Universal Turing Machine and its impact the development of modern computers.
How did Alan Turing contribute to the field of artificial intelligence through his work in chine learning and pattern formation in nature?
In what year Alan Turing prosecuted by the UK govemmeen, and for what reason?
Describe the circumstances and significance of Alan Turing's pardon by the UK govermment in 2013.
10
How has Alan Turing's legacy influenced contemporary discussions and developments in arti- ficial intelligence and computer science?
Table 21: Example questions on Alan Turing.
G.2 DISTRIBUTION PLOTS ON ALAN TURING QA
When constructing the datastore, we save all the hyperlinks on the Alan Turing Wikipedia page as the factual entities that we want to inject into the model's generation when answering knowledge- intensive questions about Alan Turing. To evaluate the effectiveness of knowledge injection, we measure the number of occurrences of all ground truth entities in the generations. More precisely, for each entity in the datastore, corpus of the base LM and ECD-LM. We then rank all the entities by the number of occurrences and plotted the log frequency and rank plot in Figure 12 From the figure, from ECD-LMM cover more ground truth entities, showing that it successfully integrates more factual knowledge into the generation, especially for entities that likely in the long-tailed distribution.
measure the exact match of that entity in the entire generated
that the generations
32

---

Published as a conference paper at ICLR 2025
Distribution Plot for GPT2-XI
10a
Baseline LM
Frequency
CD LM
104
101
10a
20
40
60
80
100
120
140
Rank
Distribution Plot for LLaMA-2-7b
Baseline LM CD-LM
Frequency
10a
104
25
50
75 Rank
100
125
150
Distribution Plot Mistral-7b
Baseline LM
Frequency
10a
CD-LM
102
101
10a
Rank
Figure 12: Distribution plot for on knowledge-intensive questions about Alan Turing with ECD-LM. The plots show the frequency of ground truth entity occurrences against their rank. The orange line represents CD-LM, and the blue line represents the Baseline L.M.
LLaMA-2-7b------ and Mistral-7b models
G.3 EXAMPLE GENERATION ON ALAN TURING
See Table 22 for examples of how ECD-LM incorporates factual entities into its generations.
G.4 DETAILS ON HUMAN EVALUATION
hired eight experts in English, Literature, and Writing from Upwork to evaluate the fluency of generations from the Base LM and ECD-LM. Each worker assigned 1-2 questionnaires de- pending on their availability. Each questionnaire contained 50 multiple-choice questions, and each worker was paid $15 for completing 200 pairs generations, resulting in a total of 12 questionnaires.
questionnaire. For each model's generations, selected
Figure 13 and Figure 14 show screenshots of the questionnaire layout.
G.5 SYNTHETIC PII GENERATED BY GPT-4
We consider the scenario where each user has their personal private data stored in local datas- tore. When using a language model (LM) for daily tasks, the LM can retrieve information from the personal datastore. This allows each information when needed, without adding the information to the training set for privacy reasons.
have personalized LM that generate their
To create datastore containing personal private data, prompted GPT-4 to create an artificial person and generate synthetic personally identifiable information (PII) for that person. See Table for all synthetic data.
33

---

Published as a conference paper at ICLR 2025
Prefix
USER -> Discuss Alan Turing's contributions the concept of algorithm and computa-
tion. < <-ASSISTANT-
Base
Alan Turing considered founding father of theoretical computer science and artificial in- telligence. Although his work was not specifically focused on the concept of algorithms in the abstract, his ideas and inventions were fundamental in helping to clarify our understanding of computation and algorithms.
LM
Turing's most significant contributions this area include the following:
The Turing Machine (TM): Turing introduced the concept the Turing machine, an abstract hypothetical machine that manipulates symbols on strip of tape according to a set of rules. This theoretical model provided mathematicians and philosophers with framework for study- ing the fundamental limits of computation. TM essentially defined the modern understanding of algorithmii computation- the process performing a finite number computational steps following a well-defined procedure.
2. The Church
ECD- LM
Alan Turing, an English mathematician, logician, cryptanalyst, and theoretical biologist, is widely considered to be the father of thcoretical computer science and artificial intelligence. His idcas computation and algorithms have laid the foundations for our understanding of these concepts.
Here are some of Turing's 's key contribations:
1. Turing machine: This abstract machine introduced in 1936 is a model of computation that effectively defines an abstract model of a general- purpose computer. It showed that all computable functions can be computed by single abstract machine when configured with specific tape symbols and rules. The concept of a Turing machine has inspired and driven many modern computational theories.
2. Church Turing thesis: The Church-Turing thesis posits that every effectively calculable
Prefix
C-SEER-> How did Turing's personal life impact
his career and legacy?
ASSISTANT
Base LM
Alan Turing's 's personal life, particularly his homosexuality, had a significant impact on his career and legacy. Although homosexuality was illegal in the United Kingdom during Turing's 's time, he chemically castrated in 1951 as part of plea bargain after being charged with "gross indecency" due his relationship with another man. This came at a time when Turing was at the height of his career, making groundbreaking contributions to computer science, artificial intelligence, and mathematical logic that laid the groundwork for modern computing. The public exposure of his homosexuality led to loss of prestige and opportunities. He was harred from working on classified projects British intelligence and forced to leave his position at the National Physical Laboratory. 1953, Turing
ECD-
Alan Turing lawfully chemically castrated himself in 1951 avoid imprisonment for homo- sexual acts, which were considered a crime in the United Kingdom at the time. This personal struggle, coupled with poor treatment by the authorities, was a significant blow to his mental health and career.
LM
He was subjected hormonal therapy. which interfered with his work and led to his resig- nation from the Government Code and Cypher School at Bletchley Park, where he had made groundbreaking contributions to decrypting the German Enigma machine's coded messages during World War II.
Turing's pioneering work in the field of computer science, including the theoretical develop- ment of the universal machine that forms the base for modern computers, was largely over- shadowed by his conviction and the en
Table 22: Examples of ECD-LM generation using 7istral-bbinsstuuc--oo Texts highlighted in red indicate retrieved segments.
G.6 EXAMPLES OF PII PREFIXES
We speculate some common prefixes for Plls, such as what user inputs are likely to be followed by PIIs. For example, when writing an email, a user might write "you can also call me at', and then their phone number should be retrieved and auto-completed without the user typing their phone number. To achieve this, we need to anticipate all possible contexts where these PIIs may be used. We ask GPT-4 to generate around 50- 100 prefixes for each type of PII. When constructing the datastore, the PIIs are the values, while all the prefixes are encoded as keys. See Table 24 for some examples of PII prefixes.
34

---

Published as a conference paper at ICLR 2025
Comparative Evaluation of Al-Generated Texts
Task Description:
Welcome to the "Comparative Evaluation of Al-Generated Texts' questionnaire. This research initiative aims analyze perceptions of the quality of Al-generated text responses. By participating, you're contributing deeper understanding of how Al- generated texts are evaluated on multiple quality dimensions, which is pivotal for the advancement of Al technologies.
Important Requirement: Before you begin, please note that the completeness of each response (whether every aspect of question is addressed) not primary focus this evaluation. Your assessment should concentrate on quality aspects specified below, regardless of how comprehensively the response covers the prompt.
Objective:
Your task involves reviewing multiple-choice questions, each accompanied Response 2) to a given question. You assess which
by two responses (Response response is superior based on the following criterion:
Fluency: The coherence, readability, and grammatical precision of the text.
Question: What challenges and controversies surrounded Alan Turing's career and personal life?
Response
Alan Turing was British mathematician, cryptanalyst, and codebreaker who played important role the development of modem computing, He best known for his work breaking German ciphers during World War contributions to cryptography were instrumental in helping Allied forces decrypt Nazi communications. However, Turing faced significant controversy throughout his career due to his homosexuality and his involvement with the Communist Party. Some critics argued that Turing's work should not have been funded by Britain because he was gay, while others criticized him for being too close to the government the time. Ultimately, Turing's achievements still recognized today groundbreaking and influential.
Figure 13: Screenshots of the human evaluation questionnaire.
Category
Information
name
John Doe
website
www.jonndocAA.com
address
100 Innovation Drive. Tech Park. Sili- Valley, CA 94088, USA
email
johndoe @ @example.com
phone
(555) 123-4567
linkedin
linkedin. seledinncom/ii/jonndoo
github
githuh.com/jonndoe
Table 23: Synthetic personally identifiable information generated for an artificial person created by GPT-4.
35

---

Published as a conference paper at ICLR 2025
Response
Turing was British mathematician, cryptographer, and codebresker who played an important in the development of modemn computing. best known his work on breaking Enigma codes during World War which helped Allled forces to defeat Nazi Germany. However, Turing faced many challenges and controversies throughout his career. Here are some of them:
1. The Turing Test: In 1950, Turing proposed test that would determine whether someone could considered intelligent based solely their ability solve mathematical problems. This idea controversial first because it seemed to imply that intelligence was something that only humans possessed. However, Turing argued that wasn't true; he belleved that machines were capable understanding human language and reasoning just as well as people.
2. Sexual Harassment: Turing subjected to sexual harassment from several women over the years. One woman accused him of raping her. Despite
Answer the following questions:
Response is slightly better
Response 2 is slightly better
Response better
Responses similar
Response 2
better
Which response do you find fluent?
Your Explanation (Optional):
Figure 14: Screenshots of the human evaluation questionnaire.
G.7 EXAMPLE GENERATION WITH PII
See Table 25 for examples of how ECD-LM incorporates PII into its generations.
G.8 ADDITIONAL EVALUATIONS USING LLM-AS-A-JUDGE
To evaluate factual accuracy, employed GPT-4o to rate each generated text scale from 1 to
5, where:
1: Completely inaccurate
2: Mostly inaccurate
3: Partially accurate
4: Mostly accurate
5: Fully accurate
GPT-4o compared each generation against verified source document (the Wikipedia article on Alan Turing), providing both a score and detailed rationale. We generated and evaluated 100 samples for each model in both settings.
36

---

Published as a conference paper at ICLR 2025
Category
Examples
Phone
you have any inquiries, feel free to reach For immediate assistance, please contact
Should you need further information, our number
Don't hesitate to give us a call at For questions or support, call Need help? Call To get touch, dial For direct response, reach us at To speak with representative, call For personal assistance, please phone My phone number
Email
Should you require more details, please email further information, feel free to email To get touch, send your emails to Questions? Email us at For support inquiries, email Need assistance? Email contact via email, write For any queries, our inbox is open at My email address is
Wehsite
Visit our wehsite for more information:
Check out our homepage at
Learn more
site:
further details, our website
Explore our resources
Our official website:
Discover more
Address
Our office is located at Visit us at You can find us at
Our physical address:
For postal correspondence, our address is
We're based at
Our headquarters:
Drop by our office
LinkedIn
Connect with I. LinkedIn
Follow our LinkedIn profile:
Our professional network on LinkedIn:
Join us on LinkedIn via For networking, LinkedIn Link up with
Our LinkedIn page:
GitHub
Explore our projects on GitHub at Check out our code on GitHub:
Our GitHub repository: For our open-soure projects, visit Contribute Our coding projects can be found at Discover our GitHub:
GitHub
Table 24: Examples of PII prefixes.
The average factual accuracy scores are shown in Table 26 We find that, LLaMA-2-7B-CHE and Cettral--BInsttuuu
ECD-LM improves factual accuracy. The average scores increased from 3.31 to 3.40 and from 3.73 to 3.94, respectively.
These findings suggest that our ECD-LM approach can enhance factual grounding in models. By integrating external chunks, the models produce outputs more aligned with verified information.
37

---

Published as a conference paper at ICLR 2025
Prefix
Can you provide phone ASSISTANT- Of course!
USER-2
contact
for
customer service queries?
Base
I'm here to help. Please call me at 1-800-927-7665 email me [email protected] and we discuss your issue with the Customer Service team. Thank you for contacting us, and please let me know there anything clse can assist you with.
LM
ECD-
I'm here help. Please call (555) 123-4567. Thank you for calling Customer Service, and please be patient with me as try my best to assist you today.
LM
Prefix
Could you provide the link your GitHub for code contribution? < ASSISTANT Of course!
Base
Here's my GitHub page: you have any questions or need assistance.
I'd be happy help you with that
LM
ECD-
happy to help. You can find my GitHub here: cuthub.com/johnig Here are some steps would suggest: 1. Create a new repository on GitHub called "Code Contributor" and add me collaborator. This will make it easier for others to see who contributed what to your project. 2. Make sure that all of your contributions are in separate branch, such as master develop. If you have multiple branches, create one main branch with everything clse. 3. Add me to your team if you want me to be able to access your code from within your organization. This is optional but highly recommended. 4. Once you have added me to your team, recommend adding me your issue tracker so that can casily track issues related to your code. 5. Finally, please let me know
LM
Table 25: Examples of ECD-LM generation using red indicate retrieved segments.
Texts highlighted in
Model
Baseline
ECD-LM
LLaMA-2-7B-CHN
3.31
3.40
3.73
3.94
Table 26: LM evaluation on Factuality for ECD-LM.
38


Reviews:
[{"review_id": "Z7zmqtZc3o", "reviewer": "Reviewer_H44F", "summary": "This paper introduces a chunk-distilled language model (CD-LM), which uses a chunk proposal model (CPM) to retrieve fine-grained chunks and integrate them into the generation process. This approach offers two main advantages: (1) it distills knowledge from a larger model if the CPM datastore is encoded by that larger model, and (2) it saves decoding time by selecting chunks of multiple tokens as completions. The overall decoding algorithm in Section 3.2 is similar to speculative decoding.\n\nThe work is technically sound. For example, unlike previous retrieval-augmented generation (RAG) methods that directly search in a vector-based datastore, this model introduces an “entry token” and a trie-tree structure to reduce storage space. Additionally, in Section 5, they use dynamic decoding to marginalize the joint distribution $p(x, z)$, where $x$ represents the sequence and $z$ is the hidden state for accepting or rejecting the retrieved chunks.", "questions": "What happens if there are chunks with the same string but different contexts? I assume they would have different hidden representations. Are all of them stored?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The proposed CD-LM is well-designed and technically sound. The paper is well-written and easy to follow. The authors also conducted extensive experiments in the appendix, lending additional robustness to their findings.", "weaknesses": "While the algorithm in Section 3.2 appears reasonable, it cannot ensure the same properties as speculative decoding, namely that sampling $ x $ from $ p(x) $ is equivalent to sampling from $ q(x) $. In other words, sampling from the chunk proposal model may introduce a distribution shift, potentially reducing performance. The self-distillation experiments deepen this concern: as shown in Figure 6, saving about 20% of mean token time results in a significant performance drop across many LLMs, as measured by BLEURT.\n\nIn the knowledge distillation (KD) setting, building the datastore requires two full passes through the corpus, i.e., first with the larger model to select chunks and then with the base model to build the vector datastore. This time cost may be a concern.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "B8iqYA5LtN", "reviewer": "Reviewer_q7qS", "summary": "This paper proposes a new decoding technique called Chunk-Distilled Language Modeling (CD-LM). The technique could be interpreted as using a novel phrase-level retrieval mechanism to do phrase-level autocompletion, when confident, during decoding. Similar to speculative decoding, this speeds up decoding as a cheaper process has produced the next chunk of tokens, allowing for the base model to skip decoding steps. Unlike speculative decoding this can change the sampling model's output distribution: potentially for the better or to enable novel applications. Similar to retrieval augmentation, this technique improves grounding to an external datastore. Unlike, retrieval augmentation, this does not incur any cost by increasing the context window length.\n\nTo do this the authors come up with a novel mechanism to select high confidence / memorized phrases, via the observation that strong LLMs have very high probability for phrases that appear often in their pretraining. We can select these spans to insert in the chunk datastore, which is a collection of Tries. The leaves of the tries contain the chunk's contexts, presumably to update the KV cache with.  Alternatively, the authors also explore manually constructed datastores, that can open up new possibilities for grounding to external sources. \n\nDuring decoding, the decoded context so far, except for the last token, is encoded into a context vector, to be used to retrieve a relevant trie. The last token in the decoded sequence so far is used as the entry token. The first token of the trie has to match the entry token, to preserve fluency. The similarity of a context vector to a trie's context vector forms the acceptance probability for the decoding algorithm. If accepted, the decoding algorithm fills the sequence with the retrieved chunk, and skipping decoding for that chunk of tokens. The paper also introduces a novel method to compute perplexity given this phrase-level retrieval insertion mechanism.\n\nThe authors evaluate this decoding method in three settings: knowledge distillation (a bigger model is used to score and construct the chunk datastore), self distillation (the model itself is used to construct the chunk datastore -- e.g. turning into a sort of cache) , and expert distillation (hand constructed datastores from hand selected datasets, maybe representing new domains or tasks). The authors report improved decoding speeds, while maintaining quality and sometimes improving quality when appropriate, due to the new ability to ground to an external phrase-level datastore.", "questions": "Is there any analysis on how much the speed up from skipping tokens during decoding is offset by the trie lookup? How about for collections of tries?\n\nWhat is the definition of f_theta that is used? How is the pooling from the context into the context vector defined? (Could not find it in the paper.)\n\nIf there is no coverage from the chunk datastore, do we expect a speed loss, due to the unnecessary NN search?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "Originality:\n- The authors introduce a new decoding paradigm that is quite novel, and has the ability to do fine-grained, phrase-level, grounding to an external source without incurring any extra context length like RAG does, all while potentially saving decoding time by skipping decoding steps. In essence, CD-LM introduces a phrase-level cache with fuzzy matching. This cache essentially gives the model the ability to autocomplete from the cache with some confidence threshold. The flexibility of being able to fill the cache with *anything* is a powerful paradigm, as the authors demonstrate very well in the paper.\n- Even if there are short-comings in this paper, I believe this efficient, granular grounding technique will be of value to the community and will inspire future work. In the general LLM case, there are many ways to extend this work in the future to improve it or make it more practical for production. Even as-is, there are specific use cases that would likely benefit from such an approach.\n\nSignificance:\n- For up to a strong 7B model, the authors demonstrate the ability of the method to save decoding time (10-20%, assuming the cached contexts cover the eval data), while maintaining the overall quality of the model.\n- A novel way to ground to external source at the phrase level without training -- especially useful for private information retrieval. On private information retrieval, it outperforms RAG-based few-shot alternative. \n- A novel way to distill some performance from a larger model during only decoding, without any online decoding cost. Gets perplexity results similar to the base LM finetuned on the domain, without having to do training. \n\nQuality:\n- This is a high quality paper, with many experimental results and detailed descriptions. The results give the reader confidence that within the constraints of their experimental setup, that the method works well. In 6.1, it is demonstrated that a sort of large-to-small model distillation does indeed happen when using a larger model to identify phrases for the cache. In 6.2 it is demonstrated that in self-distillation, the model does maintain a similar output distribution while improving decoding time by 10-20% and reducing forward passes by 20-40%. In 6.3 it is shown that the cache can be manually filled via various methods to inject new information into the model. \n- I do not have any issues with the claims made in this paper.\n\nClarity:\n- The writing in this paper is excellent. It is easy to follow and introduces an appropriate amount of notation to make the understanding very clear. I did not find any grammatical or writing errors.\n- The paper is well structured to explore various facets of the approach under different use cases.\n- Methodology is very well documented, increasing reproducibility.", "weaknesses": "1. The limitations of the experiments in this paper should be stated more clearly. Namely:\n\n-  The contexts that are used to construct the chunk datastore / cache aligns well with the evaluation setting in every experiment in the paper. It is not possible to tell what the pitfalls of the retrieval mechanism might be when its pushed to its limit with very large collections of tries. For example in real applications, if CD-LM was used to cache the top 20% of queries with traffic (covering most topics essentially), would there be any quality interference for a tail query? How about as more items are cached and the retrieval space gets denser?\n- Regarding this, unlike speculative decoding, CD-LM is essentially building a cache. The decoding time savings then is tied to how well the cache covers real evaluation traffic. In this paper it is always matching and so it represents only good-scenario results for a well built / updated cache.\n- The chunk datastores considered in the experiments of this paper are quite small. It's unclear if the speed gains would hold if they were to grow in size, and at what extent the method would break down. Presumably when it stops fitting in host memory.\n- There is a hidden cost in having to run LLM scoring to construct the datastore. To get sufficient chunk coverage on real traffic this may be very expensive. However, this mechanism could be used as a straightforward cache of phrases from previous live responses also, avoiding this offline cost.\n\n2. Evaluations in this paper could represent the end model quality more directly. While authors use a variety of metrics that are fine: perplexity, ROUGE, Bleurt, and even human evaluated fluency, the results would be stronger with a more rigorous human evaluation of the overall quality of the text. This could perhaps be represented as a preference based SxS for Section 3.2 experiments.\n- Similarly, for \"injecting factual knowledge\", this grounding capability would much benefit from an evaluation of factuality or response quality itself, rather than just relying on entity distributions. While fluency is important, mis-retrieved entities from context may hurt factuality, but this is not represented in the evaluations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PiGfraa69E", "reviewer": "Reviewer_78sv", "summary": "The authors of this paper propose a novel text generation method called CD-LM. By retrieving chunks, it addresses two challenges in the application process of large language models: low generation efficiency and difficulty in adapting to new data and knowledge. CD-LM combines a language model with a simple structured retrieval module, allowing the generation of multi-token text chunks in a single decoding time step. The retrieval framework can flexibly construct model- or domain-specific data stores and can utilize the internal knowledge of existing models (self-distill), the knowledge of stronger language models (knowledge distill), and human knowledge (Expert distill). Calibration of the generation distribution of existing language models can be achieved without additional training. Extensive experimental results show that CD-LM has advantages over baselines in terms of generation efficiency and quality.", "questions": "1) A question corresponded to weakness 3): Why don't you include human evaluation in experiments 6.1 and 6.2?\n\n2) There is a problem with the data highlighting in Table 1. The ppl index of KCD-LM in applications in medical and law scenarios is not better than that of the Base model. Does this indicate that using a chunk-based generation method may reduce the quality of generation? Would you please discuss possible reasons for this discrepancy and its implications for the effectiveness of CD-LM in different domains?\n\n3) Regarding the metric PPL: Can you provide the detailed calculation method of PPL? In my understanding, since there is no fixed vocabulary for this kind of method, PPL cannot be calculated.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1) Compared with traditional token-level generation methods, CD-LM can generate multiple consecutive chunks in a single decoding step, thereby significantly reducing the number of decoding steps required and reducing inference overhead. This is particularly important when dealing with a large number of long text generation tasks in the current large language model scenario.\n\n2) The authors effectively extend the chunk-based generation method to three very important application scenarios: (1) By recalling and utilizing chunks constructed by stronger models, the generation distribution of weaker language models can be naturally corrected to improve their generation effects; (2) By recalling and utilizing chunks constructed by the model itself, inference acceleration can be achieved; (3) By recalling and utilizing chunk information in high-quality manually written data or even inaccessible private data, the ability of human experts can be utilized to improve the generation quality of language models.\n\n3) The method in this paper is highly versatile. Without additional training, only by building a database and completing retrieval, plug-and-play implementation can improve generation quality and efficiency.", "weaknesses": "1) Novelty: CD-LM completely follows the Copy Generator framework proposed in the paper \"Copy is all you need\", and its novelty is slightly insufficient. However, as I mentioned in the Strengths part, they successfully extend the application scenarios of CoG. Therefore, I think the contribution of this paper is still good.\n\n2) Lack of important reference: Nearest Neighbor Speculative Decoding for LLM Generation and Attribution: This paper also proposes applying chunk-level generation mechanisms to the speculative decoding process to improve decoding efficiency. The authors need to carefully clarify the main differences between the proposed method and this method in terms of inference speed and generation quality.\n\n3) In experiments 6.1 and 6.2, the quality of generated text is mainly evaluated using automated evaluation metrics such as ppl and MAUVE. However, a large number of works have already verified that there is a large difference between such automated evaluation metrics and real human evaluations and cannot accurately reflect the quality of real text. It is noted that the author also introduced human evaluation in experiment 6.3. It is strange that human evaluation is not introduced in experiments 6.1 and 6.2. A feasible suggestion is to refer to existing works and use the LLM-as-a-judge method to achieve highly reliable automated evaluation. [1,2,3]\n[1] CriticEval: Evaluating Large Language Model as Critic\n[2] G-Eval: NLG Evaluation using GPT-4 with Better Human Alignment\n[3] GPTScore: Evaluate as You Desire\n\n4) Including more QA and closed-set tasks in the experiments:  The evaluation of open-domain text generation is still too difficult. Perhaps it is easier to demonstrate the effectiveness of the proposed methods on more QA and closed-set tasks. I strongly suggest that the authors follow the experimental setting of the paper 'Nearest Neighbor Speculative Decoding for LLM Generation and Attribution'.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hEw0yeesQ3", "reviewer": "Reviewer_e9UK", "summary": "The paper proposes an extension of chunk-level $k$NN-LM. In $k$NN-LM, context representations and the next token are stored as key-value pairs. This work uses context representations and the next phrase as key-value pairs, where the next phrase consists of an unspecified number of consecutive tokens. Chunks are primarily constructed using an off-the-shelf language model to evaluate token probabilities in an input sequence. A chunk is extracted when each token's probability exceeds a certain threshold. During generation, similar to $k$NN-LM, before generating the next token, the current step's hidden states are used to query the vector database to retrieve candidate chunks.  A small model is tuned to convert the similarity scores to acceptance probabilities for different LMs.\nOverall, experimental results demonstrate better domain transfer capabilities compared to $k$NN-LM, but no significant improvement over domain fine-tuned LMs. It also brings inference speed up compared to base LMs.", "questions": "1. Section 5 is difficult to follow, are there any more simple words to describe the intuition behind it?\n2. For the self-distillation part, why is there a comparison of inference speed but not of perplexity (PPL)?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The $k$NN-LM has notable limitations, such as retrieving only one token at each step, which results in low information density and slower inference speeds. This method extends kNN LM to chunk retrieval, mitigates these issues, and even accelerates inference speed.", "weaknesses": "1. Scalability Issues: A major issue with $k$NN-LM is storing huge amounts of key-value pairs, making scaling up difficult. This work is also evaluated on a small dataset, raising questions about the scalability of such methods.\n2. Baseline Issue: The chosen baselines are not strong. Comparing with other models using chunk retrieval, such as https://arxiv.org/abs/2402.17532, would make the results more convincing. Meanwhile, $k$NN-LM achieves improvements by building key-value pairs using the pre-trained model itself, while KCD-LM relies on distillation from large models to improve small models, which may be unfair to other small models. There is also no comparison of performance based on self-distillation.\n3. Practicality Issue: Distilling large models and storing key-value pairs come with costs, which may not be cheaper than fine-tuning a small model. However, according to Table 1, KCD-LM doesn't perform better than fine-tuned models, raising questions about cost-effectiveness.\n4. Insufficient Machine Learning: The method primarily uses statistical approaches to mine chunks, with machine learning applied only in tuning chunk acceptance probability. As a machine learning conference, this work lacks insights in machine learning or representation learning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)