Scaling Retrieval-Based Language Models with a Trillion-Token Datastore
Authors: Rulin Shao
Jacqueline He
Akari Asai
Weijia Shi
Tim Dettmers
Sewon Min
Luke Zettlemoyer
Pang Wei Koh
Venue: NeurIPS 2024 poster (2024)
Keywords: retrieval-augmented language model
scaling
trillion-token datastore
RAG
Primary area: natural_language_processing
TL;DR: We study datastore scaling for retrieval-augmented language models with a trillion-token datastore and an efficient pipeline.

Abstract: Scaling laws with respect to the amount of training data and the number of parameters allow us to predict the cost-benefit trade-offs of pretraining language models (LMs) in different configurations. In this paper, we consider another dimension of scaling: the amount of data available at inference time. Specifically, we find that increasing the size of the datastore used by a retrieval-based LM monotonically improves language modeling and several downstream tasks without obvious saturation, such that a smaller model augmented with a large datastore outperforms a larger LM-only model on knowledge-intensive tasks. By plotting compute-optimal scaling curves with varied datastore, model, and pretraining data sizes, we show that using larger datastores can significantly improve model performance for the same training compute budget. We carry out our study by constructing a 1.4 trillion-token datastore named MassiveDS, which is the largest and the most diverse open-sourced datastore for retrieval-based LMs to date, and designing an efficient pipeline for studying datastore scaling in an accessible manner. Finally, we analyze the effect of improving the retriever, datastore quality filtering, and other design choices on our observed scaling trends. Overall, our results show that datastore size should be considered as an integral part of LM efficiency and performance trade-offs. To facilitate future research, we open-source our datastore and code at https://github.com/RulinShao/retrieval-scaling.

Scaling Retrieval-Based Language Models with a Trillion-Toker Datastore
Rulin Shao'
Jacqueline He'
Akari Asai' Luke Zettlemoyer'
Weijia Shi' Pang Wei Koh 1.2
Tim Dettmers!
Sewon Min
Allen Institute for Al (rulins, jyyl,, akari, aakari,ssj041. (e19,41,, dettmers, seeneers,swwo,,,.. sewon, pangweii @cs washington, edu
Uniiversity Washington
Abstract
Scaling laws with respect to the amount training data and the number of parame- ters allow predict the cost-henefit trade-off's of pretraining language models (LMs) in different configurations. In this paper, we consider another dimension scaling: the amount of data available at inference time. Specifically, find that increasing the size of the datastore used by retrieval-based LM monotonically improves language modeling and several downstream tasks without obvious satu- ration, such that a smaller model augmented with large datastore outperforms a larger LM-only model on crowledge-intensive tasks. By plotting compute-optimal scaling curves with varied datastore, model, and pretraining data sizes, we show that using larger datastores can significantly improve model performance for the same training compute budget. We carry out our study by constructing 1.4 trillion- token datastore named M MASSIVEDS, which the largest and the most diverse open-sourced datastore for retrieval-based LMs to date, and designing an efficient pipeline for studying datastore scaling computationally accessible manner. Finally, we analyze the effect of improving the retriever, datastore quality filtering, and other design choices on our observed scaling trends. Overall, our results show that datastore size should be considered as an integral part of LM efficiency and performance tradc-offs. To facilitate future research, we open-source our datastore and code /tthh:::///// h/ttt::://gttub.. com wttp///innnnnaorrrriiiaaaaaaaaaa
Datastore Scaling
Compute-Optimal Scaling (TriviaQA, 5-shot) +
Language Modeling
Downstream Task
(RedPajama)
(MMLU, 5-shot)
70.0
12B
4.4
60.0
Model Size (Pythia)
Accuracy (%)
Perplexity
50.0
4.0
60
Accuracy
40.0
LM Datastore
55
30.0
20.0
3.6
LM-only
0.0
1B
10a
10a
10a
10a
10?
10
10
10
16
Datastore Size Tokens)
Compute (Pretraining FLOPs Datastore FLOPS)
Llama-2
Llama-2
Llama-3 88
- Parato Optimal w/ Retrioval
w/ Retrieval
* LM Only
-
Pareto Optimal LM-Only
Figure 1: Datastore scaling improves language modeling and downstream task performance. Left: Datastore scaling performance on language modeling and a downstream task (MMLU) with LLAMA-2 and LLAMA-3 models. Right: Compute-optimal scaling of retrieval-based language models vs. LM-only models with PYTHIA models. By considering the size of the datastore as an additional dimension of scaling,
improve model performance at lower training cost.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Introduction
The scaling of large language models (LMs) has driven tremendous performance gains across variety of tasks (Brown et al. 2020 Kaplan et al. 2020 Muennighoff et al. 2023). Current scaling laws are primarily a function of the size of the pretraining data and the number of parameters (Hoffmann et al. 2022; Muennighoff et al., 2023) Gadre ct al. 2024). In this paper. we consider another dimension of scaling: the amount of data in datastore used at inference time by retrieval-basee LMs, which can directly retrieve information from the datastore to use in context when generating output (Karpukhin)
ct al. 2020, Guu et al. 2020; Izacard & Grave 2020 Asai et al. 2024b).
Retrieval-basee LMs have a range of benefits such as improved factuality (Mallen et al. (202)). effective domain adaptation (Khandelwal et al.) (020), credit attribution (Gao et al., (023), and parametric efficiency (Min et al., 2023b). However, most prior work in retrieval-basee LMs use datastores constructed from single data source (Karpukhin et al. 2020), such as Wikipedia, with sizes on the order of few billion tokens. While there has been some work on larger datastores (Table[T), with the largest being RETRO (Borgeaud et al.) [022] Wang et al.) 2024) in the trillion-token range, these studies use proprictary datastores and custom architectures with limited evaluation suite. As such, it remains unknown how datastore scaling helps the currently dominant vetrieval-in-conee approaches on a broad categories of tasks.
We first construct MASSIVEDS, a massively multi- multi-domain datastore comprising 1.4 trillion tokens of both general web data and domain specific data ($3.1) that serves as the cornerstone for our scaling study. A key challenge in studying datastore scaling is the computational cost introduced by building datastores with all possible combinations of factors such as the datastore scale, data composition, random seed for subsampling, and different data preprocessing methods. To make our study accessible, we design an efficient datastore construction pipeline that reduces the compute needed by an order of magnitude while being equivalent to the standard pipeline ($3.2).
Using the proposed pipeline, we systematically evaluate the effects of scaling MASSIVEDS on retrieval-based LMs with varying numbers of parameters and pretraining tokens ($4). Beyond upstream language modeling, we also consider a suite of diverse downstream tasks, including general- knowledge question answering (QA), Nomain-specializze first, datastore scaling consistently improves both language modeling and some downstream tasks in a task-dependent manner (Figure [ Left), much like the widely observed data and parameter scaling trends. In fact, on knowledge-intensive tasks, a small retrieval-based LM can outperform its larger LM-only counterparts. Second, since indexing a datastore is cheaper than training on the same amount of data, retrieval-basee LMs enable better compute-optimal scaling trends, where they achieve superior performance than LM-only models at the same training cost (Figure Right).
QA, and reasoning tasks. We find that,
Through our analyses (5)) we show that retrieval-based LMs are capable of automatically retrieving documents that are in-domain to the query, which allows them to reap the benefits of larger, broader datastores. In addition, data quality filters and improved retrieval methods can further enhance our observed scaling trends.
Overall, our results show that datastore size should be considered as an integral part of LM efficiency and performance trade-offs. To spur future research, we open-source MASSIVEDS (including the raw passages, the embedding, and the index) and all code (including our evaluation suite and pipeline for ef- ficiently studying datastore scaling) at https:: //github....... httchuu.coooaaeee woo./RulinShao/eetteeeooo -esttrrieallssllnnnnee
2 Related Work
Retrieval-based LMs. Unlike parametric LMs that only use data during training, retrieval-based LMs can access data through datastore during inference (see Asai et al. (2024b) for review). We focus on strieve-in-conter
language models (RIC-LMs), which retrieves a small set of documents from the datastore and feeds concatenation of them as input to the LM (Ram et al., (203) Shi jtt al., 2023). The RIC-LM approach is simple and allows the use of off-the-shelf retrievers and LMs, even with only black-box access.
Scaling the retrieval datastore. Prior work on retrieval-basee LMs often focused on specific aspects of LMs such as factuality and attribution. In addition, they typically use limited-size, single-domain

---

Table Comparison with prior work, ordered by datastore size. '# Tokens' indicates the number
of tokens in the datastore using the LLAMA2 tokenizer (Tourron et al. 2023). The asterisk (*) denotes that the datastore is not evaluated on downstream tasks. MASSIVEDS the largest open-sourced datastore and covers a broad spectrum of domains.
Reference
# Tokens
Data Sources
Open sourced
ATLAS (Izaaad al.
2023
<5B
Wikipedia
REALM (Guu
2020
<:B
Wikipedia
RALM Ram el al. 2029
<5B
Wikipedia
SELF-R AG (Aaa el al.
2024a
<5B
Wikipedia
REPLUG (hh Cf al. 2023
47B
The Pile
RA-DIT Lin al. 20244
79B
Wikipedia, CommonCrawl
SPHERE Piktus
2022
90B
CCNet
RETRO++
Wang
330B
The Pile, CommonCrawl, RealNews, CC-Stories
INSTRUCTT
RTROO Want G1.
2024
1.2T
Wikipedia, CommonCrawl, RealNews, CC-Stories, Books
RETRO [Borgeaud
2022
1.71
Massive [Rac et al. 2022
MASSIVEDS (Ours)
1.4T
domains, listed Table
datastores such as Wikipedia, which on the order of few billion tokens (Table T Scaling the datastore remains relatively underexplored, with two notable exceptions. First, [Brgeaud et al. (2022) proposed a new RETRO transformer architecture for retrieval-based LMs and built a 1.75 trillion token datastore sourced from the proprietary data introduced in Rae et al. (2022). However, RETRO and its follow-up work, RETRO++ (Wang et al. 2023) and INSTRUCTRETRO (Wang et al., 2024h, use this trillion-token datastore solely for language modeling evaluation, while using small task-specific Wikipedia datastore for downstream task evaluation. Since RETRO-base datastores are not fully open-sourced, it is challenging replicate work based on RETRO to assess the effectiveness of datastore scaling. Second, Piktus et al. (2022) proposed SPHERE, an open-sourced 90 billion-token datastore curated from CCNet (Wenzek al. 2020). However, their evaluation downstream tasks such as KILT (Petroni et al. 20021) suggests that SPHERE does not always outperform small, in-domain datastore like Wikipedia.
In contrast, we present the first study on the downstream performance of trillion-token scale datastores, including analysis of compute-optimal scaling trends using retrieval-bact LMs with different sizes of datastores, models, and pretraining corpora. Our work is fully open-source and can be replicated on limited computational budget, enabling research on trillion-token datastores to be more broadly accessible.
MASSIVEDS and our Datastore Scaling Pipeline
3.1 MASSIVEDS: A Trillion-Token Datastore With Diverse Domain Composition
MASSIVEDS encompasses both domain- specific data and general web data (Table 2). Domain-specific data comes from special- ized sources, and tends to be smaller but higher in quality. We cull from a mix data- rich domains: books which span variety of genres (Computer) 2023): ; open-access sci- entific papers (Lo et al. 2020 Soldaini &
Table 2: The domain wise data composition of MAS-
SIVEDS. RPJ denotes REDPAJAMA V1 (Computer) 2023 CC denotes Common Crawl, Wiki denotes Wikipedia.
Domain
Datasets
Size (B)
BOOKS
RPJ Books
26.3
STEM
peS2a, RPJ ArXiv
97.7
Lo 2023)0222mp))) Computer 2023); encyclopedic
ENCYCLOPEDIA DPR 2018 Wiki, RPJ 2022 Wiki
31.9
articles (Karpukhin et al. 2020 Computer
FORUM (Q&A)
RPJ StackExchange
20.2
2023); community questions and answers from StackExchange (Computer) 2023);
CODE
RPJ Github
52.8
MATH
OpenWebMath, NaturalProofs
14.1
code from GitHub (Computer] (023)) math-
BIOMEDICAL
PubMed
6.5
ematical webpages (Paster al. 2023 and mathematical language (Welleck et al.
GENERAL WEB RPJ CC (2019-2023) RPJ C4
1191.7
2021); biomedical articles (of Medicine
Total
1441.2
2023). On the other hand, general web data sourced from Common Crawl snapshots five time periods (07/2019, 05/2020, 04/2021, 05/2022,
06/2023) and C4 (Raffel et al., 2020).
3

---

Change Composition
Change Filler
Change
Naive
Deta
(2) subsampling
[I Indesing
(4) Documeré rstrienal
[|] Topak avalsalion
Data
Filtered Raw Data
Subsampled Raw Data
Datastores
Top-k Docs
For Evaluation
Top-k Docs
Change Composition
MassiveDS (Ours)
Change p
(1) Indesing
(3 Document retrioval
Werge
(3) Data
with probaibiit p
Data suhsanping
Top nnluuaion
Data
Datastores
Per-Domain Top-K Docs
Merged
Filtered
Subsampled
Top-k Docs For Evaluation
Top- K Docs
Top-K Docs
pk Docs
Figure 2: Comparison between the MASSIVEDS pipeline and a naive pipeline for studying datastore scaling ($3.2). The green and red arrows indicate repeated operations. In the naive pipeline, these operations more expensive because they require repeating expensive steps such rehuilding the datastore index. In the MASSIVEDS pipeline, these operations are cheaper because they repeat fewer steps and are only run on the retrieved top-K documents instead of the full datastore. Datastore- level operations are represented by purple arrows, while document-level operations, repeated for every query, are represented by black arrows.
3.2 Studying Datastore Scaling with the MASSIVEDS Pipeline
Studying datastore scaling requires constructing datastores of varying sizes and varying compositions from the raw text corpus. This involves the following operations: data filtering, including deduplica- tion, decontamination, and quality filters (Soldaini et al.] 2024); data subsampling, which randomly subselects a p-fraction of the text corpus to achieve the specified size; indexing, which embeds the data using an auxiliary model and builds a searchable index; document retrieval, which uses the index to find the top-k relevant documents for each test query and top-k evaluation, which uses the top-k documents per test query to augment the retrieval-based model. A naive approach is to run these operations in the aforementioned order for each datastore, and build separate datastores for all combinations of subsampled datastore sizes, random seeds, and other experimental variations. However, this naive approach is prohibitively expensive at the trillion-token datastore scale because it repeats expensive operations, as shown in Figure 2 (top).
To make the datastore scaling study feasible, we develop the MASSIVEDS pipeline (Figure bottom). The key idea is to reorder the above operations such that the most expensive ones boes-inndxxing and retrieval- are run only once at the start and then shared across all subsequent datastore variants. Other operations with many variants-such as subsampling, deduplication, and decontamination performed as late possible to minimize repeating subsequent steps. To enable this, we first retrieve a relatively large number (K >> k) of documents for each query and then apply the subsequent operations these sets of retrieved documents, rather than the entire datastore. Altogether, this pipeline reduces compute, I/O, and storage requirements by more than an order of magnitude, enabling us conduct the datastore scaling study on a modest compute budget. In Appendix A.5 we show that the results from the MASSIVEDS pipeline are equivalent to the results from the naive pipeline with high probability, where the randomness comes from the subsampling procedure. We provide more details on the steps in the MASSIVEDS pipeline in Appendix A and detailed configuration in Appendix B.I
Note on the number of tokens in the datastore. In our figures, we plot the datastore size (on the x-axis) by multiplying the total number of tokens in the data pool by the subsampling fraction p. more accurate representation would be the number of tokens in the filtered data pool; however, we do not know the size of the filtered data pool we apply data filtering on the retrieved documents instead of the raw data for computational efficiency. As a result, the number of tokens we plot on our x-axis is proportionally larger, i.c., if a py fraction of the data is filtered out < Ps S 1), then the
Rerankiig is optionally applied, which reranks the retrieved documents based a more effective but usually more expensive heuristic (Nogueira Cho 2019 Sachan et al. (022).
'The operations involved in the naive pipeline are expensive in terms of FLOPs, I/O operations, and storage due repetitive index construction and large-scale retrieval.

---

Llama-2 7B
Liama-2 2 13B
Liama-3 8B
w/ Retrieval
x LM-Only
(a) RedPajama
(b) S2ORC
(c) TriviaQA
(d) NQ
(e) MMLU
(f) MedQA
80
36
60
7.0
35
65
Perplexity
34
Accurady
33
60
50
3.8
6.6
70
32
55
68
31
3.6
6.4
30
50
40
3.4
6.2
29
45
10
10c
10:
10:
10a
10a
10'
10a
10a
10a
10a
10a
10a
10a
10a
Datastore Size (B Tokens)
Datastore Size (B Tokens)
Figure 3: Scaling performance on upstream and downstream tasks with MASSIVEDS, in com- parison with LM-only performance. Left: Perplexity (PPL.) scaling performance REDPAJAMA (multi-domain pretraining corpus) and S2ORC (scientific papers). Right: Downstream scaling performance on TriviaQA (TQA), Natural Questions (NQ), MMLU, and MedQA.
actual number of tokens should also be multiplied by py- Since we plot datastore size on log axis, this corresponds to constant shift and does not change the scaling trends.
4 Datastore Scaling with Retrieval-Based Language Models
4.1 Experimental Setup
Model details. Following prior work (Izacard et al. 2023 Liu et al. 2023 Ram et al. 2023; Shi
et al. 2024 Xu et al. 2023 Asai et al., 2024a),
CONTRIEVER- MSMARCO (Izacard et al. 2022). which represents every document in the datastore as dense vector, as our retriever. We ablate the choice of retriever in Appendix E.I found that
performs on par with, or even better than, more recent larger retrievers. We augment input examples with retrieved documents at the granularity of 256-word chunks. We study datastore scaling perfor- mance with the LLAMA-2, LLAMA-3 (Touvron et al. 2023), PYTHIA (Biderman et al. 2023) ) and OLMo (Groeneveld et al.
2024 model families.
Evaluation. We consider both language modeling and downstream tasks for evaluation. We evaluate language modeling perplexity on data from two domains: (1) general web data sampled from REDPAJAMA (Computer 2023); (2) scientific paper data sampled from S2ORC (Lo et al. 2020 For downstream tasks, our evaluation encompasses general-knowledgee medical, math, and science domains including the following tasks. TriviaQA (TQA; Joshi et al. 2017) comprises trivia questions with answers sourced from Wikipedia and the web. Natural Questions (NQ: Kwiatkowski et al. 2019 Lee et al. 2019) comprises search engine queries and
answers from Wikipedia. Massive Multitask Language Understanding (MMLU; Henddrycss et al. 2021) comprises general-purpose, multi-tast reasoning questions. MedQA (Jin et al." 2020) comprises medical multiple-choice questions sourced from professional medical exams.
Implementation details. For evaluation with retrieval, we concatenate the top A: = 3 documents in reverse order, that higher-ranked documents positioned closer to the query. For downstream tasks, we evaluate models via 5-shot prompting, and we prepend the retrieved documents before the few-shot examples, followed by the question. We do not apply reranking for our main experiments in we study the effect of rerankers in Section 5.2 More details, including decontamination
Section[] 4 measures, are in Appendix B
4.2 Datastore Scaling Results on Language Modeling and Downstream Tasks
Finding 1: Datastore scaling significantly helps language modeling. Figures Bla) and (b) show perplexity curves as a function of datastore size on general web and scientific papers, respectively. Retrieval is strictly beneficial for language modeling: the LM-only baselines (denoted by dashed lines) show the highest perplexity across all models and evaluation datasets. Scaling up the datastore reduces perplexity without signs of saturation, suggesting that further scaling is likely yield additional
5

---

0LMo-1.7-7B
0LMo-1.7-1B
Pythia-128 -28
Pythia-6.99 6.98
Pythia-2.9B
Pythia-18
x
LM-Only
w/ Retrieval
Pareto Optimal w/ Retrieval
-x-- Pareto Optimal LM-Only
TriviaQA
Natural Questions
MMLU
MedQA
35
30
40
Accuracy (%g
Olmo
50
25
45
40
20
24
-20
15
10
35
30
10
30
28
10
10
1pic
1017
1014
1011
1014
1017
25
30
104x
1.04
1042
24
1044
10th
10th
10
TriviaQA
Natural Questions
MMLU
MedQA
70
35
88
30
50
40
Pythia
Accuracy 198)
-6
25
45
90
40
34
35
30
25
10th
10th 1gth Compute (FLOPs)
10
10
10th
10th
1077
1014
10tH Compute (FLOPs)
16th
10th
1014
2011
2014 1077
Compute (FLOPs)
Compute (FLÖPs)
Figure 4: Compute-optimal scaling curves for retrieval-based and LMMonly models of varying datastore sizes, model sizes, and pretraining corpus sizes (detailed setup in $B.4). Darker green or pink indicate larger model sizes for PYTHIA and OLMO respectively; crossmarks in matching colors represent the same model size trained with varying numbers of tokens; each crossmark corresponds to datastore scaling curve of lined dots similar to the ones in Figure The Pareto-optimal points are highlighted in red for retrieval-basee LMs and blue for LM-only. Within a fixed computational budget (represented on the x-axis), retrieval- LMs achieve superior performance, which remains unsaturated along the datastore scaling dimension. Pythia models do not exhibit meaningful scaling curves on MMLU and MedQA that require advanced reasoning abilitics.
improvements. Further, datastore scaling enables small models to outperform their larger LM-only counterparts: when retrieving from MASSIVEDS at the largest scale, LLAMA-2 7B outperforms the LM-only performance of its larger LLAMA-2-13B counterpart. Interestingly, we find 3LAMA-3 8B underperforms LLAMA-2 7B on RedPajama. This aligns with the observations in Xiao et al. (2023) and we discuss potential reasons in Appendix D
Finding 2: Datastore scaling improves performance on several downstream tasks, and the degree of improvement desendepnn.. Figure show the performance on four downstream tasks as function of datastore size. Datastore scaling brings major improvements to knowledge- intensive question answering tasks such as NQ and TQA, where retrieval-based LMs significantly outperform LM-only baselines across all scales, and performance monotonically increases with datastore seale. For instance, a LLAMA-2 7B model that retrieves from fewer than 100B tokens can outperform both its 13B LM-only counterpart and the more capable LM-only LLAMA-3 8B TQA and NQ, indicating the effectiveness of storing knowledge in the datastore.
On MMLU, multi-subject, reasoning-hean benchmark, datastore scaling monotonically improves performance across all model scales. Results are more mixed for MedQA, where only the weaker LLAMA-2 7B benefits more from datastore scaling. For both tasks, datastore scaling does not help the smaller model do better than the larger model. We suspect that this is due to task difficulty and the lack of in-domain data sources: both MMLU and MedQA are more oriented toward reasoning rather than pure factual recall, which poses bigger challenges for both the retriever and the LM. Additionally, MASSIVEDS only contains a small subset of web data and medical papers which may not cover all necessary information to answer these questions. We defer to future work to explore better data sources for these tasks.
4.3 Compute-Optimal Scaling with Retrieval-Based Language Models
Next, we study performance as a function of total training-time compute and show that retrieval-based LMs achieve superior compute-optinal performance compared to LM-only models.

---

Use of intermediate checkpoints. We use the intermediate checkpoints of PYTHIA and OLMo as an approximation of models trained on different numbers of tokens, as detailed in Appendix [..] These intermediate checkpoints share the same learning rate scheduler, with a fixed maximum number of training steps that equals or exceeds the number of steps they have been actually trained for, and therefore the performance of these intermediate checkpoints (with or without retrieval) might be lower than otherwise attainable with the same amount of compute. However, pretraining LMs from scratch for all combinations of model sizes and numbers of pretraining tokens is prohibitively expensive for an academic budget.
FLOPs calculation. We detail the FLOPs computation for datastore construction and pretraining in Appendix [B.4] Datastore construction is much cheaper than pretraining because it only requires one forward pass on all tokens in the datastore with a small retriever (177M parameters in our setup), while pretraining requires a forward pass and a backward pass on pretraining tokens with an LM that can be much larger than the retriever. As we use flat index, no additional operations are required at the indexing step, so the number of FLOPs for datastore construction equals the number of FLOPs for embedding. We note that other types of indexing, e.g., inverted file indexing (IVFADC) (Jégou et al. 2011). may require additional FLOPs during construction and fewer FLOPs at inference. We first focus on training- -iminng--tmmeeeee compute and discuss inference cost at the end of the section.
We show the scaling curves against computational cost of retrieval-based LMs and LM-only perfor- mance on downstream tasks in Figure 4 The Pareto-optimal points for retrieval-base and L.M-nly settings are highlighted in red and blue, respectively.
Finding 3: Retrieval-based LMS outperform LM-only models for the same compute budget. With the same training-time compute, strreval-bassd LMs achieves superior performance than LM- only models, indicating offloading FLOPs from pretraining to datastore construction can result in better performance. Therefore, we conjecture that storing factual knowledge in datastore is more computationally efficient than memorizing factual knowledge in model parameters at training time. We note this claim assumes the LM has enough capacity to reason with the retrieved knowledge. Otherwise, an LM may fail to utilize the retrieved knowledge, which we further discuss in Finding 5.
Finding 4: Even weak language models can benefit significantly from retrieval on knowledge- intensive tasks that measure factual recall. Surprisingly, we find that setrieval-basee PYTHIA models (trained on up to 300B tokens) and OLMo-1.7 models (trained on up to 2T token similar compute-optimal scaling trajectory on TriviaQA and NQ (left columns in Figure 4). despite PYTHIA being trained on less and lower-quality data. Both TriviaQA and NQ evaluate factual recall without complex reasoning. When the right information is provided in context using retrieval, the LM only needs to extract the answer; therefore, these results suggest that the ability to extract factual knowledge for simple factual question answering obtained early in training.
B have a
Finding 5: Retrieval shows benefits for ceasoning-intensin tasks with capable OLMO models, but it does not help when the language model is not sufficiently advanced such as PYTHIA. As shown on the right side of Figure 4 datastore scaling gives marginal benefits on MMLU and MedQA for PYTHIA models where the performance stays around random even at the 12B model size. However, OLMo, which is trained both tasks. We thus conjecture that training on higher-quality data, as OLMO applied in pretraining, could help the model benefit more from retrieved documents for reasoning-heat tasks. Beyond reasoning ability, access to the right data sources for the datastore might be critical. For example, we observe fewer benefits from retrieval on MMLU and MedQA in comparison with TriviaQA and NQ. This may indicate that MMLU and MedQA need more specific data, such as relevant textbooks for MMLU and biomedical literature for MedQA, which are currently not included in MASSIVEDS.
more and better data, consistently benefits from retrieval on
Discussion on inference cost. The compute-optimal scaling study described above focuses only on the cost of training. For inference, prepending retrieved documents in context increases inference cost due to the extended context length and additional computation required for the search. On the flip side, inference cost can be reduced by switching from a larger to a smaller LM, especially since a small LM augmented with a datastore can match or outperform its larger counterparts on
P0LM0-1.7 1B is trained up 3T tokens and OLMo-1.7 1.7 7B trained up 2T tokens.
7

---

Table 3: Downstream and upstream performance comparison between MASSIVEDS for retrieval versus single-domain datastores with LLAMA-2 7B. "SE" is short for StackExchange. The best performance is highlighted bold and the second best is underlined. We show the diverse domain coverage in MASSIVEDS consistently improve the performance across tasks.
DPR
RedPajama ArXiv
Tasks
LM-Only PubMed MATH PeS2o
MASSIVEDS
Wiki
Wiki
Books
SE
Github
TQAT NQ 7
64.1
64.5
65.5
65.6
72.6
72.9
70.5
62.3
64.7
64.2
77.0
26.6
26.7
26.4
26.9
34.6
33.8
28.0
26.4
27.0
26.4
34.6
McdQA 1
36.6
37.8
36.5
38.1
38.5
38.4
39,8
36,9
35.4
36.1
39.4
MMLU f
45.8
46.8
47.5
47.4
48.3
48.1
48.3
45.6
46.2
45.9
RedPajama (PPL)
4.09
4.06
4.08
4.08
4.06
3.99
4.01
3.87
4.01
3,95
3.50
(20RC(PPL))
7.18
7.05
7.10
6.71
7.08
7.11
7.14
6.64
7.08
7.11
6.57
some tasks. We also note that there is emerging work on accelerating retrieval search and designing efficient serving strategies for retrieval-basee LMs, such as CCo et al. (2023). We leave a study of inference-compute-ptiina scaling to future work.
Analysis
5.1 Effects of Data Composition
Finding 6: MASSIVEDS matches or outperforms single-domain datastores. The default setting in prior work to use single-domain datastore that is in-distribution to the downstream task. In practice, however, it is often difficult to determine and curate datastore that is perfectly in- distribution for downstream task, and even we can, it limits the generality of the retrieval-based model to that task.
In Table single-domain datastores. MASSIVEDS sig- nificantly outperforms these in-domain datas- tores on language modeling, as well as TQA and MMLU, while matching performance NQ and MedQA R In Figure 5) we show that the retriever tends to retrieve from the relevant domain even in the presence of out-of-domain data in the datastore: for NQ. retrieves rela- tively more frequently from Wikipedia and web sources, whereas for MedQA, it retrieves more frequently from scientific papers from peS2o (Soldaini Lo. 2023). Thus, the retriever can maintain robustness to oft-of-domain data in the datastore; this aligns with similar findings on KNN-LM (Khandelwal et al.) (2020). another type of retrieval- LM, in Shao et al. al. (2023). Overall, these results show that retrieving from broad datastores like MASSIVEDS can simul- taneously improve performance across multi- ple domains, paving the path towards general- purpose retrieval-basee based models.
compare MASSIVEDS with
MassiveDS
OPR vad
Olen R..0%
Retriee
Curerio Cutw!
Retneve
MedQA
Natural Questions
Retrieved
Retrleved
Other
CRR BNck
Eormon
(51.5%)
Common
Figure 5: Retrievers tend to retrieve from rele- vant domains. We plot the domain composition of MASSIVEDS vs. the top-1 retrieved documents for evaluation examples from MedQA and NQ. The retriever retrieves more frequently from domains that are relevant to the evaluation examples.
5.2 Effects of Reranking
Retrieving the most relevant documents from a large-scale datastore remains a challenging problem. To study how improving the retrieval process impacts datastore scaling trends, we first retrieve 500 documents from CONTRIEVER, rerank them using more computationally expensive model (Ram
NQ answers are annotated based on data from a 2018 snapshot of Wikipedia, so the Wikipedia datastore is exactly the right datastore for this task.

---

Lexical Oracle
No Reranker
LM-only
TriviaQA
Natural Questions
80
45
Accuracy (%)
75
Accuracy
40
70
35
65
30
100
1000
Datastore Size Tokens)
Datastore Size (B Tokens)
Figure 6: Scaling trends TriviaQA and NaturalQuestions using different rerankers (Sec- tion 5.2). "Lexical Oracle" represents the oracle reranker that reorders documents based on lexical overlap with the ground-truth answer. "Cross-encoder" represents a neural reranker which uses cross-encoder model. Both the oracle lexical reranker and the neural reranker boost scaling trends, indicating the potential improvement space by enhancing the retrieval quality.
w/ Aggressive Decon.
w/ Standard Decon.
w/o Decon.
on-only
Language Modeling RedPajama
Natural Questions
4.4
4.2
Accuracy (9%
34
Perplexity
4.0
32
3.8
3.6
30
3.4
10 Datastore Size (B Takens)
10
100
Datastore Token Size (B)
Figure 7: Ablation study on data decontamination. 'Aggressive Decon.' removes document as long as it has an 8-gram (i.e., 1.5% of the answer length) continuous overlap with the answer. 'Standard Decon.' default setup-removes document when it either has 32-gram (i.e., 6.2% of the answer length) continuous overlap or an 80%+ Jaccard similarity with the answer. We find decontamination impacts the language modeling performance lot but not the downstream task.
al. 2023 Sachan al. 2022 and take the final top-3 reranked documents for evaluation. Specifi- cally, we use a cross-encoder reranker, which encodes a concatenation of a query and document and returns similarity score (Nogueira & Cho 2019). We choose MINI-L.M cross-encode that is trained for passage ranking, following Izacard et al. (0022). Additionally, we use a lexical oracle reranker, which uses the gold answer, as an upper bound on the potential benefit realizable by better reranker or retriever for knowledge-intensive question-an wwwiiiinnnnserriiie -nswering.. The oracle reranker scores each document based whether the gold answer is included in the document and not, the fraction of unigram overlap between the document and the answer.
V2, BERT-based
Figure 6 reports scaling trends TQA and NQ of RIC-LM with Llama2-7B. While the cross- encoder- based reranker improves performance on TQA and NQ, a notable gap persists between the oracle reranker and the cross-encoder-basee reranker. These suggest that improving either retrieval reranking can further boost the scaling performance retrieval datastores. Improving the retriever for more reasoning-hee heavy tasks such MMLU and MedQA remains open problem (BehnamGhader et al. 2022) that we leave to future work.
5.3
Effects of Datastore Filtering
Data decontamination. Data decontamination is a crucial concern when evaluating LMs, especially in retrieval-based LMs that can retrieve the test data verbatim during inference (Borgeaud et al. 2002). By default (Section 4. we perform decontamination by filtering documents with 80+% 13-gram
html

---

Jaccard similarity for downstream tasks and 32-gram longest sequence overlap for perplexity. which we call standard decontamination. Prior work such as RETRO (Borgeaud et al. 2022) only used 80+% 13 13-gram Jaccard similarity for decontamination. However, we find the additional 32-gram longest sequence overlap decontamination is critical for removing near-identical documents.
To study the impact of varying levels of data decontamination, we compare standard decontamination with two additional methods: (1) no decontamination and (2) aggressive decontamination, which uses 8-gram longest sequence overlap for all tasks. This is a strict filter, as 8-gram overlap occurs frequently even when documents are not nearly identical.
Figure Z reports the performance of the LLAMA-2 7B model on language modeling and the Natural Questions dataset using different decontamination methods. The scaling trend shows significantly better language modeling performance without decontamination, which worsens with more aggressive decontamination methods. This suggests that the benefits in language modeling primarily arise from lexical overlap. However, retrieval continues to benefit language modeling performance even after aggressive decontamination-wweerr no more than continuous words verlap-indicating that semantically similar retrieved documents with minimal lexical overlap can still enhance language modeling. Decontamination does not significantly affect NQ performance, likely because there is less contamination of NQ in the datastore. Interestingly, decontamination decreases performance with smaller datastores, but improves final performance larger scales.
E.2 we study the impact of data quality filtering on MAS- SIVEDS, where we consider global data deduplication and combination of 3 filters adapted from DOLMA (Soldaini et al. 2024): whitespace filter; language filter, and alphanumeric filter. We find deduplication is helpful to minimizing saturation as the datastore scales on NQ: intuitively, subsam- pling with higher p increases the chance of seeing more duplicates. In addition, we observed that DOLMA quality filters have relatively limited effect. We hypothesize this is because the data sources we used in MASSIVEDS, such as RedPajama, have already gone through similar quality filtering processes and may not benefit much from applying additional filtering.
Data quality filtering. In Appendix
Limitations and Discussion
We conclude by discussing limitations and future directions. First, while our pipeline allows us to study datastore scaling efficiently, our experiments are still limited by our available compute. In particular, our compute-optimal scaling studies are limited model families like OLMo and Pythia that release intermediate model checkpoints, since full pretraining runs exceed our budget constraints. Similarly, we conduct the full scaling study with a single retriever, as changing the retriever necessitates re-indexing the entire datastore. It remains unclear how changes in the size and architecture of the retriever affect datastore scaling trends.
Second, although MASSIVEDS is large in size, it might still lack nigh-quality data for improving performance more complex, reasoning- heavy tasks such as MMLU and MedQA. Future work could study the effect of extending MASSIVEDS to more varied and higher quality data sources.
Lastly, our downstream evaluations are mostly on predefined choices or short form generations. We defer the evaluation on more tasks such as long-form generation and mathematical reasoning to future work.
tasks whose outputs are either
Despite these limitations, our research shows that increasing the scale of data available inference time can improve model performance, at lower training cost, on language modeling and variety of downstream tasks. We expect that future work on improving retrieval- datastores will lead to even larger improvements: for example, our analysis suggests that further improving the retrieval process, either through better retrievers or rerankers, could have significant impact. We also expect scaling up the datastore to be one way to achieve extremely long-context modeling (e-g.. retrieving information from a trillion-token context) and test-time scaling (e.g.. spending more test-time compute retrieve helpful information from large datastore).
models with aarge-scale
10

---

Acknowledgements
We thank Hannaneh Hajishirzi, Scott Yih, Ian Magnusson, Matt Jordan, and Rui Xin for insightful discussions. We thank Stella Li for proofreading. PWK is supported by the Singapore National Research Foundation and the National AI Group in the Singapore Ministry of Digital Development and Innovation under the AI Visiting Professorship Programme (award number AIVP-2024-001).
References
Amro Abbas, Kushal Tirumala, Dániel Simig, Surya Ganguli, and Ari S Morcos. Semdedup: Data- efficient learning at web-scale through semantic deduplication. arXiv preprint arXiv:2303.09540 2023.
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 Conference on Learning Reprexentations, 2024a.
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. arXiv, 2024b.
Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. In International Conference on Learning Representations, 2019. URL httpp://opennrevieew. net/ ByxZX20qFQ
Parishad BehnamGhader, Santiago Miret, and Siva Reddy. Can retriever-augmented language models reason? the blame game between the retriever and the language model. arXiv preprint arXiv:2212.09146. 2022.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pp. 2397-2430. PMLR, 2023.
Piotr Bojanowski, Edouard Grave, Armand Joulin, and Tomas Mikolov. Enriching word vectors with subword information. Transactions of the Association for Computational Linguistics, 5:135-146, 2017. ISSN 2307-387X.
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Mil- lican, George Bm Van Den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, Diego De Las Casas, Aurelia Guy, Jacob Menick, Roman Ring, Tom Hennigan, Saffron Huang, Loren Maggiore, Chris Jones, Albin Cassirer, Andy Brock, Michela Paganini, Ge- offrey Irving, Oriol Vinyals, Simon Osindero, Karen Simonyan, Jack Rac, Erich Elsen, and Laurent Sifre. Improving language models by retrieving from trillions of tokens. In Kama- lika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. https://proceeddngg mlr ppress/v12//borrrgeeee
PMLR, 17-23 Jul 2022. URL
html
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Miererrt-Vsss. Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Ben- jamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot leamner. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems. volume 33, pp. 1877-1901. Curran Associates, Inc., 2020. URL httputtts//pppceeeiingg http://owceeiiinnn.nie.o.ppp.pp. www.eess......p...e.eeeeeeeeeeee cc/paper
Qingqing Cao, Sewon Min, Yizhong Wang, and Hannaneh Hajishirzi. Btr: Binary token repre- sentations for efficient retrieval augmented language models. arXiv preprint 10.00323,..01329, 2023.
11

---

Together Computer. Redpajama: An open source recipe to reproduce llama training dataset, 2023.
URL https://gittubb. coom/ooogtterrcooptteee.eee RedPaj ama amaa
Samir Yitzhak Gadre, Georgios Smymis, Vaishaal Shankar, Suchin Gururangan, Mitchell Wortsman, Rulin Shao, Jean Mercat, Alex Fang, Jeffrey Li, Sedrick Keh, et al. Language models scale reliably with over-training and on downstream tasks. arXiv preprint (xXix:2403.08544 2024.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang. Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy, The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00000 2020.
Tianyu Gao, Howard Yen, Jiatong Yu, and Danqi Chen. Enabling large language models to generate text with citations. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 6465-6488, Singapore, December 2023. Association for Computational Linguistics. doi: 10. (0.11855//11/ 18653//12023..mmmppmmm 398. URL
2223..2222222222 emnlp-maile mwim. 398
Dirk Groeneveld, Iz Beltagy, Pete Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya Harsh Jha, Hamish Ivison, Ian Magnusson, Yizhong Wang, al. Olmo: Accelerat- ing the science of language models. arXiv preprint arXiv:2402.00838, 2024.
Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Ming-We Wei Chang. Realm: retrieval- augmented language model pre-training. In Proceedings of the 37th International Conference Machine Learning, ICML 20. JMLR.org, 2020.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300, 2020.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. Proceedings of the International Conference on Learning Representations (ICLR), 2021.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, al.
Training compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022.
Gautier Izacard and Edouard Grave. Leveraging passage retrieval with generative models for open domain question answering, 2020. URL htttt::////aaattt www./ba//////aaaeeo ab0/ .0128
Gautier Izacard, Mathilde Caron, Lucas Hosseini, Sebastian Riedel, Piotr Bojanowski, Armand Joulin, and Edouard Grave Unsupervised dense information retrieval with contrastive learn-
ing. Transactions on Machine Learning Research, 2022. ISSN 2835-8856. URL https //openreview.net// wtt.://ooorrmmddddddd
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, 2023. URL
httt:://jml:
www.o////../..........o/. http:/apers/v22///2 0037 html
Hervé Jégou, Romain Tavenard, Matthijs Douze, and Laurent Amsaleg. Searching in one billion vectors: re-rank with source coding. In 201/ IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 861-864. IEEE, 2011.
Di Jin, Eileen Pan, Nassim Oufattole, Wei-Hung Weng, Hanyi Fang, and Peter Szolovits. What disease does this patient have? a large-scale open domain question answering dataset from medical exams. arXiv preprint arXiv:2009. 1308I, 2020.
Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. arXiv preprint 1705../005.03.3 2017.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv arXiv:2001.00 08361, 2020.
12

---

Vladimir Karpukhin, Barlas Oğuz. Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen tau Yih. Dense passage retrieval for open-domain question answering, 2020.
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer. and Mike Lewis. Generalization through memorization: Nearest neighbor language models, 2020.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: benchmark for question answering research. Transactions of the Association for Computational 2019.
Linguistics,
Katherine Lee, Daphne Ippolito, Andrew Nystrom, Chiyuan Zhang, Douglas Eck, Chris Callison- Burch, and Nicholas Carlini. Deduplicating training data makes language models better. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.). Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Pupers), pp. 8424-8445, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 10. 1/2022.acl-ongggg 577. URL //
org/ wtgg/2222acll1oortn acl-long. 577
Kenton Lee, Ming-Wei Wei Chang, and Kristina Toutanova. Latent retrieval for weakly supervised open domain question answering. In Anna Korhonen, David Traum, and Lluís Mårquez (eds.), Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/ v1/P19-1612. URL [ttLlttts///ccaattiioee
org/P19-1612
Sheng-Chies Chigg-Chihhuaa Lin, Akari Asai, Minghan Li, Barlas Oguz, Jimmy Lin, Yashar Mehdad, Wen-tau Yih, and Xilun Chen. How to train your dragon: Diverse augmentation towards generalizable dense retrieval. arXiv preprint arXiv:2302.07452. 2023.
Xi Victoria Lin, Xilun Chen, Mingda Chen, Weijia Shi, Maria Lomeli, Rich James, Pedro Ro- driguez, Jacob Kahn, Gergely Szilvasy, Mike Lewis, Luke Zettlemoyer, and Scott Yih. RA-DIT: Retrieval-augmeented dual instruction tuning. In The Twelfth International Conference Learning Representations, 2024. URL htttttt:/s/pprreeeeeeee
Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts, 2023, SXXiv:2307.0317.
Kyle Lo, Lucy Lu Wang, Mark Neumann, Rodney Kinney, and Daniel Weld. S2ORC: The semantic scholar open research corpus. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 4983, Online, July 2020. Association for Computational Lin- guistics. doi: 10. 10.18653/v1/2020..o main. 447. URL https: //wp::://wwww... aclweb. org/ anthology/ 20200 acl - main.ainn... 447)
Ian Magnusson, Akshita Bhagia, Valentin Hofmann, Luca Soldaini, Ananya Harsh Jha, Oyvind Tafjord, Dustin Schwenk, Evan Pete Walsh, Yanai Elazar, Kyle Lo, et al. Paloma: A benchmark for evaluating language model fit. arXiv preprint arXiv:2312. 10523, 2023.
Alex Mallen, Akari Asai, Victor Zhong, Rajarshi Das, Daniel Khashabi, and Hannaneh Hajishirzi. When not to trust language models: Investigating effectiveness of parametric and non-parametric memories. In Anna Rogers, Jordan Boyd-Grabel and Naoaki Okazaki (eds.), Proceedings of the sl Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 9802-9822, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: URL httpss::::/// httppp::/aaaaaathooooooo org/ 2023 acl-long. 546
Sewon Min, Suchin Gururangan, Eric Wallace, Hannaneh Hajishirzi, Noah A. Smith, and Luke Zettlemoyer. Silo language models: Isolating legal risk in nonparametric datastore, 2023a.
Sewon Min, Weijia Shi, Mike Lewis, Xilun Chen, Wen-tau Yih, Hannaneh Hajishirzi, and Luke Zettlemoyer. Nonparametric masked language modeling. In Anna Rogers, Jordan Goyd-Graber and Naoaki Okazaki (eds.), Findings of the Association for Computational Linguistics: ACL 2023, pp. 2097-2118, Toronto, Canada, July 2023b. Association for Computational Linguistics. doi: 10. 1166386111/222 findings-acl.1333 J.R..https://aclatthee
www.orgg/3023 findings-acl
132
13

---

Niklas Muennighoff, Alexander M Rush, Boaz Barak, Teven Le Scao, Aleksandra Piktus, Nouamane Tazi, Sampo Pyysalo, Thomas Wolf, and Colin Raffel. Scaling data-constrained language models. arXiv preprint arXiv:2305. /6264, 2023.
Jianmo Ni, Chen Qu, Jing Lu, Zhuyun Dai, Gustavo Heráández Abrego, Ma, Vincent Y Zhao,
Yi Luan, Keith B Hall, Ming- Chang, et al. Large dual encoders are generalizable retrievers arXiv preprint arXiv:21 12.07899, 2021.
Rodrigo Nogueira and Kyunghyun Cho. Passage re-ranking with bert.
arXiv preprint
1rX1::1901.04085, 2019.
National Library of Medicine. Pubmed baseline 2023 repository, 2023. URL https://lhnch: nlm. mih. htt.g//i////oiorratiooooo /MBRRR html
Keiran Paster, Marco Dos Santos, Zhangir Azerbayev, and Jimmy Ba. Openwebmath: An open dataset of high-quality mathematical web text, 2023.
Guilherme Penedo, Quentin Malartic, Daniel Hesslow, Ruxandra Cojocaru, Alessandro Cappelli, Hamza Alobeidli, Baptiste Pannier, Ebtesam Almazrouei, and Julien Launay. The refinedweb dataset for falcon IIm: outperforming curated corpora with web data, and web data only. arXiv preprint ar arXiv:2306.0/ 11,, 2023.
Fabio Petroni, Aleksandra Piktus, Angela Fan, Patrick Lewis, Majid Yazdani, Nicola De Cao, James Thorne, Yacine Jernite, Vladimir Karpukhin, Jean Maillard, Vassilis Plachouras, Tim Rocktäschel, and Sebastian Riedel. KILT: a benchmark for knowledge intensive language tasks. In Kristina Toutanova, Anna Rumshisky, Luke Zettlemoyer, Dilek Hakkani-Tur, Iz Beltagy, Steven Bethard, Ryan Cotterell, Tanmoy Chakraborty, and Yichao Zhou (eds.). Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 2523-2544, Online, June 2021. Association for Computational Linguistics. doi: 10.18653/v 1ttp://1/22111a2 www.aaaaaa.aaaa... main.200. URL https///ccaaahhlogy.ooom 2rg/ 2021
naacl-mate 200
Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Dmytro Okhonko, Samuel Broscheit, Gautier Izacard, Patrick Lewis, Barlas Oğuz, Edouard Grave, Wen tau Yih, and Sebastian Riedel. The web is your oyster- knowledge-intene intensive nlp against a very large web corpus, 2022.
Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, Eliza Rutherford, Tom Hennigan, Jacob Menick, Albin Cassirer, Richard Powell, George van den Driessche, Lisa Anne Hendricks, Maribeth Rauh, Po-Sen Huang. Amelia Glaese, Johannes Welbl, Sumanth Dathathri, Saffron Huang. Jonathan Uesato, John Mellor, Irina Higgins, Antonia Creswell, Nat McAleese, Amy Wu, Erich Elsen. Siddhant Jayakumar, Elena Buchatskaya. David Budden, Esme Sutherland, Karen Simonyan, Michela Paganini, Laurent Sifre, Lena Martens, Xiang Lorraine Li, Adhiguna Kuncoro, Aida Nematzadeh, Elena Gribovskaya, Domenic Donato, Angeliki Lazaridou, Arthur Mensch, lean-Baptiste Lespiau, Maria Tsimpoukelli, Nikolai Grigorev, Doug Fritz, Thibault Sottiaux, Mantas Pajarskas, Toby Pohlen, Zhitao Gong. Daniel Toyama, Cyprien de Masson d' Autume, Yujia Li. Tayfun Terzi, Vladimir Mikulik, Igor Babuschkin, Aidan Clark, Diego de Las Casas, Aurelia Guy, Chris Jones, James Bradbury, Matthew Johnson, Blake Hechtman, Laura Weidinger, lason Gabriel, William Isaac, Ed Lockhart, Simon Osindero, Laura Rimell, Chris Dyer, Oriol Vinyals, Kareem Ayoub, Jeff Stanway, Lorrayne Bennett, Demis Hassabis, Koray Kavukcuoglu, and Geoffrey Irving. Scaling language models: Methods, analysis & insights from training gopher, 2022.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1-67, 2020. URL
http://jmlr prg/papprs/vv//220 074 000.-0744htt html
Ori Ram, Yoav Levine, Itay Dalmedigos, Dor Muhlgay, Amnon Shashua, Kevin dcyton-Brown and Yoav Shoham. In-context retrreeallaugmmented language models. Transactions of the Association for Computational Linguistics, 1 11:1316-1331, 2023. doi: 10.1 1 122ttacl a 00605. URL httpss
org/2023. //aclantholo .rg/ 2023. tac1-1.75
14

---

Devendra Sachan, Mike Lewis, Mandar Joshi, Armen Aghajanyan, Wen-tau Yih, Joelle Pineau, and Luke Zettlemoyer. Improving passage retrieval with zero- question generation. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang (eds.), Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 3781-3797, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 10. 10.18653/v120022 emnlp-main.244 249.
httpe//aannnoooooogg/222222 emnlp-main. 249
Rulin Shao, Sewon Min, Luke Zettlemoyer, and Pang Wei Koh. Retrieval- language models using multi-domain datastore. In NeurlPS 2023 Workshop Distribution Shifis: New Frontiers with Foundation Models, 2023. httptttpss///ppeeeeeeeeeeee http://t...nntttttrrrooodccceee
Weijia Shi, Sewon Min, Michihiro Yasunaga, Minjoon Seo, Rich James, Mike Lewis, Luke Zettle- moyer, and Wen tau Yih. Replug: Retrieval-augmented black box language models, 2023.
Weijia Shi, Sewon Min, Maria Lomeli, Chunting Zhou, Margaret Li, Gergely Szilvasy, Rich James, Xi Victoria Lin, Noah A. Smith, Luke Zettlemoyer, Scott Yih, and Mike Lewis. In-context pretraining: Language modeling beyond document boundaries, 2024.
Luca Soldaini and Kyle Lo. peS2o (Pretraining Efficiently on S2ORC) Dataset. Technical report, Allen Institute for AI, 2023. ODC-By, https::/githuub. com/ al allenai/pes20
Luca Soldaini, Rodney Kinney, Akshita Bhagia, Dustin Schwenk, David Atkinson, Russell Authur, Ben Bogin, Khyathi Chandu, Jennifer Dumas, Yanai Elazar, Valentin Hofmann, Ananya Harsh Jha, Sachin Kumar, Li Lucy. Xinxi Lyu, Nathan Lambert, Ian Magnusson, Jacob Morrison, Niklas Muennighoff, Aakanksha Naik, Crystal Nam, Matthew E. Peters, Abhilasha Ravichander, Kylc Richardson, Zejiang Shen, Emma Strubell, Nishant Subramani, Oyvind Tafjord, Pete Walsh, Luke Zettlemoyer, Noah A. Smith, Hannaneh Hajishirzi, Iz Beltagy. Dirk Groeneveld, Jesse Dodge, and Kyle Lo. Dolma: an open corpus of three trillion tokens for language model pretraining research, 2024.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov. Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.092288 2023.
Boxin Wang. Wei Ping, Peng Xu, Lawrence McAfee, Zihan Liu, Mohammad Shoeybi, Yi Dong, Oleksii Kuchaiev, Bo Li, Chaowei Xiao, Anima Anandkumar, and Bryan Catanzaro. Shall we pretrain autoregressive language models with retrieval? a comprehensive study. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empir- ical Methods in Natural Language Processing, pp. 7763-7786, Singapore, December 2023. Association for Computational Linguistics. doi: 10.188533v1/20222 wwtp::2....mmmmmmmmmm wwn.p-maii.n48222222 URL https://accaanhoolgyy www.oologo.oog//2220 wwt.:r//02....mmmmmaaaa.......aa 482
Boxin Wang, Wei Ping, Lawrence McAfee, Peng Xu, Bo Li, Mohammad Shoeybi, and Bryan Catanzaro. Instructretro: Instruction tuning post retrieval -augeetedd pretraining. 2024. URL
httiss://penrreviev. httpeeerrriieeeeeeeeeeeeee ooromiimmmidd dd
Sean Welleck, Jiacheng Liu, Ronan Le Bras, Hannaneh Hajishirzi, Yejin Choi, and Kyunghyun Cho. Naturalproofs: Mathematical theorem proving natural language. Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Truck (Round 1), 2021. URL
https://openreviii wtcpiiiwwweeeeeeeeee www.:nnoooomddddieeeent for.ormiiddd..aaaaaa3331111
Guillaume Wenzek, Marie-Anne Lachaux, Alexis Conneau, Vishrav Chaudhary, Francisco Guzmán, Anmand Joulin, and Edouard Grave. CCNet: Extracting high quality monolingual datasets from web crawl data. In Nicoletta Calzolari, Frédéric Béchet, Philippe Blache, Khalid Choukri, Christopher Cieri, Thierry Declerck, Sara Goggi, Hitoshi Isahara, Bente Maegaard, Joseph Mariani, Hélène Mazo, Asuncion Moreno, Jan Odijk, and Stelios Piperidis (eds.), Proceedings of the Twelfth Language Resources and Evaluation Conference, pp. 4003-4012 2, Marseille, France, May 2020. European Language Resources Association. ISBN 979-10-95546-344 URL https //aclanthology. org/
wrg./020....... 2020 1rec- 1 194
Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference Machine Learning, pp. 38087-38099. PMLR, 2023.
15

---

Fangyuan Xu, Weijia Shi, and Eunsol Choi. Recomp: Improving retrieval-augmented Ims with compression and selective augmentation, 2023.
16

---

Appendix
Datastore Construction Pipeline
18
A.1
Distributed Indexing
18
A.2
Distributed Document Retrieval
18
A.3
Domain Merging
19
A.4
Data Filtering and Reranking
19
A.5
Data Subsampling
21
6
Evaluation
23
Implementation Details
24
B.1
Pipeline
24
B.2
Language Modeling Evaluation Setup
24
B.3
Downstream Evaluation Setup
24
B.4
Compute-Optimal Scaling Setup
25
Complete Datastore Scaling Results
25
D Discussion on the Inferior LLAMA-3 PPL Performance on RedPajama Data
29
Additional Analysis
30
Ablation on the Retriever
30
Effect of Data Quality Filtering
30
17

---

A Datastore Construction Pipeline
The MASSIVEDS pipeline entails running the following operations in order (Figure
1. Distributed indexing. We split the documents from each domain into separate shards and construct an index for each shard. Distributing the index allows us to parallelize indexing and retrieval processes, and to study different combinations of data sources more easily.
2. Distributed document retrieval. For each test query, search for the top-K documents over each shard index in parallel. The searched results are first merged within each domain. We cache the per-domain searched results for data composition analysis.
3. Domain merging. For target combination of domains in the datastore, we merge the top-K documents from the target domains for that query. This leads to merged pool DK retrieved documents, where D is the number of target domains. From this pool, we re-select the top-K documents.
Data filtering and reranking. We then apply the data filtering and reranking steps, as described above, to only the lopK documents (for each test query). This allows us to experiment with different data filters and rerankers without having to rerun indexing or retrieval; moreover, we only need to filter/rerank the retrieved results instead of the whole datastore, and we do not need to repeat these processes for different subsampling ratios.
Data subsampling. We subsample the filtered and reranked documents for each test query. Specifically, for a particular subsampling ratio p. we select each document for inclusion i.i.d. with probability p. To try different subsamples (based on different random seeds), we only need to restart the pipeline from this step, and we can reuse computation from all previous steps.
6. Evaluation. We use the top-k documents for each test query and prepend these documents in front of the query and few-shot examples for evaluation.
In practice, we position the reranking step after data subsampling to reduce the number of documents that require reranking. The commutativity of reranking and subsampling is demonstrated Ap- pendix A.5 Furthermore, we collect the top-K' documents (where K' K)) from the deduplicated and decontaminated top-K K set for reranking. We set K" = & when reranking is not applied. Next, we describe each step of our the equivalence between our MASSIVEDSS pipeline and the naive pipeline. Below is a table of notations for easy reference.
datastore construction pipeline in detail and demonstrate
Symbol Deseription
Symbol
Description
Parameters of the reader model Number of documents for evaluation Number of documents for reranking
Parameters of the retriever model
K
Number of documents before deduplication Retrieval datastore
K"
D
A.1
Distributed Indexing
Our pipeline starts with raw data that is split into fixed-size documents. For each document, run one forward pass over it with retriever model & and save the final-layer representation as an embedding for that document. We store the embeddings of all documents the datastore for similarity-basee search. In practice, we split the documents into separate shards and embed each shard in parallel. As we use the uncompressed embedding for retrieval searches, our indexing does not require additional operations over the saved embeddings. We treat each shard of embeddings as one sharded index. This indexing step executed only once over all data, while the subsequent steps are performed for each query at inference. For simplicity, we describe the subsequent steps for single query in the sections below.
A.2 Distributed Document Retrieval
At inference, we first retrieve the top-K documents from each sharded index in parallel to accelerate the search process. Specifically, we convert the query into an embedding and compute the inner- product similarity scores between this query embedding and every document embedding. We then rank the documents from each shard based on these similarity scores and collect all top- documents for subsequent merging.
18

---

A.3 Domain Merging
Assuming we have m sharded indices, we merge the m sets of top-K retrieved documents from all indices based on their similarity scores to obtain the final top-K documents. This process is formally defined below as m-sharded distributed element-wise top- K retrieval. We demonstrate that is equivalent to directly retrieving the top-K documents from a single index built on all data.
Formally, we introduce two definitions: element-wise top-K retrieval and m-shard element-wise lopKK retrieval, the latter of which is what our pipeline uses.
Definition A.1 (Element-wise top-K retrieval). Assume a datastore of N documents: 'D {di..... ,dx.. Given a query element-wise top-K retrieval uses a retriever to compute the similarity score 8; = g0lg, di) of the query and each document d, independently, and then returns the documents with the top- K highest retrieval scores.
Definition A.2 (m-shard distributed element top-K retrieval). Consider sharding strategy , }. such that each shard contains a disjoint subset
that splits the datastore into m shards {D1, ... Dm of documents D (i.e., Dj U D2 U... D,n D: D; D = 0, when + j). In m--haad distributed element-wise retrieval, we fetch the top-K documents from each shard in parallel (if there are fewer than K documents in shard, then all documents in the shard are returned). Then, we merge these mK documents and return the top-K K highest-scoring documents from the merged pool.
m-shard distributed dement-wise top-K retrieval is equivalent to retrieving from single unsharded index (m 1).
Lemma A.1. Changing m, the number of shards used in distributed element- top-K retrieval, does not impact the final retrieved results.
Proof. Let the top- K documents obtained by the element-wise top-K K retrieval and m-shaadd dis- tributed element-wise top-K K retrieval be Dx and D'x respectively. Since document that is ranked the top-k K across all documents must be ranked in the top-K within any individual shard, we have Dr C 'Dxg Because (D!.... Dm) } is disjoint union of D. for any document d in D'x there are no more than K documents in the m shards, i.c., all documents, that have higher scores than d. Then we have d Dr. Therefore, D'x Dr. Finally, we have Dx D'x for any choice of m E N
In our experiments, we use an indexing method 0 that computes the similarity score between the query and every document, independent of other documents, i.e., element- wise retrieval. Therefore, our method for distributed search equivalent to building and retrieving from a single non-distributed index.
A.4 Data Filtering and Reranking
Our pipeline includes two basic data processing steps data deduplication and decontamination. These steps are applied post-hoc on the retrieved documents. Reranking, which is optional, used to enhance the quality of retrieval. Detailed descriptions of deduplication, decontamination, and optional reranking are provided in this section.
A.4.1 Post Hoc Datastore Deduplication
Although REDPAJAMA (Computer 2023). ) the pretraining corpora used to build MASSIVEDS. has already been deduplicated, we still noticed many duplicates in the retrieved results, particularly from the web domain. This is because Computer (2023) only performs local deduplication within each data shard; globally, many duplicates between shards remain.
"We INDEXFLATIP implemented in FAISS: http:://ggttu.cco//ccotttttrrrrr wiki/Faiss- -inaassss-nneeesssssssee
Additional data quality filters, such as those from Dolma Soldaini et al. 2024 are discussed in our analysis These filters applied the same stage deduplication.
section 5.3
19

---

There are two widely adopted approaches for deduplication over large-scale dataset: the first uses Bloom filter and the second uses MinHas/ In short, the Bloom filter approach is typically used to remove exact duplicates and can also remove near duplicates; the MinHash approach is used to remove near duplicates detected based on n-gram Jaccard similarity. Running global deduplication over 1 1.4 trillion tokens using a Bloom filter requires 1.5TB RAM memory and 933 CPU hours The MinHash approach requires even more memory and CPU hours than the Bloom filter approach. Therefore, deduplicating over the entire raw text of the datastore is computationally expensive, particularly when it needs to be repeated with every experimental modification.
To get around this, we add post-hoe MinHash deduplication step which is applied to a large pool of top-K retrieved documents. Following Computer (2023), a document pair is duplicate if its 13-gram Jaccard similarity score is at least 80%. Note that deduplication on the retrieved results obtained from the entire corpus is an affordable alternative to running global deduplication. The only risk is that we may not have enough documents for reranking (requiring K" documents) or evaluation (requiring & documents) after deduplication. Therefore, we choose a large K (K > K" and K > k), i.e. K 1000, to mitigate this risk.
The original implementation of MinHash deduplication skips the chunks with less than 13 grams. Qualitatively, after deduplicating with 13-gram Jaccard similaritily, we still find many short duplicates or nonsensical phrases under 13 words in the deduplicated data pool. Thus, we remove all documents with less than 13 words as well at the deduplication step.
A.4.2 Post Hoc Datastore Decontamination
Expanding the datastore to the Internet-scale incurs the risk of data contamination. Similar to global deduplication, running pre-decontamination the entire datastore against every possible evaluation set is inconvenient and computationally expensive. Therefore, we implement a post-hoc decontamination method to remove contamination from the retrieved documents instead of the large datastore. For downstream tasks, we compare the 13-gram Jaccard similarity between each question and the retrieved document, and remove the documents with a high similarity score that is at least 80%. For language modeling (Figure (). we adopt a stricter approach: we calculate both the 13-gram Jaccard similarity and the longest sub-sentence between the document and the answer, marking the document as contaminated if it either has at least 80% 13-gram similarity or contains a continuous 13 13-gram overlap with the answer. We split sentences into grams based on whitespace.
We show in Lemma A.2 that our method for post-deduplication and post-decontamination is equivalent to running deduplication and decontamination on the raw data prior to retrieval.
Lemma A.2. Running post hoc exact deduplication and decontamination over the top-K retrieved documents before laking the top-K' (where K' S K) documents is equivalent to retrieving the top-K' K' documents from deduplicated and decontaminated datastore.
Proof. Given a datastore D, let the deduplicated datastore DD' Denote the top-K documents from D as Dx and the top-K' documents D' as D' Then, denote the top- documents retrieved from Dæ as DR. Since both D' K' and DK' contain the top-K' ranked documents from all those retained after deduplication or decontamination, and since the removal of data is deterministie, we have D' K' DK
For approximate deduplication, there might be edge cases where the two documents are near duplicates but only one of those documents is retrieved in the top-K results. However, the property that none of the final top-k retrieved results are (near-) pdppicaep of each other will still hold.
A.4.3 Reranking
Reranking is optionally applied to enhance retrieval quality in our pipeline. Given the top-K retrieved documents, we apply stronger retrieval model to embed each document and the query,
A typical implementation is hstttss://gittuu.com///
9 A typical implementation is htttps://ekkhuu wtt.//oo//aaaaaa /atasketatt
10 requirement of Bloom Filter deduplication is estimated using the implementation by A12: https
/ //gtthb..cc com/ wttp////alnaaa/tt//eeeeeaa http://tfff/treeeeai//eee hardware setup.
20

---

(soed neriie documeens)
Gold lest
perpleitg)
Compare the similartty between each document and test data
Shrane liadinng aowen when bMIp hestwayyyy diss and inboste hase 7007 towarsee tho and for nearty lan years oeatee new ligit bownd, or theme someeh mg atoultt conowmed stout?
beflery has been Insed...
This Warkinil collet Sleep, Appll relabookk will kowp the EDAM conteess allww (dleep Pnc jaspon)
red the IAM conterte from thes, and has you gat the progrees ber you aew now my Apple thos about is here...
Check similarity is below threshold
Doc.
Apple nolebooiks will keep the RAM contenis ahte (sleap PC jeggnn) needs
road the HAM contanta from dial, and thas you gat the progress but you ure sou ng. Apple dos about it hare....
Mullinnkiing
backgmund progame
meertey cribswey?
X
Doc.
Keep
Rumove
Why doss older Macstook the ahow bar arter the battryy complesey expletes and restart
3 Filtor the documents
seeing
Apple teptops batheyy system Upon enging restanting,
RUM, rosu ling the prog usaa solice...
o How disoble or get the sta hup sound my Mac makes? Everyine tur on my Macbool Pro it maliss start clsble
annowing there volume iy Lof (just don't sound to play
A: Open Temminal.app and typo: sude is. Give the posswort when askodd for..
Figure 8: The post-hoe decontamination process for perplexity evaluation.
and recalculate the similarity scores to reorder the documents. Since reranking is an element-wiss operation, which we demonstrate be commutable with subsampling in the next section, we apply reranking after subsampling in practice.
A.5 Data Subsampling
To study how model behavior changes with datastore scale, we must subsample datastores of different sizes (i.e., different sampling ratios) and with different random seeds. We first describe a naive way of subsampling from a large datastore (Algorithm I and show that it is computationally expensive to scale the datastore in this manner. We then introduce our efficient datastore subsampling strategy (Algorithm(2) to make datastore scaling more affordable. We also provide theoretical justification to show that our method is equivalent with high probability to naive datastore subsampling, and additionally provide a computational comparison of the two approaches. Finally, we discuss the commutativity of subsampling and other operations, demonstrating that users can flexibly reorder some of these operations in practice.
First, we define a few functions used in our algorithm.
Function SUBSAMPLE(D, s). Given data corpus D with N elements, sample each element i.i.d. following a Bernoullip)) distribution with probability and a random seed set such that for any given (p. s), the same æ will always either be sampled or not. The total number of sampled data follows B.,.
Function SEARCH(S, D, k). Given a query g and a datastore D, this function returns the ks docu- ments from the datastore with the highest similarity to q.
Function GETTOP(D, k). This function takes in descending ordered list D and returns the first elements.
A naive way of studying datastore scaling is to subsample from the raw data for every combination of subsampling ratio and random seed, as Algorithm shows. However, this approach is computationally inefficient as it would entail repeatedly running subsampling over the raw data, building index, and running retrieval search.
To affordably study datastore scaling trends, we propose efficient way to subsample over the datastore. Instead of subsampling from the raw data, for each query, we first retrieve the top- K documents from the entire datastore, and then subsample from this set. We then take the top-k from
21

---

Algorithm Naive implementation of datastore scaling
1: Input: Data corpus Dv, list of subsampling ratios P. list of random seeds S, a list of queries
Q
2: Output: Retrieved top-k documents [D] (D)a P.QEE
3:
4: function MAIN(DN, Q)
5: for E S do
6: for E do
7:
DiN + SUBSAMPIEEDNN
Complexity: O(N)
8:
Build an index over Dr
Complexity: O(pN M 11
9:
for 4 E Q do
10:
D ((....) + SEARCII( (jpp)))
4 Complexity: O
11: end for
end for
13: end for
return [D(..P.* Q.Q.EE
15:
end function
Algorithm Our efficient datastore scaling implementation
1: Input: Data corpus Dv, list of subsampling ratios P, a list of random seeds, a list of queries
Q. number of intermediate retrieved documents K (K < N)
2: Output: Retrieved top-k documents [D(..P.)] g E E P,
3:
4: function MAIN(DN, P, S, Q, K)
5: Build an index over DN
Complexity: O(N 4 Complexity: O
12
6: {D* q E Q} + SEARCH(Q, Ni,N: K)
7:
for E S do
8:
for E P do
9:
for q E Q do
10:
D(g.p.8)
<- SUBSAMPLE P.P,))
< Complexity: O(K)
(1.P..) + GETTOP(D ((a.p.s) kk)
11:
k
Complexity: O(1)
12:
end for end for
14:
end for
return [D (D(....>)
1 q P.se
16:
end function
the pool of K documents for final evaluation. Algorithm shows the pseudocode for proposed approach.
Computation comparison. The complexity of naive datastore subsampling is O(11 M IQ1) S|N). where M is the number of parameters of the retriever model, Q is the number of queries, |P] is the number of subsampling ratios, |S] the number of random seeds, and N is the number of documents in the datastore. The complexity of naive datastore subsampling is dominated by PSS N) practice, which the cost of repetitively building the datastore for all combinations of subsampling ratios and random seeds. The complexity of our subsampling strategy only O (M + [Q)) + K (() which is dominated by O(MN), representing the cost of one-time indexing.
Lemma A.3. Subsampling from the retrieved top-K K documents with probability p and then taking the top-k (k & from the subsampled documents (Algorithm is equivalent with high probubility to directly retrieving top-k documents from a datastore that is subsampled from the entire raw text datastore with probability (Algorithm 7.. The equivalence holds as long there are enough
22

---

documents left after subsampling. The chance of failure, i.e., not having enough documents left, is exponentially small K.
Proof. Assume fixed random seed set in way that whether each document in the raw data pool will be included not is determined. If document is subsampled in the naive approach, it will also be sampled in our approach, because this document has retrieval score falls in top-k and it is determined to be sampled. Therefore, as long as there are at least documents left after our subsampling over the top- K documents, the results are guaranteed to be the same. Since the number of documents remaining after subsampling follows a binomial distribution with parameters K and p. we can use standard tail bounds to show that the failure of not having enough documents is exponentially low in K.
Lemma A.3 holds with exponentially small probability for failure. The probability for failure is very small in practice, i.c., less than 1% if we have 1000 documents for subsampling. To guarantee equivalence,
use fallback mechanism where if we do not have enough documents left are subsampling, we try again with a larger value of K. However, in practice, we do not use this fallback as the failure rate very low in our experiments.
Commutability of Operations. We next discuss the commutability between subsampling and other operations. To begin, we distinguish the data operations into two sets: element-level operations and set-level operations.
Definition A.3 (Element-level operation). An element-level operation is conditioned on a single element, iee., document in our context. For example, assigning element-wise score to each document during reranking is an element-level operation.
Definition A.4 (Set-level operation). A set-level operation refers is conditioned on of elements containing at least two elements. For example, deduplication is set-level operation.
Lemma A.4. Independent element-level operations are commutable with each other. Set-level operations are not commutable.
Proof. Since the results of independent element-level operations not impacted by their order of execution, they are commutative with each other. However, this does not hold true for set-level operations, which are order-sensitive.
We note that both merging and subsampling can be considered independent element-level operations if we regard the removed documents as being classified by either process to be masked out. As results, operations such as data quality filters, data decontamination, reranking can be moved around before or after post-hoc merging, which made it possible for to efficiently evaluate the effect of their variants by moving them to after retrieval and merging.
Proposition A.1. Our MASSIVEDS pipeline is equivalent the naive pipeline, shown in Figure with high probability.
Proof. Lemmal shows that the distributed indexing and retrieval is equivalent to unsharded indexing and retrieval. Lemmas 2-4 show that the operations that we reordered between the naive pipeline and the MASSIVEDS pipeline commute without changing the returned results, with failure probability exponential in K, where the randomness is due to subsampling. Thus, the pipelines are equivalent with high probability.
A.6 Evaluation
After the aforementioned operations, we select the top-k documents from those retained for evaluation. We refer to the next section for detailed evaluation setup.
23

---

B Implementation Details
B.1 Pipeline
Sharding. We split the raw data of each domain into m shards, with m determined based on domain size. Specifically, m 32 for each time slice of CommonCrawl and C4, and m for the other domains.
Chunking. Prior to datastore construction, we chunk the raw data in each shard into fixed-length passages of at most 256 words each.
Deduplication. Deduplication is a process that removes documents with high overlap within data pool. Following Borgeaud et al. 22022) Magnusson et al. (023), we deduplicate the retrieved top-K documents based on 13-gram Jaccard similarity between document pairs. Document pairs that share at least 80% similarity are marked as duplicates, and the document with the lower retrieval score is removed.
Decontamination. Decontamination is a process that removes documents that share high similarity with evaluation data. For upstream language modeling, we apply combination of two decontamina- tion methods: 13-gram Jaccard decontamination and 32-gram longest decontamination Specifically, 13-gram Jaccard decontamination computes the 13-gram Jaccard similarity between the document and the answer, and the document is removed if it shares at least 80% similarity score with the answer. 32-gram longest decontamination removes documents that overlap with the answer by a contiguous sequence of at least 32 tokens. Note that we use 512-token answer to compute perplexity, so the decontamination ratio 0.0625 of the answer length. For downstream tasks, we apply 13-gram Jaccard decontamination and remove retrieved documents with at least 80% similarity to the test data.
Our pipeline has three main syper-parametentt k. K, and p. k is the number of documents used for evaluation. K is the number of documents retrieved before subsampling, the subsampling ratio which controls the size of the datastore. We consider & = 3, K = 1000, and p [0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 1]. We also set different random seeds for the subsampling process. We run each subsampling with three seeds (100, 101, 102) to obtain the confidence intervals in our scaling analyses. We provide lookup table of the tail bound for our subsampling algorithm's failure to provide enough documents for evaluation in Table 4 This indicates that our setup is approximately equivalent to performing expensive subsampling the raw data first based Proposition A.3
Plyper-paramette
Table 4: Lookup table of the tail bound for a binomial distribution Binomial K, K,p) with at least m number of successes.
p 0.01
p p= 0.05
p 0.1
p
0.25
p= 0.5
p 0.75
K 1000
0.9973
1.000
1.000
1.000
1.000
1.000
B.2 Language Modeling Evaluation Setup
Following Baevski & Auli (2019); Khandelwal et al. (2020); Min et al. (2023a) we split evaluation data into into lexed-length chunks of 1,024 tokens, with stride of 512 tokens. For each chunk, the first half is used as prefix and retrieval query, and the second half as the target sequence for calculating perplexity.
B.3 Downstream Evaluation Setup
Table shows the domain, metric, and sample count for each downstream task. We evaluate all downstream tasks a 5-shot setting and prepend the top-3 documents for retrieval-based LM evaluation. We adapt the Im-evalluation-harnes 13 widely used LM evaluation suite, for
downstream evaluation.
I httppp///gthhaocco///tttttttttte
24

---

Table Downstream evaluation tasks.
Task
Type
Domain
Metric
Sample Count
TQA (Joshi et al. 2017
Open- open-domain QA
Wikipedia
Exact Match
17,944
NQ (Kwiatkowski
2019
Open-domain QA
Wikipedia
Exact Match
3,610
MEDQA Jin al 2020)
Science QA
Medical
Accuracy
1,273
MMLU Henddrycks ct al. 2020
Knowledge Reasoning
Varied
Accuracy
14,042
Prompt format. We tested two variants of prompt format. The first format starts with the few-shot examples and is followed by retrieved documents. The second format starts with the retrieved documents followed by the few-shot examples. We found the LM can learn the few-shot pattern better when the few-shot exmaples are closer to the question, leading superior performance than the other way. Therefore, we stick to the second format through the paper.
B.4 Compute-Optimal Scaling Setup
Intermediate checkpoints. Thanks prior works that release intermediate checkpoints, we can study the computational scaling behaviors approximately without pretraining LMs from scratch. In particular, we consider the intermediate checkpoints provided by Pythia (Biderman al. 2023) and OLMo (Groeneveld et al.) 2024). Pythia provides checkpoints
sizes trained on up to 300B tokens from the Pile (Gao al., (020). We consider Pythia models that have at least IB parameters, i.e. PYTHIA- 1B, PYTIIA-22 8B, PYTIIA-6.3 9B, and PYTHIA-1 12B, in favor of their capability of handling complex downstream tasks. Additionally,
consider OLMO- 1.7-IB and OLMo- 7.7-BB which are trained on 3T and 2T tokens from Dolma (Soldaini et al. 2024), respectively. For Pythia, we use checkpoints trained on 1/30, 1/15, 1/10, 1/5, 174, 173, 1/2, and all of the full corpus. For OLMo, which only has models of two sizes, we select additional checkpoints trained on 1/50, 1/40, 1/20, 1/9, 1/8, 1/7, and 1/6 of the full corpus.
FLOPs calculation. Following (Kaplan et al. 2020 Hollmann et al. approximate the FLOPs needed for one forward pass as FLOPSppd FLOPspwd R 4ND, where N is the number of parameters and D is the number of tokens. Pretraining requires one forward pass and one backward pass on every token in the pretraining corpus. Denote the size of LM as NiM and the size of the pretraining corpus as Dreetrain The FLOPs for pretraining can be approximated as FLOPSpremmin Rs 6.N.M pretrinn Datastore construction requires one forward pass on every token in the datastore corpus during the embedding step. Denote the size of the retriever as Nietriever and the size of the datastore corpus as Ddttastore The FLOPs for datastore construction can be approximated as FLOPsdatastore 2 retrieever Because use flat index. i.e., no additional operations are required at the indexing step. the FLOPs for datastore construction equal the FLOPs for embedding. We note that other types of indexing, e.g., inverted file index (IVFADC) (Jégou et al. 2011), may require additional FLOPs during construction and fewer FLOPs at inference.
2022 Gadre et al. 2024). we
2ND and for backward as
C Complete Datastore Scaling Results
In this section, we supplement the datastore scaling performance of PYTHIA and OLMO models, in addition the LLAMAmodels shown in Figure Specifically, we present the datastore scaling performance for the following model families:
Pythia (Biderman et al. 2023) of 4 sizes: PYTHIA-IB, PYTHIA-2.8B,, PYTHIA-6.9B, and P2T
OLMo-IB and OLMo-7B. 2023 which we consider LLAMA-2 7B, LLAMA-2 13B, and -Laama Touvron et al.
OLMo (Groeneveld et al. 2024 of 2 sizes:
LLAMA-3 8B.
The complete datastore scaling results for TriviaQA, Natural Questions, MMLU, and MedQA are shown in Figures TO IT and respectively. We find retrieval benefits LMS of varied sizes across different LM families. In particular, the results show that small LMs outperform larger LMs of the
25

---

same model architectures when augmented with MASSIVEDS, Surprisingly, Pythia- 1B matches Pythia- 12B when augmented with only 100B tokens knowledge- -nteesivee tasks such as TriviaQA and Natural Questions, and it outperfoms Pythia-128 I2B when further increasing the size of datastore.
TriviaQA
Pythia-1B Retrieval
0.8
Pythia-1B LM-Only
Pythia- Retrieval
Pythia-2.88 LM-Only
Pythia Pythia-6.9B 6.9B Retrieval
Pythia-6.98 LM-Only
0.7
Pythia-128 Retrieval
Pythia-128 LMM-Only
OLMo 1B Retrieval
OLMo-1.7 1B LLM-Only
OLMc-1.7 7B Retrieval
OLMo 1.7 7B -LM-Only
0.6
Llama 7B Retrieval
Llama-2 7B -LM-Only
Llama-2 13B Retrieval
Llama 13B LL-Only
Llama-3 Retrieval
Llama-3 8B LM-Only
Accuracy
0.5
0.4
0.3
0.2
0.1
0
10
100
1000
Datastore Size (B Tokens)
Figure Complete datastore scaling performance on TriviaQA with PYTHIA OLMo and LLAMAmodels.
26

---

Natural Questions
0.40
Pythia-1B Retrieval
Pythia-1B LM-Only
Pythia-2.88 Retrieval
Pythia-2.88 LM-Only
0.35
Pythia-6.98 Retrieval
Pythia-6.98 LM-Only
Pythia-12B Retrieval
Pythia-128 LM-Only
OLMo-1.7 18 Retrieval
0.30
OLMo-1.7 1B -LM-Only
OLMo-1.7 Retrieval
OLMo-1.7 7B -LM-Only
Llama-2 7B Retrieval
Llama-2 -LM-Only
0.25
Llama-2 13B Retrieval
Llama-2 13B -LM-Only
Llama-3 Retrieval
Accuracy
Llama-3 8B LM-Only
0.20
0.15
0.10
0.05
0.00
10
100
1000
Datastore Size (B Tokens)
Figure 10: Complete datastore scaling performance on Natural Questions with PYTIIA OLMO and LLAMAmodels.
27

---

MMLU
0.70
Pythia-1B Retrieval
Pythia-1B LM-Only
Pythia-2.8B Retrieval
Pythia-2.88 LM-Only
0.65
Pythia-6.9B Retrieval
Pythia-6.98 LM-Only
Pythia-12B Retrieval
*
Pythia-128 LM-Only
0.60
OLMo-1.7 18 Retrieval
OLMo-1.7 1B -LM-Only
OLMo-1.7 Retrieval
OLMo-1.7 7B -LM-Only
Llama-2 7B Retrieval
0.55
Llama-2 -LM-Only
Llama-2 13B Retrieval
Llama-2 13B -LM-Only
Llama-3 Retrieval
Accuracy
0.50
Llama-3 8B LM-Only
0.45
0.40
0.35
0.30
0.25
0
10
100
1000
Datastore Size (B Tokens)
Figure 11: Complete datastore scaling performance on MMLU with PYTII OLMO and LLAMAmodels.
28

---

MedQA
Pythia-1B Retrieval
0.60
Pythia-1B LM-Only
Pythia-2.88 Retrieval
Pythia-2.88 LM-Only
Pythia-6.98 Retrieval
Pythia-6.98 LM-Only
0.55
Pythia-12B Retrieval
*
Pythia-128 LM-Only
OLMo-1.7 18 Retrieval
OLMo-1.7 1B -LM-Only
OLMo-1.7 Retrieval
0.50
OLMo-1.7 7B -LM-Only
Llama-2 7B Retrieval
Llama-2 -LM-Only
Llama-2 13B Retrieval
Llama-2 13B -LM-Only
0.45
Llama-3 Retrieval
Accuracy
Llama-3 8B LM-Only
0.40
0.35
0.30
0.25
0.20
0
100 Datastore Size (B Tokens)
10
1000
Figure 12: Complete datastore scaling performance on MedQA with PYTIIA OLMO and LLAMAmodels.
D Discussion the Inferior LLAMA-3 PPL Performance RedPajama Data
We show in Figure that LLAMA-3 8B has worse PPL scores than LLAMA-2 2 7B on RedPajama data with and without retrieval augmentation, which contradicts the intuition that LLAMA-3 is stronger than LLAMA-2 7B. In fact, Xiao et al. (2023) also reported worse PPL scores for LLAMA-3 8B than LLAMA-2 7B. We note that many factors can contribute to the difference in PPL. evaluation, such as the pretraining data and context length. For example, LLAMA-3 incorporated post-training process using instruction-tuning data, which aims enhance model alignment and output quality for complex tasks but could shift performance metrics away from those optimal for simple PPL evaluations. In
29

---

addition, LLAMA-3 was trained with significantly more data a domain (such as RedPajama) could be down-weighted in larger corpus, leading to less memorization for this certain domain during pretraining.
E Additional Analysis
In this section, we provide additional analyses of the impact of retriever and data quality filtering.
E.1 Ablation on the Retriever
In this section, we ablate the choice retriever. Since datastore construction expensive, we subsample 10% from the full corpus of MASSIVEDS for this ablation study. We build datastores with RRAGGN-RROBRTTH ROBERTA (Lin) 3 different retrievers, CONTRIEVER MSMARCO (Izacard et al., (022). DRAGON
al. 2023). SETR-55-BAR (N et al. 202I), and evaluate upstream perplexity and downstream tasks with them separately. subset of the tasks, i.e., RedPjama for language modeling evaluation, and NQ and MMLU for downstream evaluation, is used compare different retrievers. The evaluation results are shown in Table 6 We find the retrievers perform similarly.
Table 6: Ablation on the retriever. We evaluate different retrievers using 10% randomly sampled MASSIVEDS. We evaluate with LLAMA-2 7B on language modeling with RedPajama data and downstream tasks Natural Questions and MMLU, The best performance is highlighted in bold, and the second best is underlined.
Retriever
Perplexity L
Natural Questions t
MMLU f
Name
Type
Size
Contriever
dense
177M
4.2210
0.3321
0,4922
DRAGON dense
110M
4.2373
0.3399
0,4875
GTR-Base
dense
110M
4.2146
0.3080
0.4934
We empirically find the implementations of Contriever" and DRAGON 15 sentence-transfomiee implementations, e.g., GTR-Base As result, choose Contriever in full-size scaling study with a consideration of both performance and efficiency.
run much faster than the
E.2 Effect of Data Quality Filtering
Setup. We consider three data quality filters applied to the DOLMA Soldaiii et al. 2024 corpus: (I) whitespace filter which counts the number of chitespace-sseparaeed tokens each document, and filters out documents with counts under manually defined threshold: (2) a language filter which uses FastText (Bojanowski et 2017) model to detect the language of the document and filters out those with low model confidence; (3) an alphanumeric filter which only retains documents which contain at least one alphanumeric character, and do not contain span of all punctuation characters (the length of the span is user-defined).
Data deduplication. Removing data duplicates has proven effective for pretraining more compute- optimal language models (Lee et al. 2022). Duplicates are especially undesirable the context of retrieval augmentation as they repeat the same information while increase the inference cost, raising a need for global deduplication. In our default setting (Section (4. we perform global datastore deduplication based on 13-gram Jaccard similarity, similar to Lee et al. 2022). Additionally, report results without global deduplication for comparison.
Figure 13 (b) and (e) report the results on language modeling perplexity (on RedPajama) and on NQ. respectively. We find negligible impact of global deduplication in language modeling perplexity. On
hhtprr///gghhhaee///////bbbereer httpe://sbert.net/
30

---

[a] Decontaminatinn on RPj
(t) Dedupliction on RP
Tel Doima r fiters RP)
4.2
Perpoott
Perpickty
4.2
3.6
Uiinsty
trouk
35
Witt ded Aiciiii
-6-
Wittett buima fihers
agpeesbve sztinna
x
x LYory
staddrd secontamrraton
Dalma fiters
34
55
10
100
1020
10
103
1000
12
120
1000
Datastore Tokeen Size (B)
Datastore Toiien Size (B)
Datastore Taken Size (B)
[d) Deconzaminetion
Deduplication
fitters
14.00
0.52
0.35
Witlett Juma Milar
US only
Liooy
With lcation
Detma fihers
Etet Metsh
0.33
Exact Match
0.37
Feact Match
0.33
0.32
0.31
0.30
0.30
Ui-enty
With approstive decontani netion With standard bocrtanninaiin
<
K
100
103
100 Datactore Talion Size (B)
Datastora Tolkon
Datastore
Figure 13: Analysis on the effect by deduplication, decontamination, and quality filters, from left to right (Section 5533.. The first row corresponds to language modeling on REDPAJAMA and the second row shows QA results on Natural Questions (NQ).
NQ. deduplication is crucial to minimizing saturation as the datastore scales; intuitively, subsampling with higher increases the chance of seeing more duplicates 18
DOLMA Quality filtering. To study the impact of quality filtering, consider data filter adapted from DOLMA (Soldaini et al. (2024) which uses a combination of three filters: a whitespace filter; a language filter, and alphanumeric filter (detailed in Appendix Section E.2).
Figure 13 (c) and (f) indicate that the quality filter has a relatively limited effect. We hypothesize that the data sources we used in MASSIVEDS, such as RedPajama, have already gone through similar quality filtering processes and may not benefit much from applying additional filtering. While not explored in this paper, recent studies indicate that more computationally expensive but higher-quality filters can further enhance pre-training (Abbas et al. 2023) Penedo et al. 2023); we suggest future work to explore such filters for potential performance improvements.
Removing Small Chunks from MASSIVEDS We show a few examples of the top-1 retrieved documents before and after removing the short chunks that have less than 13 words in Figure [4 Without removing these chunks, we find retriever tend to retrieve documents with verbatim text overlap to the question, but do not provide helpful information about the answer, leading to degradation in end-task performance. Figure 1 indiiices that removing short chunks can significantly improve the retrieval-basee LM performance with large datastore, which is more likely to contain short, but meaningless word chunks.
1'ehe source corpus of MASSIVEDS has been applied moderate deduplication. For example, RedPajama preprocessed the data using several filters: httpp::/giithi: com/ ceammptter/eePPae. RRddaaaaaaaa jama- Data/ tree//pp.vl/data_prey However, we still find many duplicates in the retrieved results.
31

---

"question": "Answer these questions: when did the eagles win last super bowL7/nss", "before": "Eagles won the Super Boul.", who Philly about five years, agree about city(u2019s greatness
"after":
"question"; "Answer these questionsiNmkn00 who sang ran all the way hone?\nA=", "before": "AlL the way Home", "after":; "I Ran All the May Hone contains the vocal group's two hits: the snash title track and the ninor hit Oh,
"question": "Answer these questions:\alnn0 who plays gran the and the restless?\nn=",
"before"s "on The Young and the Resiless", "after": "Wichael's dad
Baldwin played Michael Gross. Wallpaper and background images the
questions:/niien developed the concept total quality menagement?Vvmt"" "before"s "Introducing total quality nanagenent.", "after": "Total quality managenent it may have been first coined in the United States by the Naval, Air Systems Con
"question":
"question"; "Answer these questions:Nnkide who wrote cant get you out of my head tyrics?\nA="' "before"s "lu2018Can/uZe19t You My Head1 "after"; "the need just can't tt't out of ny head
La, la, la, Songuriters: Cathy Dennis
"question": "Answer these questions:\nnnd last episode of what happens my family?\nA=",
"before"s "What happened to the fenily?", "after": "As child, used tove to watch the television seriiss Colunbo because the alnost every ep
"question"" "Answer these questionss\aivn who wrote it's a long long way to pasadena?lnks", "before": "Long "after": "[see 1975 music), and was written Angus Young, Malcoonn Young and Bon Scott. it's long to the
"question": "Answer these questions:Nakn0:: who the färst wife on sister Nives7\nA:", "before": "first wife?", "after": "Q: for relation between nan)u2619s
relation
two man
"question": these questiorszVnknn: does dlc
siege cone out7\nA:",
"before": "for Rainbow Six: Siege", "after": "the first-person shooter, as well as increasing some of the strategic choke-points. The DLC will enter:
"ques tinnt "Answer these quest questiuns:\ntnn; where does the last name galvez cone fron?'nn'", "before"s "Galvez", "after": "The nountainous borders contain origins prestigtous
The earliest for you get cashiers from?\nA:",
"question":
questions:Nnkkn0
"before": "Get Cashier\u2019s Check From Any Bank?", "after": "Find out where to get a cashierku2019 check and what steps to take to get one from bank, cashierluz
"question": "Answer these guestions:knivng what is the neaning of cc and bcc7Vnt:", nean security issues.", "after"; "Co stands for carbon and it that whoever
"before":
after the Cet will
of the
Figure 14: Retrieved documents for NQ when short chunks (< 13 words) are not removed. "questions" are inputs from NQ: "before" refers to the top-1 retrieved document when we do not remove short chunks; "after" refers to the top-1 document after short-chunk removal step is added to our pipeline. The retriever tends to retrieve short chunk that may have high lexical overlap with the question, but does not provide any useful information for the answer.
32

---

0.35
w/ short chunk removal w/o short chunk removal
LM-only
0.34
0.33
Exact Match
0.32
0.30
0.29
X
0
10
100
1000
Datastore Token Size (B)
Figure 15: Comparison of NQ performance with and without the removal of short chunks from MASSIVEDS. We use LLAMA-2 7B as the reader model.
33

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We ensure that the claims in our abstract and introduction closely reflect the contributions of our paper. We also carefully describe our scope (c.g., datastore scaling for particular class of retrieval-based LMs) and limitations.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations of our work in ??.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-specfiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
34

---

Answer: [Yes]
Justification: Appendix ?? covers to the theoretical contributions of our pipeline imple- mentation. We justify step-by-step approach, and sketch out the pseudocode of algorithms to provide intuition.
Guidelines:
The answer NA means that the paper does not include theoretical results, All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We provide low-level implementation details in subsection 4.1 and Appendix B including dataset construction, models used for analyses, all hyperparameters, etc.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
35

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: We are planning to open-source all contributions in this work, including MASSIVEDS indices, instructions for datastore construction, retrieved files for upstream and downstream analyses, and all code.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://niis... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-sourge benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
guwdes/ CodeeSumiiiseeeeee onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Our experiments are training-free. but we include all details for inference and analysis in Section 4.1 and Appendix B writing, and will release code to facilitate future study.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: All of the experiments in the main section of our paper are ran with three random seeds (100, 101, 102) and we report the average with 95% confidence intervals.
Guidelines:
The answer NA means that the paper does not include experiments.
36

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative eтror rates).
show in tables or
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide compute details Appendix B
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: We have reviewed the NeurIPS Code of Ethics, and strictly follow it in this work
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Yes, discuss broader impacts of work in ??.
37

---

Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (c.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., denloymemt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: We do not believe this question is applicable: our datastore is amassed entirely from open-source data (e.g., RedPajama) for which safety filters (e.g.. removal of PII or NSFW content) have already been applied.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (c.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Yes. We acknowledge the codebases have referred to, cite and provide the urls of datasets and models in sesctiinn
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
38

---

The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, pwwpsssipppsseeeee thcode com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We provide careful documentation of the datastore construction process in our paper.
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details the dataset/code/moodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: We do not involve any crowdsourcing human subjects in our study, and therefore this question is not applicable.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
supplemental material is fine, but if the main contribu-
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
39

---

Justification: Our study did not involve any crowdsourcing or human subjects, and thus this question is not applicable.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
40


Reviews:
[{"review_id": "ddixAHNCXH", "reviewer": "Reviewer_F6AB", "summary": "The paper studies the effects of scaling retrieval datastore in retrieval augmented language models. The authors present the impacts of various design choices such as data size and data selection. A testbed dataset MassiveDS is also introduced.", "questions": "- It was discussed in the paper that MASSIVEDS can outperform in-domain datastores on two eval sets. I wonder how this happens. Are there pieces of information easier to retrieve and/or read from out-of-domain? If so, what forms can this information take to be easier to read and retrieve than the original information?\n - The authors claim to observe an empirical log-linear scaling trend in the RAG system. It would be interesting to see some discussion on why this is the case and the implications of this. How does this happen? Does this mean retrieval datastore scaling has an extremely diminishing gain?\n- The authors use token counts extensively across the paper. While this is common in describing pre-training data, the actual retrieval units, i.e. passage/chunk is arguably a better measurement to describe the statistics of the retrieval task. Can the authors provide some statistics based on chunks/passages?", "limitations": "n.a.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The authors present a substantially larger retrieval datastore for retrieval augmented LMs.\n- Extensive experiments have been conducted to study how datastore designs/properties can impact RAG results.\n- Detailed analysis has been conducted on several eval datasets and models.", "weaknesses": "- Large-scale embedding-based search is not a new topic but has been studied for years. Arguably it is a mature technology ready in production [1, 2].  The authors have made a rather arbitrary design in the index design. Information like hardware latency and memory/CPU footprint were not reported; comparisons with popular techniques like DiskANN [2] are not presented. This could be confusing and/or misleading to people reading the paper: should they adopt the search pipeline introduced in the paper, or just use a more developed technology? Also, this makes it difficult to assess the contribution of the entire paper.\n- While the paper shows a general improving trend in RAG tasks and with relatively capable LMs, the models in the paper still seem to underperform the popular 2-year-old RAG model ATLAS [3] by a significant margin. Specifically, the design of the RAG system used in the paper can be somewhat basic. How much the results transfer to stronger system/system designs remains relatively unclear.\n\n\n[1] Huang, J., Sharma, A., Sun, S., Xia, L., Zhang, D., Pronin, P., Padmanabhan, J., Ottaviano, G., & Yang, L. (2020). Embedding-based Retrieval in Facebook Search. Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining.\n\n[2] Subramanya, S.J., Devvrit, Kadekodi, R., Krishaswamy, R., & Simhadri, H.V. (2019). DiskANN : Fast Accurate Billion-point Nearest Neighbor Search on a Single Node.\n\n[3] Izacard, G., Lewis, P., Lomeli, M., Hosseini, L., Petroni, F., Schick, T., Yu, J.A., Joulin, A., Riedel, S., & Grave, E. (2022). Few-shot Learning with Retrieval Augmented Language Models. ArXiv, abs/2208.03299.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for pointing out the comprehensiveness of our work! We would like to address the concerns and questions below, and we will edite the paper accordingly.\n\n**Weakness 1.** Large-scale embedding-based search is not a new topic but has been studied for years. Arguably it is a mature technology ready in production [1, 2]. Information like hardware latency and memory/CPU footprint were not reported; comparisons with popular techniques like DiskANN [2] are not presented. This could be confusing and/or misleading to people reading the paper: should they adopt the search pipeline introduced in the paper, or just use a more developed technology? Also, this makes it difficult to assess the contribution of the entire paper.\n\n**A1.** The work listed by the reviewer focuses on the search efficiency at inference with a fixed datastore configuration, but we would like to clarify that this is orthogonal to the focus of the current paper. Our main goal is to conduct a thorough analysis of the effect of scaling the datastore for retrieval-based language models and the impact of various datastore factors, where the inference speed isn’t a key bottleneck—even with the state-of-the-art nearest neighbor search used in productions such as DiskANN, studying the effect of various factors of the datastore like the size, quality filters, and decontamination methods remains expensive. This is because such experiments require rebuilding search indices for every combination of the factors, unrelated to which search algorithm is being used. Therefore, we design a new pipeline for efficient experimentation of various factors of the datastore by removing the need for repetitive datastore construction and repetitive large-scale retrieval, which are orthogonal and can be combined with prior work on efficient search algorithms. \n\n---\n\n**Weakness 2.** While the paper shows a general improving trend in RAG tasks and with relatively capable LMs, the models in the paper still seem to underperform the popular 2-year-old RAG model ATLAS [3] by a significant margin. Specifically, the design of the RAG system used in the paper can be somewhat basic. How much the results transfer to stronger system/system designs remains relatively unclear.\n\n**A2.** Despite active research on advancing RAG designs, there have been few scientific studies of the scaling properties of datastores with more than a few billion tokens. As an initial open-source study on datastore scaling, we picked the most basic but widely adopted RAG design instead of any particularly optimized system. In Figure 5, we presented an analysis of the impact of improved retrieval on the datastore scaling performance, which shows that adopting a more advanced retrieval technique, i.e., a reranker, can further improve the scaling performance, indicating a positive sign that our conclusions could potentially be extended to more recent RAG designs. Overall, we focus on conducting a comprehensive analysis of datastore scaling and sharing our thoughts on how future works could further improve the scaling curves, rather than achieving state-of-the-art scores on specific datasets with a particularly optimized RAG design.\n\nThe purpose of Atlas is to finetune the retriever for best performance. Atlas finetunes the retriever and the LM to adapt to every specific task using the full task data (>10k) for their best-reported scores, while we evaluate the model in a training-free fashion by prepending only 5 examples in context. The purpose of our work is to disentangle the factors that affect performance and study them in detail. As such, Altas is orthogonal work. The insights from our work can be combined with Atlas, but here we ask more basic questions about RAG scaling.\n\n---\n\n**Question 1.** It was discussed in the paper that MASSIVEDS can outperform in-domain datastores on two eval sets. I wonder how this happens. Are there pieces of information easier to retrieve and/or read from out-of-domain?\n\n**A3.** MassiveDS is a hybrid multi-domain datastore that includes both “in-domain datastores” as subsets and general web data to enhance its knowledge base. The key finding here is that the model is robust to other out-of-domain data in the same datastore and the retriever still preferentially retrieves from the right domain (Figure 3), so it eliminates the need to develop domain-specific datastores for each task, which is often a costly and complex process. Consequently, users can establish a single, general-purpose datastore that is effective across various tasks, even in cases where creating a task-specific in-domain datastore is unfeasible. This flexibility greatly simplifies data management and expands the utility of MassiveDS in diverse applications.\n\n---\n**Question 2.** The authors claim to observe an empirical log-linear scaling trend in the RAG system. It would be interesting to see some discussion on why this is the case and the implications of this. How does this happen? Does this mean retrieval datastore scaling has an extremely diminishing gain?\n\n**A4.** Many previous scaling law studies have shown significant gains by scaling the pretraining data and the model size, and they all present performance gains with the x-axis (the number of tokens for pretraining or the number of parameters) in the log scale [1,2]. We find similar performance gains by scaling the datastore to the gains by scaling pretraining data. Therefore, our results indicate that the datastore could be another dimension to scale in addition to the pretraining data and the model size. \n\n[1] Gadre, Samir Yitzhak, et al. \"Language models scale reliably with over-training and on downstream tasks.\"\n\n[2] Kaplan, Jared, et al. \"Scaling laws for neural language models.\"\n\n---\n**Question 3.** Can the authors provide some statistics based on chunks/passages?\n\n**A5.** We have 4B passages in our datastore, which average around 360 tokens. We will report detailed statistics in the next version of the paper.", "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": "bTckmrhbCg", "reviewer": "Reviewer_cEKQ", "summary": "The paper \"Scaling Retrieval-Based Language Models with a Trillion-Token Datastore\" explores a new dimension of scaling language models (LMs) by considering the amount of data used during inference. The study focuses on retrieval-based LMs, which access a large external datastore during inference, and examines the impact of scaling the datastore. The authors introduce MASSIVEDS, a 1.4 trillion-token datastore, and an efficient pipeline to study various datastore features such as size, data filters, and decontamination strategies. The experiments reveal that datastore scaling follows a log-linear trend and significantly improves performance on various tasks. The authors also highlight the importance of aggressive decontamination, data deduplication, and improved retrieval techniques for better scalability. If accepted, the dataset MASSIVEDS will be the largest and also most diverse open-sourced datastore which will benefit the research community.", "questions": "1. Previous studies (https://arxiv.org/html/2307.07164v2) have shown that lack of diversity adversely impacts the results. What was the diversity of results retrieved by the underlying retrieval?", "limitations": "The authors adequately address the limitations of their work, such as focusing on a specific class of retrieval-based LMs, limited model sizes, and dense retrievers. They also acknowledge the need for further exploration of different retrieval architectures and a wider range of evaluation tasks.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "- **Originality**: The paper tackles the important question of how scaling retrieval-based LMs by the size of the inference datastore impacts upstream and downstream task performance which is novel and adds a new dimension to LM scaling laws.\n- **Quality**: The research is rigorous, with well-designed experiments and robust methodologies.\n- **Clarity**: The paper is well-organized, and the key contributions are clearly articulated.\n- **Significance**: The open-sourced resources will facilitate further research in this direction.\n- **Experiment Design**: The design of the experimental set-up is novel and theoretically grounded which will be interest to the research community on how to efficiently experiment with multiple dataset configurations", "weaknesses": "The following are some of the limitations of the paper (some of these have already been highlighted by the authors):\n- **Task Diversity**: The evaluation is primarily focused on QA and knowledge-intensive tasks where one can expect the retrieval based LMs to work well. Inclusion of more diverse tasks, such as commonsense reasoning or open-ended text generation, would provide a more comprehensive assessment of the datastore's scalability.\n- **Limited Scope of Analysis**: This is a generic comment about the overall analysis results in the paper. For each of the analysis dimension, the authors have scratched the surface by experimenting with a basic technique which I am afraid makes the results limited in their scope and generalizability. Here are a few examples to illustrate the point - a) quality filtering: they explore basic filtering from DOLMA but leave the higher-quality filters for future-work b) Decontamination: 13-gram models for decontamination and leave out some techniques such as Min-K% Prob which are known to be stronger baselines c) Retrieval Model: The paper uses a specific retrieval model (CONTRIEVER-MSMARCO) without exploring alternative models. This might be due to the fact that a major portion of the paper was focused designing the scalable experimental set-up for this analysis.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for acknowledging the novelty and solidness of our work! We would like to address the concerns and questions below:\n\n**Weakness 1.** The evaluation is primarily focused on QA and knowledge-intensive tasks where one can expect the retrieval based LMs to work well. Inclusion of more diverse tasks, such as commonsense reasoning or open-ended text generation, would provide a more comprehensive assessment of the datastore's scalability.\n\n**A1.** We agree that more diverse tasks would make our study more comprehensive. Our study takes a step in this direction as it is the first analysis, to the best of our knowledge, of datascore scaling on downstream tasks. Previous studies (RETRO, Borgeaud, et al.) examine retrieval-based language models with datastores exceeding 1 trillion tokens have focused solely on language modeling, and it remains unclear how datastore scaling and different datastore factors impact the performance of downstream tasks. \nIn addition to knowledge-intensive tasks, we also include tasks that involve reasoning on top of knowledge. In particular, we included MMLU, which is widely recognized as a standard metric for assessing the performance of modern LMs, in our evaluation. Given our resource constraints, it was challenging to cover more downstream tasks, as we discussed in our limitation sections. We will emphasize this point in the next revision, and we hope that by open-sourcing our datastore and pipeline, we will enable others to build upon our work and evaluate large-scale retrieval-based models on a more diverse set of downstream tasks.\n\n---\n\n**Weakness 2.** For each of the analysis dimensions, the techniques are basic: E.g., a) quality filtering: basic filtering from DOLMA but leave the higher-quality filters b) Decontamination: 13-gram models for decontamination and leave out some techniques such as Min-K% Prob c) Retrieval Model. This might be due to the fact that a major portion of the paper was focused designing the scalable experimental set-up for this analysis.\n\n**A2.** We thank the reviewer for acknowledging that our main focus is on the design of scalable experimental setup, as it is challenging to cover all possible datastore configurations given our computational constraints. Nevertheless, our analyses have covered several key aspects that people care about when designing the datastores, and our findings are inspiring.\n\nHere are some discussions on the points mentioned by the reviewer: \n\n(1) We agree with the reviewers that the DOLMA filters are basic and that higher-quality filters, such as fasttext filtering [1], are not covered in our study. We focus on standard filters that have been applied to both traditional and modern pretraining corpora such as RedPajama, C4, Dolma, etc. We believe these analyses are interesting to a broad community. In addition, we show that it is easy to test the effect of different quality filters using our pipeline and defer the study on more recent filters to future works. \n\n(2) For data decontamination, we considered 13-gram Jaccard similarity-based decontamination because it has been widely adopted by many works, such as GPT-3, RETRO, Dolma, etc.  Beside this, we also applied another decontamination method, longest-string decontamination, which enables us to experiment with different strictness levels of decontamination, as shown in Figure 4. By tuning the hyper-parameters of the decontamination methods, we show datastore scaling can benefit the model’s performance across different decontamination levels (Figure 4), so our conclusions hold regardless of the strictness of the decontamination method. Separately, Min-K% prob is designed for detecting if a given LM has been trained on certain data, and isn’t intended to be a data decontamination method. \n\n(3) We indeed only used a single base retriever model in the submission, so we add a new ablation on the base retriever using 10% of MassiveDS and supplement the results in the table below. We chose DRAGON and GTR-Base because they rank higher than Contriever on the MTEB benchmark and are similar in size. Interestingly, the results indicate that these base retrievers exhibit similar performance on MassiveDS. We hypothesize that this is due to the differing domain compositions of MassiveDS compared to the datasets used in MTEB.\n\n| Name       | Perplexity ↓ | Natural Questions ↑ | MMLU ↑ |\n|------------|--------------|----------------------|--------|\n| Contriever | 4.2210       | 0.3321               | 0.4922 |\n| DRAGON     | 4.2373       | 0.3399           | 0.4875 |\n| GTR-Base   | 4.2146   | 0.3080               | 0.4934 |\n\nDue to limited compute, we leave it to future work to explore the performance with a larger and more capable retriever such as GRIT-7B. But we note that such larger embedding models are often prohibitively expensive to scale up to a trillion-token datastore (Wang et al., 2024). However, we are optimistic that significantly improving the retriever will lead to better scaling trends. This expectation is supported by the evidence in Figure 5, which shows that enhancing retrieval quality with a cross-encoder reranker improves performance.\n\n[1] Li, Jeffrey, et al. \"DataComp-LM: In search of the next generation of training sets for language models.\"\n\n---\n\n**Q3.** Previous studies (https://arxiv.org/html/2307.07164v2) have shown that lack of diversity adversely impacts the results. What was the diversity of results retrieved by the underlying retrieval?\n\n**A3.** As shown in Figure 3, the retriever retrieves from diverse sources in MassiveDS, and it tends to retrieve from relevant domains more frequently than other domains. Our results also indicate that increasing diversity helps improve the performance—MassiveDS outperforms single-domain datastores on the evaluated tasks, as shown in Table 3. We will discuss this in our next version.", "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": "Tpukitdd6x", "reviewer": "Reviewer_TsAT", "summary": "The paper introduces MASSIVEDS, the largest and most diverse open-sourced datastore for retrieval-based language models, containing 1.4 trillion tokens. The authors design a MASSIVEDS pipeline to efficiently explore the impact of different datastore features by reordering the datastore construction operations so that the most expensive operations, such as indexing and retrieval, are only run once. Extensive experiments demonstrate that model performance improves as the datastore size increases.", "questions": "Please see my comments in Weakness.", "limitations": "The limitations are properly stated in the Conclusion and Limitations section.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "(1) To my knowledge, this is the first work studying scaling laws regarding datastore size. The author introduces the largest and most diverse open-sourced datastore for retrieval-based language models, containing 1.4 trillion tokens. This large-scale datastore can serve as a promising testbed for developing new retrieval-based LMs.\n\n(2) The author promises to release not only the data but also the code to reproduce the experiments. This could greatly facilitate research efforts within the community.\n\n(3) The author proposes a simple yet effective data processing pipeline that enables efficient investigation of the impact of different datastore features.\n\n(4) Extensive experiments show that the performance of retrieval-based LMs benefits from scaling datastore sizes. The evaluation includes both language modeling and downstream question answering. The results on MMLU and MedQA also highlight the need for developing retrieval-augmented LMs that can excel in tasks requiring reasoning abilities.\n\n(5) Some ablation studies provide interesting and insightful findings, such as data deduplication being a crucial factor for enhancing language model performance when the datastore size is extremely large.\n\nOverall, I believe this is a good paper studying an important problem supported by solid experiments.", "weaknesses": "Overall, I do not find any significant weaknesses in the paper. The only minor issues I noticed are:\n\n(1) The retriever and reranker used in this paper may be somewhat outdated. Recent models on the MTEB leaderboard might offer better performance. The benefits of scaling datastore sizes could be more significant with a more powerful retriever. However, this is a very minor point. Considering the computational cost of indexing all the data, trying other retrievers is not needed.\n\n(2) Typo: L223 \"Datascore\" should be \"Datastore\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for acknowledging the importance of our study and the solidness of our work! We would like to address the only concern by the reviewer on the choice of retriever below.\n\n**Q1.** The retriever and reranker used in this paper may be somewhat outdated. Recent models on the MTEB leaderboard might offer better performance. The benefits of scaling datastore sizes could be more significant with a more powerful retriever. However, this is a very minor point. Considering the computational cost of indexing all the data, trying other retrievers is not needed.\n\n**A1.** We thank the reviewer for understanding our computational constraints, and we agree with the reviewer that the base retriever we used in this paper isn’t new. We chose Contriever-MSMACRO because it was used in many previous works cited in L178-179. We agree that benefits of scaling datastore sizes could be more significant with a more powerful retriever, which is consistent with our conclusion in Section 6.2 that improved retrieval could enhance scaling trends. In the table below, we tried 2 more base retrievers that outperform Contriever on the MTEB leaderboard using 10% randomly sampled MassiveDS. Our results show that, interestingly, the performance of the base retrievers of similar sizes on our general web data does not necessarily align with the ranking on MTEB. We hypothesize that this low correlation is because the domain compositions of MassiveDS and MTEB are different. Due to limited computational resources, we defer the study with larger retrievers, such as GRIT-7B, to future work. But we note that such larger embedding models are often prohibitively expensive to scale up to a trillion-token datastore (Wang et al., 2024). \n| Name       | Retriever Type | Size  | Perplexity ↓ | Natural Questions ↑ | MMLU ↑ |\n|------------|----------------|-------|--------------|----------------------|--------|\n| Contriever | dense          | 177M  | 4.2210       | 0.3321               | 0.4922 |\n| DRAGON     | dense          | 110M  | 4.2373       | **0.3399**           | 0.4875 |\n| GTR-Base   | dense          | 110M  | **4.2146**   | 0.3080               | **0.4934** |\n\n---\n\n**Q2.** Typo: L223 \"Datascore\" should be \"Datastore\".\n\n**A2.** We thank the reviewer for pointing out the typo for us. We will fix it in our next version. We are happy to address more questions from the reviewer during the rebuttal period.", "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": "L5oWITqcgn", "reviewer": "Reviewer_omwP", "summary": "This paper studies the impact of scaling the datastore (retrieval dataset) on retrieval-based language models.\n\nThe contributions are:\n- MASSIVEDS a 1.4 trillion-token datastore for retrieval-based LMs that will be made open-source.\n- A pipeline to study the impact of the datastore scaling on the language models at inference time\n- An analysis of the impact of scaling and other pre-processing steps (decontamination, deduplication, and quality filtering) on the language models at inference time", "questions": "1) CODE NOT AVAILABLE?\n\nUnless there is an oversight on my part, I could not find the implementation (no link to anonymous GitHub and no supplementary material) and I could not find the code attached to the paper. Although the authors indicated that the code and data will be made available, having an anonymous GitHub repo with the code already available would improve the strength of this submission.\n\n2) RETRIEVAL AUGMENTED GENERATION?\n\nSince the LMs used in the experiment are all generative language models, why is the term Retrieval Augmented Generation not used to describe this work? ]\n\n3) DATA DECONTAMINATION/DEDUPLICATION\n\nCan more details be provided about data decontamination/deduplication? In practice, what is the difference between the two? Since data decontamination filters out \"documents with 80+% 13-gram Jaccard similarity or 32-gram longest sequence overlap\", how is deduplication made? is it only on documents matching exactly at the string level?\n\n4) FIGURE 4 (small detail):\n\nIs there a reason why the baselines LM-only on the subplots (d), (e), and (f) are not extended with a dashed line?", "limitations": "The authors discuss several limitations of their work:\n- Focus on a particular class of retrieval-based LMs (RIC-LMs), \n- Do not study very large LMs (only 1B and 7B parameters)\n- Only consider dense retrievers. \n- Only considers QA tasks\n\nThere is only one limitation to this work that was not mentioned by the authors: the lack of analysis of the impact on inference time (see weaknesses).", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1) OPEN SOURCE\n\nBoth the models and data used in the work are open source and when the code is made available, it will facilitate future work and additional research on retrieval-based LM at scale.\n\n2) WRITING/PRESENTATION\n\nThe paper is easy to read, the contributions and findings are explicitly stated, and the tables and figures are easy to read.\n\n3) IMPACT\n\nGiven the importance of retrieval augmented generation for making LM more trustworthy and adapted to a specific domain/environment, releasing a large-scale dataset and pipeline for retrieval-based LMs can facilitate future research in that domain.", "weaknesses": "INFERENCE TIME ANALYSIS\n\nThis work does not study the impact of the scaling of the datastore on inference time. Since scaling the data store can make retrieval slower and therefore, make the generation slower, this is a crucial aspect to study, that was not addressed in this work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for highlighting our open-source contribution and acknowledge the potential impact on the community! We would like to address the reviewer’s concerns below:\n\n**Weakness 1.** This work does not study the impact of the scaling of the datastore on inference time.\n\n**A1.** Our main focus is to see the effect of dataset scaling on upstream and downstream model performance. While we agree that a larger datastore can introduce additional latency, there's active research on improving the speed of nearest neighbor search; we believe our contributions are orthogonal. In addition, we show a small LM with retrieval augmentation could match or even outperform larger LM-only models, as supplemented in Figure R1 in the general reply, which indicates a potential reduction of inference time by using a small LM augmented with MassiveDS. Given the complexity of inference speed optimization and its loose relationship with our focus, we defer the study of inference-time efficiency-performance tradeoff to future work.\n\n---\n\n**Question 1.** The codes are not available by the time of submission.\n\n**A2.** We thank the reviewer for reminding us to provide anonymous codes with the submission. We uploaded our code in an anonymous github repository and sent it to the AC according to the rebuttal rules. We hope the reviewer could easily get access to the code.\n\n---\n\n**Question 2.** Why don’t the authors use the term “retrieval augmented generation (RAG)”?\n\n**A3.** We define retrieval-based LMs as a general family of LMs that leverage large-scale text data at inference time [1], and RAG [2] is one such model. We believe that our finding is generally applicable to other types of retrieval-based LMs such as kNN-LM style models, and therefore decided to use the term. More specifically, we focus on retrieval-in-context LMs as mentioned in Section 2, which is often used interchangeably with RAG [3] but more specifically describes off-the-shelf LMs that use retrieved context at inference time.\n\n[1] ACL 2023 Tutorial: Retrieval-based Language Models and Applications https://acl2023-retrieval-lm.github.io/\n\n[2] Lewis, Patrick, et al. \"Retrieval-augmented generation for knowledge-intensive nlp tasks.\"\n\n[3] Min, Sewon, et al. \"Silo language models: Isolating legal risk in a nonparametric datastore.\"\n\n---\n\n**Question 3.** Can more details be provided about data decontamination/deduplication? \n\n**A4.** The motivation of deduplication for a datastore is to remove documents that are exactly or approximately the same such that the retrieved top-k documents won’t contain repetitive information. To achieve this, a similarity score is computed between every two documents in the datastore, and one document is removed from every duplicate pair. While decontamination aims at removing the same or near-duplicate documents in the datastore that directly contain the information about the test samples. The motivation is to avoid having test samples directly in the datastore such that the model can achieve high evaluation performance by cheating. The test sample is compared against every document in the datastore, and the documents with high similarity scores are removed from the datastore for this certain task.\nIn our paper, the 13-gram Jaccard similarity score is one metric we use to measure the overlapping rate between two documents (in deduplication, it’s 2 documents from the datastore; in decontamination, one document is the test sample and another document is from the datastore). We set the threshold to 80% for both deduplication and decontamination that we remove a document if it has a similarity score that is higher than 80% with another document. Since decontamination plays an important role in evaluation, we combined the 13-gram Jaccard similarity metric with another metric for decontamination to make our datastore less likely to be contaminated. More details can be found in Appendix B.1.\n\n---\n\n**Question 4.** Is there a reason why the baselines LM-only on the subplots (d), (e), and (f) are not extended with a dashed line?\n\n**A5.** We thank the reviewer for pointing out this detail! We forgot to add a dashed line for these 3 subfigures and we will add them back in the next revised version!", "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": "2mFkzVALyy", "reviewer": "Reviewer_vyTD", "summary": "This paper studies the effect of scaling datastores for retrieval-based language models. A trillion-token datastore, MassiveDS, is constructed and then filtered to remove contaminated and duplicate documents. A distributed pipeline is proposed to index and retrieve from MassiveDS with a modest compute budget. Evaluation on language modeling tasks and downstream QA tasks shows that scaling the datastore brings clear improvements, and careful datastore filtering is crucial for model performance.", "questions": "1. What is the recall rate of the retriever for different datastore sizes?", "limitations": "The limitations section is included and adequately discussed.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "1. This paper studies an important research question: the scaling of datastores for retrieval-based LMs.\n\n2. This paper open-sources the 1.4T token datastore for retrieval-based LMs, which consists of diverse domains and is carefully filtered. This datastore could be a valuable resource for future retrieval-based LM studies.\n\n3. Evaluation on multiple tasks demonstrates the importance of datastore scaling and highlights some interesting questions, such as how to improve retrieval performance and how to filter datastores.", "weaknesses": "1. The claim of proposing the 'largest' datastore could be reconsidered. For LM pre-training, there are larger datastores like RedPajama-v2 with 30T tokens, and for IR, there are corpora like ClueWeb22 with 16T tokens. The proposed MassiveDS is mainly sourced from the RedPajama, making it similar to other existing collections.\n\n2. The proposed distributed index and retrieval pipeline is interesting but seems inflexible. How does it compare with existing search engine frameworks like ElasticSearch and Weaviate?\n\n3. This paper lacks an analysis of retriever. Contriever-msmarco is used, whose training data may differ from MedQA and MMLU. Considering the limited generalizability of dense retrievers (as shown by the BEIR paper), stronger retrievers like those in MTEB could be used.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewers for acknowledging the importance of our research question and the value of our findings and open-sourced resources! We would like to address the reviewer’s concerns below:\n\n**Weakness 1.** The claim of proposing the 'largest' datastore could be reconsidered as there are other pretraining corpus that exceed even 10T tokens.\n\n**A1.** We would like to clarify that a datastore is different from a pretraining corpus: by “datastore”, we refer to an index that is ready for retrieval [1]; while a pretraining corpus only contains the raw data. Building a datastore on an existing pretraining corpus requires further data cleaning, data chunking, embedding, and indexing over the raw data, which is a non-trivial effort for the retrieval community. It has been challenging to conduct datastore scaling research at this scale before and existing open-sourced datastores are much smaller than MassiveDS, as shown in Table 1. Althogh there is pretraining corpus that is larger than our datastore in terms of raw text, we are the first to make this attempt to construct and open-source a datastore with over 1 trillion tokens.\n\n[1] Asai, Akari, et al. \"Reliable, adaptable, and attributable language models with retrieval.\"\n\n---\n\n**Weakness 2.** The proposed distributed index and retrieval pipeline is interesting but seems inflexible. How does it compare with existing search engine frameworks like ElasticSearch and Weaviate?\n\n**A2.** We would like to note that our contribution is orthogonal to these existing frameworks: prior work, such as ElasticSearch and Weaviate, focuses on efficiency at inference, while our work focuses on the scientific evaluation of retrieval performance which requires not efficient serving, but efficient experimentation of various factors of the datastore. Even with the state-of-the-art nearest neighbor search used in production, studying the effect of various factors of the datastore like the size, quality filters, and decontamination methods remains expensive. This is because such experiments require rebuilding search indices for every combination of the factors, orthogonal to what search algorithm is being used. Our study focuses on removing the needs for such repetitive rebuilding of indices. The outcome is the comprehensive study of the impact of various datastore factors, as we demonstrated as part of the results, which is the novel part of our work.\n\n---\n\n**Weakness 3.** This paper lacks an analysis of retriever. Contriever-msmarco is used, whose training data may differ from MedQA and MMLU. Considering the limited generalizability of dense retrievers (as shown by the BEIR paper), stronger retrievers like those in MTEB could be used.\n\n**A3.** We agree with the reviewer that we indeed only investigated one base retriever due to being limited by academic computational resources, and we studied the effect of improved retrieval by applying a cross-encoder reranker which is more computationally efficient in the paper. In the table below, we run the evaluation using 2 more base retrievers of similar sizes that perform better than Contriever on the MTEB benchmark using 10% randomly sampled MassiveDS. Our results show that, interestingly, the performance of the base retrievers of similar sizes on our general web data does not necessarily align with the ranking on MTEB. We hypothesize that this low correlation is because the domain compositions of MassiveDS and MTEB are different. Due to limited computational resources, we defer the study with larger retrievers, such as GRIT-7B, to future work. But we note that such larger embedding models are often prohibitively expensive to scale up to a trillion-token datastore (Wang et al., 2024). \n| Name       | Retriever Type | Size  | Perplexity ↓ | Natural Questions ↑ | MMLU ↑ |\n|------------|----------------|-------|--------------|----------------------|--------|\n| Contriever | dense          | 177M  | 4.2210       | 0.3321               | 0.4922 |\n| DRAGON     | dense          | 110M  | 4.2373       | **0.3399**           | 0.4875 |\n| GTR-Base   | dense          | 110M  | **4.2146**   | 0.3080               | **0.4934** |\n\n---\n\n**Question 1.** What is the recall rate of the retriever for different datastore sizes?\n\n**A4.** In our evaluation setup, both upstream language modeling and downstream evaluation sets do not provide gold documents for the questions. Therefore, we can only report the end-to-end performance on these tasks.", "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": "GRwvotU8p2", "reviewer": "Reviewer_AKfc", "summary": "The paper introduces a substantial datastore named MASSIVEDS for retrieval-in-context language models (RIC-LM). It details a comprehensive construction pipeline, with a notable deviation from the traditional datastore construction sequence. Specifically, it places indexing and retrieval at the initial stage, followed by the merging and filtering of retrieved results. The evaluation indicates that RIC-LM with MASSIVEDS, outperforms standalone language models in knowledge-intensive tasks. Additionally, the authors analyze the scaling behavior of the datastore under various configurations.", "questions": "1. Could you elaborate on how few-shot exemplars are selected for QA tasks?\n\n2. For downstream evalaution, how is the language model configuration determined? Is there any inherent randomness in the LM inference process, and if so, what measures are taken to mitigate its effects? \n\n3. Following W3, have the authors verified whether the downstream evaluation dataset overlaps with the datastore? How is this ensured, and how do you analyze the extent to which observed improvements might be influenced by potential data leakage?\n\n4. Given that many of the stages within MASSIVEDS, such as parallel indexing and aggregating retrieval results from different sources, are common practices in retrieval pipelines, how do you substantiate that the contribution of MASSIVEDS is significant and not merely an assembly of various data segments?\n\n5. Suggestion: given the emphasis on streamlining the datastore construction pipeline, proposing multiple variants of the datastore setup (such as variations in source and chunk size) might be more advantageous than focusing solely on creating the largest and most diverse datastore. This could facilitate more comprehensive analyses and support broader conclusions for future works in this domain.", "limitations": "In the limitation section, the authors did acknowledge their limited analysis concerning the systematic range of retrievers and language models, which is indeed crucial. However, it could be problematic to draw broad conclusions on the datastore scaling trend from such a restricted dataset and limited scope of experiments.  A more comprehensive analysis is necessary, or alternatively, the conclusions should be more narrowly defined to ensure precision and avoid misleading interpretations.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The exploration of how datastore configurations impact retrieval-based language models is both intriguing and significant. Some of the findings provide valuable contributions to the advancement of these systems.\n\n2. The paper is well-written and accessible, with each step of the implementation thoroughly explained and presented.\n\n3. The proposed datastore holds potential for related research in this field.", "weaknesses": "1. The paper attempts to draw critical and general conclusions, yet the scope and depth of their experiments lack the robustness required to fully support these assertions. The study focuses on the impact of various datastore configurations, such as size, filters, and decontamination strategies. Evaluations are carried out on OpenQA tasks using a retriever specifically trained for QA. However, several key questions remain unaddressed. For example, it is unclear whether the findings would remain consistent using a single data source, whether the conclusions apply to specific tasks that are distinct from broader tasks like QA, and how the interplay between different types of retrievers and language models might affect the results. Without exploring these areas, claiming a definitive trend is premature.\n\n2. Regarding the proposed MASSIVEDS, there is a lack of a comparative analysis and comparson with other existing datastores [1, 2]. Merely assembling a large and diverse datastore does not inherently indicate a significant contribution or guarantee enhanced performance. Evaluating the proposed datastore against these established frameworks could help justify claims of superiority or uniqueness. \n\n3. There are several major assertions may be incorrect or lack sufficient justification. For instance:\n    - The claim that indexing the entire datastore initially reduces computational overhead is debatable. In practice, dense vectors can be stored and reused, and there is no need to re-index any data twice. Therefore, the upper bound of computational cost remains indexing the entire datastore. This aspect could be also alleviated significantly by considering alternative indexing strategies [3, 4].\n    - The overlap between the datastore and the downstream evaluation dataset raises concerns about data leakage.  The presence of RedPajama data in MASSIVEDS (Table 2), and its use in downstream evaluation (line 191) necessitates a clearer delineation of this overlap and its impact on the results. Clarifying this could address the validity of the claim that \"Retrieval is strictly beneficial for language modeling\" (line 214).\n    - The use of models such as TinyLlama and Llama-7B to conclude that \"retrieval effectively closes the performance gap between small and large models\" (line 225) appears inconclusive in the context of current advancements in LLM. Reevaluating this claim with additional models might provide a clearer picture.\n\nTypos:\n- line193 into into -> into\n- line726: exmaples -> examples\n\n\n[1] Soldaini, Luca, et al. \"Dolma: An open corpus of three trillion tokens for language model pretraining research\" \n\n[2] Gao, Leo, et al. \"The pile: An 800gb dataset of diverse text for language modeling\" \n\n[3] Li, Haitao, et al. \"Constructing tree-based index for efficient and effective dense retrieval\"\n\n[4] Zhou, Jiawei, et al. \"Semi-Parametric Retrieval via Binary Token Index\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for acknowledging that our work makes valuable contributions to the advancement of these systems! \n\n**Weakness 1.** It remains uncertain if the findings would hold with a single data source, whether they apply to specific tasks beyond general ones like QA, and how interactions between various retrievers and language models might influence the results.\n\n**A1.** We will carefully adjust our claims in our final draft to highlight our choice of models and retrievers, and that our downstream tasks focus on QA. Given resource constraints, we chose to focus on the most widely adopted RAG designs and evaluation setups; we agree that there are many more questions to be studied, and we hope that our open-source data and code will facilitate such work in the future.\n\nOn the specific comments: \n\n(1) It is hard to scale single data sources like Wikipedia to the trillion token regime, so we focus on multi-domain web-scale data; however, we compare with single in-domain data sources in Table 3.  \n\n(2) We agree it would be good to study more downstream tasks beyond QA, and we will emphasize this in the discussion. We note that it is common in prior work (RePlug, Shi, et al; Atlas, Izacard, et al.) to focus on QA as well, and that we are the first to study downstream tasks at large datastore scales.\n\n(3) We have added results on new LMs (Llama3, Olmo, and Pythia families) and new retrievers (DRAGON and GTR) in Figure R1 and Table R1 in the general reply, respectively; these results are consistent with our submission. \n\n---\n\n**Weakness 2.** There is a lack of comparative analysis and comparison of MASSIVEDS with other existing datastores [1, 2] (Dolma and Pile).\n\n**A2.** We would like to clarify that a datastore, an indexed resource ready for retrieval, differs from a pretraining corpus, which contains only raw data, like Dolma and Pile. Constructing a datastore from a pretraining corpus involves data cleaning, chunking, embedding, and indexing, which take substantial effort, and as such there are no comparable open-source datastores to MassiveDS.\n\nThat said, the reviewer raises an interesting point about comparing different datastore sources. We ran an additional experiment using datastores constructed from DCLM-baseline and FineWeb-Edu (which have been shown to perform better than Pile and Dolma). As Table R2 in the general reply shows, MassiveDS achieves comparable performance.\n\n---\n\n**Weakness 3.** The claim that indexing the entire datastore initially reduces computational overhead is debatable. This aspect could be alleviated significantly by alternative indexing strategies [3, 4].\n\n**A3.** The reviewer is correct that embedding should be done once to optimize efficiency; this is also the core of our pipeline. However, there are additional challenges: e.g., the retrieval process is I/O bound and often takes significant time to load passages from a large document pool. Also, using existing software like FAISS requires building separate indices for different sets of dense embeddings, significantly increasing storage demands for various configurations. Our open-source pipeline provides an efficient, ready-to-use solution that overcomes these challenges and that can be used as a foundation for future research. We appreciate the reviewer’s point and will expand our discussion of this in the paper. \n\nThe alternative indexing strategies cited by the reviewer do not fully overcome the challenges: [3] speeds up the search process but doesn't lower the cost of index construction; [4] improves datastore construction efficiency, but needs additional training and doesn't avoid redundant storage for different configurations. We will discuss them in our next revision.\n\n---\n\n**Weakness 4 and Q3.** The overlap between the datastore and the downstream evaluation dataset raises concerns about data leakage.\n\n**A4.** We agree that data leakage is a potential concern. In Appendix B.1, we describe the strategies we used to avoid it. Specifically, we applied a strict data decontamination method to remove documents that overlap with the evaluation samples, which is stricter than common practice such as Dolma and RedPajama. In Figure 2(b), we also assessed perplexity (PPL) scaling on M2D2, which is not part of MassiveDS, and found scaling curves similar to those using decontaminated RPJ.\n\n---\n\n**Weakness 5.** The author should reevaluate the claim with additional models.\n\n**A5.** We thank the reviewer for suggesting this. We have added more results with different LMs in Figure R1 in our general reply and will include them in our paper.\n\n---\n\n**Q1.** How few-shot exemplars are selected for QA tasks?\n\n**A6.** Following the popular evaluation repository LM-Evaluation-Harness, we use few-shot examples randomly sampled from the development set, which do not overlap with test samples.\n\n---\n\n**Q2.** How is the language model configuration determined? Is there any inherent randomness?\n\n**A7.** We follow LM-Evaluation-Harness, which uses greedy generation for TriviaQA and NQ, and log-likelihood for MMLU and MedQA. Therefore, there is no randomness in the inference. We’ll clarify this in the paper.\n\n---\n\n**Q3.** (see Weakness 4)\n\n---\n\n**Q4.** How do you substantiate that the contribution of MASSIVEDS is significant and not merely an assembly of various data segments?\n\n**A8.** Our primary contribution is enabling the scientific study of >1T-token RAG scaling on an academic budget, which extends beyond merely assembling the datastore and differs from previous systems that aimed to enhance search speed. Our identification of the commutability of different operations and the strategic ordering of these operations are crucial for facilitating accessible datastore scaling studies. Beyond this, we highlight our other contributions in presenting the first datastore scaling trends and analyses on various datastore design factors.\n\nWe thank the reviewer for their valuable suggestions and we’ll update our draft accordingly.", "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:
We appreciate the reviewers' strong support for the contributions of the paper and their insightful comments. This general response outlines how we have responded to their concerns and provides the requested supplementary results.

**Summary of common concerns and our corresponding response.**
* Reviewer F6AB and Reviewer TsAT asked to compare our proposed pipeline with existing retrieval systems for efficient search at inference time, such as DiskANN. We would like to clarify that our pipeline is optimized for a different goal and is orthogonal to the existing systems.
    * Our pipeline is designed to facilitate accessible datastore scaling studies, where the key challenge is how to efficiently study datastores of different configurations, such as datastore sizes, quality filters, data decontamination methods, etc. While prior work such as DiskANN focuses on search efficiency at inference time given a fixed configuration, we believe these are complementary to our work. Specifically, even with the state-of-the-art nearest neighbor search used in production, studying the effect of various factors of the datastore remains expensive. This is because such experiments require rebuilding search indexes for every combination of the factors, orthogonal to what search algorithm is being used. Our work focuses on removing the need for such repetitive rebuilding of indexes. This allows us to conduct  a comprehensive study of the impact of various datastore factors, which is a novel contribution of our work.

* Reviewer TsAT and Reviewer AKfc asked about comparisons of MassiveDS with other datasets such as Dolma. We would like to clarify that the definition of a datastore is different from a pretraining dataset.
    * By “datastore”, we refer to an index that is ready for retrieval; a pretraining corpus only contains raw data. Building a datastore on an existing pretraining corpus requires additional data chunking, embedding, and indexing over the raw data, which is a non-trivial effort. It has been challenging to conduct datastore scaling research at this scale before and the existing open-sourced datastores are much smaller than MassiveDS, as shown in Table 1 of the submission. Prior larger scale pre-training data is only available in the format of raw text, while we open source both the raw text as well as the resulting embedding and index to search from the entire 1.4 trillion corpora. In addition, we open source the codebase to use our released index for future study.
    * Inspired by the reviewer, we further compare the performance of datastores built with different sources in Table R2 below, which hasn’t been examined by any prior work. See the response below for more details.

* Requested supplementary results.

    * As requested by Reviewer omwP and Reviewer AKfc, we supplement more scaling results with more language models, such as Llama3, Olmo, and Pythia models, in Figure R1. The new results show our conclusions hold across different LMs.

    * As requested by Reviewer cEKQ, Reviewer TsAT, Reviewer vyTD, and Reviewer AKfc, we supplement the evaluation results using 3 different base retrievers in Table R1. The results indicate that retrievers of similar sizes have comparable performance on general-web data and that Contriever-MSMACRO is a reasonable pick for our main experiments.

    * As requested by Reviewer AKfc, we compare the performance of datastores constructed using different data sources (DCLM-baseline and FineWeb-Edu) in Table R2. The results indicate that MassiveDS matches or even outperforms the datastores constructed with the latest high-quality pretraining data sources.

* Code: We uploaded our code to an anonymized repository and have sent the link to the AC.

 We will go into more detail about each of these in the relevant individual responses below.


**Table R1. Results of different base retrievers on 10% randomly sampled MassiveDS evaluated with Llama2-7B.**
| Name       | Retriever Type | Size  | Perplexity ↓ | Natural Questions ↑ | MMLU ↑ |
|------------|----------------|-------|--------------|----------------------|--------|
| Contriever-MSMACRO | dense          | 177M  | 4.221       | 33.2               | 49.2 |
| DRAGON [1]     | dense          | 110M  | 4.237       | 33.9           | 48.7 |
| GTR-Base [2]   | dense          | 110M  | 4.214   | 30.8               | 49.3 |

[1] Lin, Sheng-Chieh, et al. "How to train your dragon: Diverse augmentation towards generalizable dense retrieval."

[2] Ni, Jianmo, et al. "Large dual encoders are generalizable retrievers."

**Table R2. Results of different data sources for the datastore evaluated with Llama2-7B.**
| Source         | #Tokens (B) | NQ    | MMLU  |
|----------------|-------------|-------|-------|
| LM-only        | /           | 26.6  | 45.8  |
| DCLM-baseline [3]  | 100         | 31.4 | 49.0    |
| FineWeb-Edu [4]    | 100         | 30.4 | 49.4 |
| MassiveDS (ours)     | 100         | 33.2  | 48.8  |

[3] Li, Jeffrey, et al. "DataComp-LM: In search of the next generation of training sets for language models."

[4] Penedo, Guilherme, et al. "The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale."

**Figure R1. (See PDF)  Scaling performance on downstream tasks with different language models, corresponding to Figure 2 (c)-(f) with additional models included. The trends remain consistent.**

Meta-review: This work considers the problem of scaling data used at inference time in retrieval based LLMs. It presents MASSIVEDS, a 
1.4 trillion-token datastore along with a pipeline to study efficiently the impact of different datastore features, such as data size, data filters, and decontamination strategies on language modeling as well as a few retrieval-based tasks that consume data retrieved from the datastore. It does a systematic study of the effect of scaling data on the performance and finds that scaling produces loglinear improvement without saturation. 

Reviewers agree that scaling retrieval-based language models is an interesting problem, the paper is written well, experiments are systematic and demonstrate the value of scaling, code and data store can be potentially useful in future research in this field when open-sourced. Reviewers also found that some of the insights from the experiments are useful.

Reviewers expressed concern on the choice of RAG models used in the study, whether scaling would benefit more powerful RAG models, impact of scaling on efficiency of retrieval, lack of diversity in the tasks included in the scaling study, larger and more recent retrievers not considered, interplay between different types of retrievers and language models not studied, lack of comparative analysis of other open-source datastores.

I agree with the authors and the reviewers that the work will facilitate datastore scaling studies for retrieval-based tasks.
Decision: Accept (poster)