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_cc_2024_Conference
Published: 2024-05-15 18:12:35

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.

Main paper: 
 introduction: 
 The scaling of large language models (LMs) has driven tremendous performance gains across a 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 et al., 2024) . In this paper, we consider another dimension of scaling: the amount of data in a datastore used at inference time by retrieval-based LMs, which can directly retrieve information from the datastore to use in context when generating output (Karpukhin et al., 2020; Guu et al., 2020; Izacard & Grave, 2020; Asai et al., 2024b) .Retrieval-based LMs have a range of benefits such as improved factuality (Mallen et al., 2023) , effective domain adaptation (Khandelwal et al., 2020) , credit attribution (Gao et al., 2023) , and parametric efficiency (Min et al., 2023b) . However, most prior work in retrieval-based LMs use datastores constructed from a single data source (Karpukhin et al., 2020) , such as Wikipedia, with sizes on the order of a few billion tokens. While there has been some work on larger datastores (Table 1 ), with the largest being RETRO (Borgeaud et al., 2022; Wang et al., 2024) in the trillion-token range, these studies use proprietary datastores and custom architectures with a limited evaluation suite. As such, it remains unknown how datastore scaling helps the currently dominant retrieval-in-context approaches on a broad categories of tasks.We first construct MASSIVEDS, a massively 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 generalknowledge question answering (QA), domain-specialized QA, and reasoning tasks. We find that, first, datastore scaling consistently improves both language modeling and some downstream tasks in a task-dependent manner (Figure 1 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-based LMs enable better compute-optimal scaling trends, where they achieve superior performance than LM-only models at the same training cost (Figure 1 Right).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 efficiently studying datastore scaling) at https://github.com/RulinShao/retrieval-scaling. 
related work: 
 Retrieval-based LMs. Unlike parametric LMs that only use data during training, retrieval-based LMs can access data through a datastore during inference (see Asai et al. (2024b) for a review). We focus on retrieve-in-context language models (RIC-LMs), which retrieves a small set of documents from the datastore and feeds a concatenation of them as an input to the LM (Ram et al., 2023; Shi et 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-based LMs often focused on specific aspects of LMs such as factuality and attribution. In addition, they typically use limited-size, single-domain Table 1 : Comparison with prior work, ordered by datastore size. '# Tokens' indicates the number of tokens in the datastore using the LLAMA2 tokenizer (Touvron et al., 2023) . The asterisk (*) denotes that the datastore is not evaluated on downstream tasks. MASSIVEDS is the largest open-sourced datastore and covers a broad spectrum of domains.# Tokens Data Sources Open sourced ATLAS (Izacard et al., 2023) <5B Wikipedia ✗ REALM (Guu et al., 2020) <5B Wikipedia ✗ RALM (Ram et al., 2023) <5B Wikipedia ✓ SELF-RAG (Asai et al., 2024a) <5B Wikipedia ✓ REPLUG (Shi et al., 2023) 47B The Pile ✓ RA-DIT (Lin et al., 2024) 79B Wikipedia, CommonCrawl ✗ SPHERE (Piktus et al., 2022) 90B CCNet ✓ RETRO++ (Wang et al., 2024) 330B* The Pile, CommonCrawl, RealNews, CC-Stories ✗ INSTRUCTRETRO (Wang et al., 2024) 1.2T* Wikipedia, CommonCrawl, RealNews, CC-Stories, Books ✗ RETRO (Borgeaud et al., 2022) 1.7T* MassiveText (Rae et al., 2022 ) ✗1.4T 8 domains, listed in Table 2 ✓ datastores such as Wikipedia, which is on the order of a few billion tokens (Table 1 ). Scaling the datastore remains relatively underexplored, with two notable exceptions. First, Borgeaud 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., 2024) , use this trillion-token datastore solely for language modeling evaluation, while using a small task-specific Wikipedia datastore for downstream task evaluation. Since RETRO-based datastores are not fully open-sourced, it is challenging to 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 et al., 2020) . However, their evaluation on downstream tasks such as KILT (Petroni et al., 2021) suggests that SPHERE does not always outperform a small, in-domain datastore like Wikipedia.In contrast, we present the first study on the downstream performance of trillion-token scale datastores, including an analysis of compute-optimal scaling trends using retrieval-based LMs with different sizes of datastores, models, and pretraining corpora. Our work is fully open-source and can be replicated on a limited computational budget, enabling research on trillion-token datastores to be more broadly accessible.3 MASSIVEDS and our Datastore Scaling Pipeline MASSIVEDS encompasses both domainspecific data and general web data (Table 2 ). Domain-specific data comes from specialized sources, and tends to be smaller but higher in quality. We cull from a mix of datarich domains: books which span a variety of genres (Computer, 2023) ; open-access scientific papers (Lo et al., 2020; Soldaini & Lo, 2023; Computer, 2023) ; encyclopedic articles (Karpukhin et al., 2020; Computer, 2023) ; community questions and answers from StackExchange (Computer, 2023); code from GitHub (Computer, 2023); mathematical webpages (Paster et al., 2023) and mathematical language (Welleck et al., 2021) ; biomedical articles (of Medicine, 2023). On the other hand, general web data is sourced from Common Crawl snapshots at five time periods (07/2019, 05/2020, 04/2021, 05/2022, 06/2023) and C4 (Raffel et al., 2020) .  
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 deduplication, 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 query1 ; 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 2 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 2 The key idea is to reorder the above operations such that the most expensive ones-indexing 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-are performed as late as 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 to 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 to 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.1.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 raw data pool by the subsampling fraction p. A 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 as 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.e., if a p f fraction of the data is filtered out (0 < p f ≤ 1), then the actual number of tokens should also be multiplied by p f . Since we plot datastore size on a log axis, this corresponds to a 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) , we use CONTRIEVER-MSMARCO (Izacard et al., 2022) , which represents every document in the datastore as a dense vector, as our retriever. We ablate the choice of retriever in Appendix E.1; we found that CONTRIEVER-MSMARCO 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 performance 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-knowledge, 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 human-annotated answers from Wikipedia. Massive Multitask Language Understanding (MMLU; Hendrycks et al. 2021 ) comprises general-purpose, multi-task 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 k = 3 documents in reverse order, so that higher-ranked documents are 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 Section 4; we study the effect of rerankers in Section 5.2. More details, including decontamination measures, are in Appendix B. 
datastore scaling results on language modeling and downstream tasks: 
 Finding 1: Datastore scaling significantly helps language modeling. Figures 3(a ) 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 to yield additionalFigure 4 : Compute-optimal scaling curves for retrieval-based and LM-only 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 a datastore scaling curve of lined dots similar to the ones in Figure 3 . The Pareto-optimal points are highlighted in red for retrieval-based LMs and blue for LM-only. Within a fixed computational budget (represented on the x-axis), retrieval-based 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 abilities. 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 LLAMA-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 is task-dependent. Figure 3(c)-(f ) show the performance on four downstream tasks as a function of datastore size. Datastore scaling brings major improvements to knowledgeintensive 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 scale. 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 on TQA and NQ, indicating the effectiveness of storing knowledge in the datastore.On MMLU, a multi-subject, reasoning-heavy 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.Next, we study performance as a function of total training-time compute and show that retrieval-based LMs achieve superior compute-optimal 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 B.4. 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 a 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-time 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 performance on downstream tasks in Figure 4 . The Pareto-optimal points for retrieval-based and LM-only 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, retrieval-based LMs achieves superior performance than LMonly models, indicating offloading FLOPs from pretraining to datastore construction can result in better performance. Therefore, we conjecture that storing factual knowledge in a 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 knowledgeintensive tasks that measure factual recall. Surprisingly, we find that retrieval-based PYTHIA models (trained on up to 300B tokens) and OLMO-1.7 models (trained on up to 2T tokens3 ) have a 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 is obtained early in training.Finding 5: Retrieval shows benefits for reasoning-intensive 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 on more and better data, consistently benefits from retrieval on 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-heavy 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.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 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 in bold and the second best is underlined. We show the diverse domain coverage in MASSIVEDS consistently improve the performance across tasks. some tasks. We also note that there is emerging work on accelerating retrieval search and designing efficient serving strategies for retrieval-based LMs, such as Cao et al. (2023) . We leave a study of inference-compute-optimal scaling to future work. 
effects of data composition: 
 Finding 6: MASSIVEDS matches or outperforms single-domain datastores. The default setting in prior work is to use a single-domain datastore that is in-distribution to the downstream task.In practice, however, it is often difficult to determine and curate a datastore that is perfectly indistribution for a downstream task, and even if we can, it limits the generality of the retrieval-based model to that task. In Table 3 , we compare MASSIVEDS with single-domain datastores. MASSIVEDS significantly outperforms these in-domain datastores on language modeling, as well as TQA and MMLU, while matching performance on NQ and MedQA. 4 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, it retrieves relatively 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 out-of-domain data in the datastore; this aligns with similar findings on kNN-LM (Khandelwal et al., 2020) , another type of retrieval-based LM, in Shao et al. (2023) .Overall, these results show that retrieving from broad datastores like MASSIVEDS can simultaneously improve performance across multiple domains, paving the path towards generalpurpose retrieval-based models.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 a more computationally expensive model (Ram 2% of the answer length) continuous overlap or an 80%+ Jaccard similarity with the answer. We find decontamination impacts the language modeling performance a lot but not the downstream task. et al., 2023; Sachan et al., 2022) , and take the final top-3 reranked documents for evaluation. Specifically, we use a cross-encoder reranker, which encodes a concatenation of a query and document and returns a similarity score (Nogueira & Cho, 2019) . We choose MINI-LM-L12 V2, a BERT-based cross-encoder5 that is trained for passage ranking, following Izacard et al. (2022) . Additionally, we use a lexical oracle reranker, which uses the gold answer, as an upper bound on the potential benefit realizable by a better reranker or retriever for knowledge-intensive question-answering. The oracle reranker scores each document based on whether the gold answer is included in the document and if not, the fraction of unigram overlap between the document and the answer.Figure 6 reports scaling trends on TQA and NQ of RIC-LM with Llama2-7B. While the crossencoder-based reranker improves performance on TQA and NQ, a notable gap persists between the oracle reranker and the cross-encoder-based reranker. These suggest that improving either retrieval or reranking can further boost the scaling performance of retrieval datastores. Improving the retriever for more reasoning-heavy tasks such as MMLU and MedQA remains an open problem (BehnamGhader et al., 2022 ) that we leave to future work.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., 2022) . By default (Section 4), we perform decontamination by filtering documents with 80+% 13-gram 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-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 7 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-where no more than 8 continuous words overlap-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 at larger scales.Data quality filtering. In Appendix E.2, we study the impact of data quality filtering on MAS-SIVEDS, where we consider global data deduplication and a 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, subsampling with higher p increases the chance of seeing more duplicates. In addition, we observed that DOLMA quality filters have a 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. 
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 to 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 high-quality data for improving performance on 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 question-answering tasks whose outputs are either predefined choices or short form generations. We defer the evaluation on more tasks such as long-form generation and mathematical reasoning to future work.Despite these limitations, our research shows that increasing the scale of data available at inference time can improve model performance, at lower training cost, on language modeling and a variety of downstream tasks. We expect that future work on improving retrieval-based models with large-scale 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 a 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 to retrieve helpful information from a large datastore).Fangyuan Xu, Weijia Shi, and Eunsol Choi. Recomp: Improving retrieval-augmented lms with compression and selective augmentation, 2023.The MASSIVEDS pipeline entails running the following operations in order (Figure 2 ):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, we 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 a target combination of domains in the datastore, we merge the top-K documents from the target domains for that query. This leads to a merged pool of DK retrieved documents, where D is the number of target domains. From this pool, we re-select the top-K documents.4. Data filtering and reranking. We then apply the data filtering and reranking steps, as described above, to only the top-K 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.5. 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 in Appendix A.5. Furthermore, we collect the top-K ′ documents (where K ′ < K) from the deduplicated and decontaminated top-K set for reranking. We set K ′ = k when reranking is not applied. Next, we describe each step of our efficiency-oriented datastore construction pipeline in detail and demonstrate the equivalence between our MASSIVEDS pipeline and the naive pipeline. Below is a table of notations for easy reference. Our pipeline starts with raw data that is split into fixed-size documents. For each document, we run one forward pass over it with a retriever model ϕ and save the final-layer representation as an embedding for that document. We store the embeddings of all documents in the datastore for similarity-based 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 is 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 a single query in the sections below.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 innerproduct 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-K documents for subsequent 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 it 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 top-K 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 = {d 1 , • • • , d N }.Given a query q, element-wise top-K retrieval uses a retriever g ϕ to compute the similarity score s i = g ϕ (q, d i ) of the query and each document d i independently, and then returns the documents with the top-K highest retrieval scores.Definition A.2 (m-shard distributed element-wise top-K retrieval). Consider a sharding strategy that splits the datastore intom shards {D 1 , • • • , D m }, such that each shard contains a disjoint subset of documents D (i.e., D 1 ∪ D 2 ∪ • • • ∪ D m = D; D i ∩ D j = ∅, when i ̸ = j).In m-shard distributed element-wise retrieval, we fetch the top-K documents from each shard in parallel (if there are fewer than K documents in a shard, then all documents in the shard are returned). Then, we merge these mK documents and return the top-K highest-scoring documents from the merged pool.m-shard distributed element-wise top-K retrieval is equivalent to retrieving from a single unsharded index (m = 1).Lemma A.1. Changing m, the number of shards used in distributed element-wise top-K retrieval, does not impact the final retrieved results.Proof. Let the top-K documents obtained by the element-wise top-K retrieval and m-shard distributed element-wise top-K retrieval be D K and D ′ K , respectively. Since a document that is ranked in the top-K across all documents must be ranked in the top-K within any individual shard, we haveD K ⊆ D ′ K . Because {D 1 , • • • , D m }is a disjoint union of D, for any document d in D ′ K , there are no more than K -1 documents in the m shards, i.e., all documents, that have higher scores than d. Then we have d ∈ D K . Therefore, D ′ K ⊆ D K . Finally, we have D K = D ′ K for any choice of m ∈ N + .In our experiments, we use an indexing method6 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 is equivalent to building and retrieving from a single non-distributed index.Our pipeline includes two basic data processing steps7 : data deduplication and decontamination. These steps are applied post-hoc on the retrieved documents. Reranking, which is optional, is used to enhance the quality of retrieval. Detailed descriptions of deduplication, decontamination, and optional reranking are provided in this section.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.There are two widely adopted approaches for deduplication over a large-scale dataset: the first uses a Bloom filter8 and the second uses MinHash9 . 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.4 trillion tokens using a Bloom filter requires 1.5TB RAM memory and 933 CPU hours10 . 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 a post-hoc MinHash deduplication step which is applied to a large pool of top-K retrieved documents. Following Computer (2023), a document pair is a 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 k 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.Expanding the datastore to the Internet-scale incurs the risk of data contamination. Similar to global deduplication, running pre-decontamination on 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 8 ), 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-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 taking the top-K ′ (where K ′ ≤ K) documents is equivalent to retrieving the top-K ′ documents from a deduplicated and decontaminated datastore.Proof. Given a datastore D, let the deduplicated datastore as D ′ . Denote the top-K documents from D as D K and the top-K ′ documents D ′ as D ′ K ′ . Then, denote the top-K ′ documents retrieved from D K as D K ′ . 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 deterministic, we haveD ′ 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-)duplicates of each other will still hold.Reranking is optionally applied to enhance retrieval quality in our pipeline. Given the top-K retrieved documents, we apply a stronger retrieval model to embed each document and the query, Algorithm 1 Naive implementation of datastore scaling 1: Input: Data corpus D N , a list of subsampling ratios P , a list of random seeds S, a list of queries for q ∈ Q do 10:Q 2: Output: Retrieved top-k documents [D (q,p,s) k | q ∈ Q, p ∈ P, s ∈ S]D (q,p,s) k ← SEARCH(q, D s p|D| , k) ◁ Complexity: O(pN ) 11:end for 12:end for 13:end for 14:return [D (q,p,s) k| q ∈ Q, p ∈ P, s ∈ S] 15: end function Algorithm 2 Our efficient datastore scaling implementation 1: Input: Data corpus D N , a 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 (q,p,s) k | q ∈ Q, p ∈ P, s ∈ S]{D q K | q ∈ Q} ← SEARCH(Q, D s piN , K) ◁ Complexity: O(N |Q|) 7:for s ∈ S do 8:for p ∈ P do 9:for q ∈ Q do Lemma A.3. Subsampling from the retrieved top-K documents with probability p and then taking the top-k (k ≪ K) from the subsampled documents (Algorithm 2) is equivalent with high probability to directly retrieving top-k documents from a datastore that is subsampled from the entire raw text datastore with probability p (Algorithm 1). The equivalence holds as long as there are enough k 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 = 8 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 a data pool. Following Borgeaud et al. (2022) ; Magnusson et al. (2023) , 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 a combination of two decontamination 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 an 512-token answer to compute perplexity, so the decontamination ratio is 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.Hyper-parameters. Our pipeline has three main hyper-parameters: k, K, and p. k is the number of documents used for evaluation. K is the number of documents retrieved before subsampling. p is the subsampling ratio which controls the size of the datastore. We consider k = 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 a 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 on the raw data first based on Proposition A.3. Table 4 : Lookup table of the tail bound for a binomial distribution Binomial(K, p) with at least m = 3 number of successes. p = 0.01 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.000Following Baevski & Auli (2019) ; Khandelwal et al. (2020) ; Min et al. (2023a) , we split evaluation data into into fixed-length chunks of 1,024 tokens, with a stride of 512 tokens. For each chunk, the first half is used as a prefix and retrieval query, and the second half as the target sequence for calculating perplexity.Table 5 shows the domain, metric, and sample count for each downstream task. We evaluate all downstream tasks in a 5-shot setting and prepend the top-3 documents for retrieval-based LM evaluation. We adapt the lm-evaluation-harness,13 a widely used LM evaluation suite, for downstream evaluation. (Kwiatkowski et al., 2019) Open-domain QA Wikipedia Exact Match 3,610 MEDQA (Jin et al., 2020) Science QA Medical Accuracy 1,273 MMLU (Hendrycks et al., 2020) 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 to a superior performance than the other way. Therefore, we stick to the second format through the paper.Intermediate checkpoints. Thanks to 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 et al., 2023) and OLMo (Groeneveld et al., 2024) . Pythia provides checkpoints of 9 sizes trained on up to 300B tokens from the Pile (Gao et al., 2020) . We consider Pythia models that have at least 1B parameters, i.e., PYTHIA-1B, PYTHIA-2.8B, PYTHIA-6.9B, and PYTHIA-12B, in favor of their capability of handling complex downstream tasks. Additionally, we consider OLMO-1.7-1B and OLMO-1.7-7B 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, 1/4, 1/3, 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; Hoffmann et al., 2022; Gadre et al., 2024) , we approximate the FLOPs needed for one forward pass as FLOPs fwd ≈ 2N D and for a backward as FLOPs bwd ≈ 4N D, 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 N LM and the size of the pretraining corpus as D pretrain . The FLOPs for pretraining can be approximated as FLOPs pretrain ≈ 6N LM D pretrain . Datastore construction requires one forward pass on every token in the datastore corpus during the embedding step. Denote the size of the retriever as N retriever and the size of the datastore corpus as D datastore . The FLOPs for datastore construction can be approximated as FLOPsdatastore ≈ 2N retrieverD datastore . Because we use a 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.In this section, we supplement the datastore scaling performance of PYTHIA and OLMO models, in addition to the LLAMAmodels shown in Figure 3 . Specifically, we present the datastore scaling performance for the following model families:• Pythia (Biderman et al., 2023) of 4 sizes: PYTHIA-1B, PYTHIA-2.8B, PYTHIA-6.9B, and PYTHIA-12B. • OLMo (Groeneveld et al., 2024) of 2 sizes: OLMO-1B and OLMO-7B.• Llama (Touvron et al., 2023) , which we consider LLAMA-2 7B, LLAMA-2 13B, and LLAMA-3 8B.The complete datastore scaling results for TriviaQA, Natural Questions, MMLU, and MedQA are shown in Figures 9, 10, 11, and 12 , 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 same model architectures when augmented with MASSIVEDS. Surprisingly, Pythia-1B matches Pythia-12B when augmented with only 100B tokens on knowledge-intensive tasks such as TriviaQA and Natural Questions, and it outperfoms Pythia-12B when further increasing the size of datastore. We show in Figure 3 that LLAMA-3 8B has worse PPL scores than LLAMA-2 7B on RedPajama data with and without retrieval augmentation, which contradicts the intuition that LLAMA-3 8B 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 a post-training process using instruction-tuning data, which aims to enhance model alignment and output quality for complex tasks but could shift performance metrics away from those optimal for simple PPL evaluations. In addition, LLAMA-3 was trained with significantly more data-a domain (such as RedPajama) could be down-weighted in a larger corpus, leading to less memorization for this certain domain during pretraining.In this section, we provide additional analyses of the impact of retriever and data quality filtering.In this section, we ablate the choice of retriever. Since datastore construction is expensive, we subsample 10% from the full corpus of MASSIVEDS for this ablation study. We build datastores with 3 different retrievers, CONTRIEVER-MSMARCO (Izacard et al., 2022) , DRAGON-ROBERTA (Lin et al., 2023) , GTR-T5-BASE (Ni et al., 2021) , and evaluate on upstream perplexity and downstream tasks with them separately. A subset of the tasks, i.e., RedPjama for language modeling evaluation, and NQ and MMLU for downstream evaluation, is used to compare different retrievers. The evaluation results are shown in Table 6 . We find the 3 retrievers perform similarly. We empirically find the implementations of Contriever14 and DRAGON15 run much faster than the sentence-transformer16 implementations, e.g., GTR-Base17 . As a result, we choose Contriever in our full-size scaling study with a consideration of both performance and efficiency.Setup. We consider three data quality filters applied to the DOLMA (Soldaini et al., 2024) corpus:(1) a whitespace filter which counts the number of whitespace-separated tokens in each document, and filters out documents with counts under a manually defined threshold; (2) a language filter which uses a FastText (Bojanowski et al., 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 a span of all punctuation characters (the length of the span is user-defined).Data deduplication. Removing data duplicates has proven effective for pretraining more computeoptimal language models (Lee et al., 2022) . Duplicates are especially undesirable in 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, we 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. OnFigure 13 : Analysis on the effect by deduplication, decontamination, and quality filters, from left to right (Section 5.3). 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 p increases the chance of seeing more duplicates 18 .DOLMA Quality filtering. To study the impact of quality filtering, we consider a data filter adapted from DOLMA (Soldaini et al., 2024) which uses a combination of three filters: a whitespace filter; a language filter, and an 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 14 .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 a degradation in end-task performance. Figure 15 indicates that removing short chunks can significantly improve the retrieval-based LM performance with a large datastore, which is more likely to contain short, but meaningless word chunks.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 a short-chunk removal step is added to our pipeline. The retriever tends to retrieve a short chunk that may have high lexical overlap with the question, but does not provide any useful information for the answer. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: Yes, we discuss broader impacts of our work in ??.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We do not believe this question is applicable: as 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 no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: Yes. We acknowledge the codebases we have referred to, cite and provide the urls of datasets and models in section 3.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [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 new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: We do not involve any crowdsourcing or 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 in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: 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 to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.Reranking is optionally applied, which reranks the retrieved documents based on a more effective but usually more expensive heuristic(Nogueira & Cho,2019;Sachan et al., 2022).2 The operations involved in the naive pipeline are expensive in terms of FLOPs, I/O operations, and storage due to repetitive index construction and large-scale retrieval.1B is trained on up to 3T tokens and OLMO-1.7 7B is trained on up to 2T tokens.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.https://www.sbert.net/docs/pretrained-models/ce-msmarco.htmlWe use INDEXFLATIP implemented in FAISS: https://github.com/facebookresearch/faiss/ wiki/Faiss-indexes.Additional data quality filters, such as those from DolmaSoldaini et al. (2024), are discussed in our analysis section §5.3. These filters are applied at the same stage as deduplication.A typical implementation is https://github.com/allenai/bff/tree/main.A typical implementation is https://ekzhu.com/datasketch/.The requirement of Bloom Filter deduplication is estimated using the implementation by AI2: https: //github.com/allenai/bff/tree/main on our hardware setup.https://github.com/EleutherAI/lm-evaluation-harnesshttps://github.com/facebookresearch/contrieverhttps://huggingface.co/facebook/dragon-roberta-query-encoderhttps://sbert.net/https://huggingface.co/sentence-transformers/gtr-t5-baseThe source corpus of MASSIVEDS has been applied moderate deduplication. For example, RedPajama preprocessed the data using several filters: https://github.com/togethercomputer/RedPajama-Data/ tree/rp_v1/data_prep. However, we still find many duplicates in the retrieved results. 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- 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.


[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.

[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.

[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.
questions: 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?
 - 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?
- 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: limitations: n.a.
paper_summary: paper_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.
strength_weakness: strength_weakness: The following are some of the limitations of the paper (some of these have already been highlighted by the authors):
- **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.
- **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.
questions: 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: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: Overall, I do not find any significant weaknesses in the paper. The only minor issues I noticed are:

(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.

(2) Typo: L223 "Datascore" should be "Datastore".
questions: questions: Please see my comments in Weakness.
limitations: limitations: The limitations are properly stated in the Conclusion and Limitations section.
paper_summary: paper_summary: This paper studies the impact of scaling the datastore (retrieval dataset) on retrieval-based language models.

The contributions are:
- MASSIVEDS a 1.4 trillion-token datastore for retrieval-based LMs that will be made open-source.
- A pipeline to study the impact of the datastore scaling on the language models at inference time
- An analysis of the impact of scaling and other pre-processing steps (decontamination, deduplication, and quality filtering) on the language models at inference time
strength_weakness: strength_weakness: INFERENCE TIME ANALYSIS

This 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.
questions: questions: 1) CODE NOT AVAILABLE?

Unless 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.

2) RETRIEVAL AUGMENTED GENERATION?

Since the LMs used in the experiment are all generative language models, why is the term Retrieval Augmented Generation not used to describe this work? ]

3) DATA DECONTAMINATION/DEDUPLICATION

Can 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?

4) FIGURE 4 (small detail):

Is there a reason why the baselines LM-only on the subplots (d), (e), and (f) are not extended with a dashed line?
limitations: limitations: The authors discuss several limitations of their work:
- Focus on a particular class of retrieval-based LMs (RIC-LMs), 
- Do not study very large LMs (only 1B and 7B parameters)
- Only consider dense retrievers. 
- Only considers QA tasks

There 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).
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

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?

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.
questions: questions: 1. What is the recall rate of the retriever for different datastore sizes?
limitations: limitations: The limitations section is included and adequately discussed.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

2. 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. 

3. There are several major assertions may be incorrect or lack sufficient justification. For instance:
    - 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].
    - 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).
    - 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.

Typos:
- line193 into into -> into
- line726: exmaples -> examples


[1] Soldaini, Luca, et al. "Dolma: An open corpus of three trillion tokens for language model pretraining research" 

[2] Gao, Leo, et al. "The pile: An 800gb dataset of diverse text for language modeling" 

[3] Li, Haitao, et al. "Constructing tree-based index for efficient and effective dense retrieval"

[4] Zhou, Jiawei, et al. "Semi-Parametric Retrieval via Binary Token Index"
questions: questions: 1. Could you elaborate on how few-shot exemplars are selected for QA tasks?

2. 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? 

3. 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?

4. 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?

5. 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: 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.

Comments: title: Reviewer-Author Discussions
Response to Authors
Response to Rebuttal from Authors
Response to the follow-up questions
Response to Authors
comment: Dear Reviewer,

As the discussion deadline approaches, we want to ensure that all your concerns and questions have been fully addressed. To address the weaknesses listed in the original review, we have clarified our contributions and the detailed settings of our experiments. We will revise our paper accordingly.

We deeply appreciate your insights and the time you have invested in reviewing our work. Your feedback is invaluable in refining our research and ensuring its quality and relevance. Please let us know if you have any further questions!
Thank the authors for further explanations. I do not have other concerns and will keep my original positive rating.
We thank the reviewer for their insightful suggestions, which are very valuable in helping us further improve the quality of our paper. We also appreciate the reviewer's raised score. We will definitely include these discussions in our next revised version. Thank you for taking the time to review our manuscript!
Dear reviewers: as you are aware, reviewer-author discussions phase is ending on Aug 13. We request you to kindly make use of the remaining time to contribute productively to these discussions. If you have not read and/or responded to author rebuttal, please do it asap so that the authors get a chance to respond to you. If you have more questions to ask or want further clarification from the authors, please feel free to do it. 

Dear authors: please read the responses of the reviewers and if there are any pending questions/request-for-clarification, please attend to them.
Dear Authors,

Thank you for addressing my questions. A1 has resolved most of my concerns regarding the generalization of the scaling trend, and I recommend incorporating it into the main paper, as it is well justified. Regarding W3, in my initial reading, I assumed the computational overhead was related to GPU resources. It would be helpful to clarify this point as discussed in previous A3. Additionally, it would be beneficial to discuss whether these computational challenges are theoretically inherent to the method or if they stem from the choice of index type and third-party libraries. I have made the necessary changes to the score based on our discussion.
Dear Authors, Thanks for responding to my comments and also including the supplementary results. The response and the proposed revisions to the next version adequately addresses the concerns I had related to some of the weakness of this work. After going through their responses, I have no other major concerns and due to this I am revising my rating.
We thank the reviewer for raising the score, which encourages us. Here are some discussions on the follow-up questions:

**Q1.** What are the critical factors to consider when expanding a datastore with diverse data sources to ensure performance improvements in RAG?

**A1.** This is a great question. In Section 5.2, we  visualize the source distribution of retrieved results in Figure 3 and compare the performance of the multi-domain MassiveDS and single-domain datastores (Table 3). Our results show that the retriever tends to retrieve from relevant sources, and having more OOD data in the datastore does not have a negative impact on performance. Based on these observations, we hypothesize that having useful information in the datastore is the most critical factor. Therefore, our rule of thumb for data source selection is to **maximize the chance that there exists data in the datastore that can provide helpful information for future queries**. In other words, data sources that could potentially contain useful information for future queries are all desirable for better performance.

Based on the above principle, we manually selected two types of data sources: 
The first type is general web data, such as data from CommonCrawl, which is a large and diverse data source that potentially covers various topics that a user may query at inference time.
The second type is domain-specific data, which is selected based on our prior knowledge of the evaluation domain. For example, we intentionally added more scientific data, such as pes2o scientific papers, math datasets, and biomedical data, to the datastore because we want the RAG system to perform better on scientific benchmarks such as MMLU and MedQA.

In summary, our data source selection was done empirically and we believe that exploring automatic data selection methods for datastore construction is a promising future direction. For instance, one potential follow-up work could involve training classifiers for source-level or document-level data selection to reduce the size of the datastore while maintaining its effectiveness for targeted tasks. We will include these discussions in the updated paper.


**Q2.** Regarding the application of MASSIVEDS to other complex tasks, would you recommend maintaining the same datastore while adapting the retriever and language model to be more task-specific?

**A2.** As discussed in A1, when adapting to new tasks, we recommend optimizing the datastore composition to include more data sources that are potentially helpful for these new tasks. Prior knowledge about the task distribution can help guide targeted data selection. For example, you may want to include more code data when building a datastore for coding tasks, such as HumanEval.

**Q3.** Could additional experiments be conducted on TriviaQA, using a merged corpus such as <MedQA corpus + TriviaQA corpus>, to assess whether such integration enhances performance?

**A3..** We would like to clarify that including more data in the datastore is likely, but not guaranteed, to improve performance. However, you can still try including everything in a datastore because it is robust to OOD data if there is no storage or computational constraints for the datastore size. To illustrate this, we conducted additional experiments where we compared the performance of different datastore combinations using the setup suggested by the reviewer. The results are shown in the below table.

| Datastore     | LM-only | Wikipedia | Wikipedia + C4 | Wikipedia+PubMed | MassiveDS |
|---------|------|---|-------|---------|------|
| TQA             |    64.1   | 72.6      | 75.8           | 72.6              | 77.0      |

As shown in the table, adding a new data source that is a hybrid of OOD data and potentially helpful data, such as C4, can further improve the performance of the in-domain datastore. Meanwhile, including a data source that only contains OOD data, such as PubMed, does not improve or decrease the TQA performance. However, we note that having PubMed in the datastore is helpful for biomedical tasks. Therefore, we recommend increasing the diversity of the datastore to enhance the performance on both single or multiple tasks. We will clarify that the data sources should still be carefully chosen with a consideration on what sources would be potentially helpful to the targeted tasks.

We thank the reviewer for the insightful discussion. We are happy to follow up if the reviewer has any other questions!
Dear Authors,

Thank you for your detailed response. I appreciate the efforts to address my initial concerns, and I have accordingly updated my review score. However, I still have some follow-up questions:

1. What are the critical factors to consider when expanding a datastore with diverse data sources to ensure performance improvements in RAG?

2. Regarding the application of MASSIVEDS to other complex tasks, would you recommend maintaining the same datastore while adapting the retriever and language model to be more task-specific?

3. Could additional experiments be conducted on TriviaQA, using a merged corpus such as <MedQA corpus + TriviaQA corpus>, to assess whether such integration enhances performance?