Scaling Retrieval-Based Language Models with a Trillion-Token Datastore

We thank the reviewer for acknowledging that our work makes valuable contributions to the advancement of these systems!

Weakness 1. It remains uncertain if the findings would hold with a single data source, whether they apply to specific tasks beyond general ones like QA, and how interactions between various retrievers and language models might influence the results.

A1. We will carefully adjust our claims in our final draft to highlight our choice of models and retrievers, and that our downstream tasks focus on QA. Given resource constraints, we chose to focus on the most widely adopted RAG designs and evaluation setups; we agree that there are many more questions to be studied, and we hope that our open-source data and code will facilitate such work in the future.

On the specific comments:

(1) It is hard to scale single data sources like Wikipedia to the trillion token regime, so we focus on multi-domain web-scale data; however, we compare with single in-domain data sources in Table 3.

(2) We agree it would be good to study more downstream tasks beyond QA, and we will emphasize this in the discussion. We note that it is common in prior work (RePlug, Shi, et al; Atlas, Izacard, et al.) to focus on QA as well, and that we are the first to study downstream tasks at large datastore scales.

(3) We have added results on new LMs (Llama3, Olmo, and Pythia families) and new retrievers (DRAGON and GTR) in Figure R1 and Table R1 in the general reply, respectively; these results are consistent with our submission.

---

Weakness 2. There is a lack of comparative analysis and comparison of MASSIVEDS with other existing datastores [1, 2] (Dolma and Pile).

A2. We would like to clarify that a datastore, an indexed resource ready for retrieval, differs from a pretraining corpus, which contains only raw data, like Dolma and Pile. Constructing a datastore from a pretraining corpus involves data cleaning, chunking, embedding, and indexing, which take substantial effort, and as such there are no comparable open-source datastores to MassiveDS.

That said, the reviewer raises an interesting point about comparing different datastore sources. We ran an additional experiment using datastores constructed from DCLM-baseline and FineWeb-Edu (which have been shown to perform better than Pile and Dolma). As Table R2 in the general reply shows, MassiveDS achieves comparable performance.

---

Weakness 3. The claim that indexing the entire datastore initially reduces computational overhead is debatable. This aspect could be alleviated significantly by alternative indexing strategies [3, 4].

A3. The reviewer is correct that embedding should be done once to optimize efficiency; this is also the core of our pipeline. However, there are additional challenges: e.g., the retrieval process is I/O bound and often takes significant time to load passages from a large document pool. Also, using existing software like FAISS requires building separate indices for different sets of dense embeddings, significantly increasing storage demands for various configurations. Our open-source pipeline provides an efficient, ready-to-use solution that overcomes these challenges and that can be used as a foundation for future research. We appreciate the reviewer’s point and will expand our discussion of this in the paper.

The alternative indexing strategies cited by the reviewer do not fully overcome the challenges: [3] speeds up the search process but doesn't lower the cost of index construction; [4] improves datastore construction efficiency, but needs additional training and doesn't avoid redundant storage for different configurations. We will discuss them in our next revision.

---

Weakness 4 and Q3. The overlap between the datastore and the downstream evaluation dataset raises concerns about data leakage.

A4. We agree that data leakage is a potential concern. In Appendix B.1, we describe the strategies we used to avoid it. Specifically, we applied a strict data decontamination method to remove documents that overlap with the evaluation samples, which is stricter than common practice such as Dolma and RedPajama. In Figure 2(b), we also assessed perplexity (PPL) scaling on M2D2, which is not part of MassiveDS, and found scaling curves similar to those using decontaminated RPJ.

---

Weakness 5. The author should reevaluate the claim with additional models.

A5. We thank the reviewer for suggesting this. We have added more results with different LMs in Figure R1 in our general reply and will include them in our paper.

---

Q1. How few-shot exemplars are selected for QA tasks?

A6. Following the popular evaluation repository LM-Evaluation-Harness, we use few-shot examples randomly sampled from the development set, which do not overlap with test samples.

---

Q2. How is the language model configuration determined? Is there any inherent randomness?

A7. We follow LM-Evaluation-Harness, which uses greedy generation for TriviaQA and NQ, and log-likelihood for MMLU and MedQA. Therefore, there is no randomness in the inference. We’ll clarify this in the paper.

---

Q3. (see Weakness 4)

---

Q4. How do you substantiate that the contribution of MASSIVEDS is significant and not merely an assembly of various data segments?

A8. Our primary contribution is enabling the scientific study of >1T-token RAG scaling on an academic budget, which extends beyond merely assembling the datastore and differs from previous systems that aimed to enhance search speed. Our identification of the commutability of different operations and the strategic ordering of these operations are crucial for facilitating accessible datastore scaling studies. Beyond this, we highlight our other contributions in presenting the first datastore scaling trends and analyses on various datastore design factors.

We thank the reviewer for their valuable suggestions and we’ll update our draft accordingly.

We thank the reviewers for acknowledging the importance of our research question and the value of our findings and open-sourced resources! We would like to address the reviewer’s concerns below:

Weakness 1. The claim of proposing the 'largest' datastore could be reconsidered as there are other pretraining corpus that exceed even 10T tokens.

A1. We would like to clarify that a datastore is different from a pretraining corpus: by “datastore”, we refer to an index that is ready for retrieval [1]; while a pretraining corpus only contains the raw data. Building a datastore on an existing pretraining corpus requires further data cleaning, data chunking, embedding, and indexing over the raw data, which is a non-trivial effort for the retrieval community. It has been challenging to conduct datastore scaling research at this scale before and existing open-sourced datastores are much smaller than MassiveDS, as shown in Table 1. Althogh there is pretraining corpus that is larger than our datastore in terms of raw text, we are the first to make this attempt to construct and open-source a datastore with over 1 trillion tokens.

[1] Asai, Akari, et al. "Reliable, adaptable, and attributable language models with retrieval."

---

Weakness 2. The proposed distributed index and retrieval pipeline is interesting but seems inflexible. How does it compare with existing search engine frameworks like ElasticSearch and Weaviate?

A2. We would like to note that our contribution is orthogonal to these existing frameworks: prior work, such as ElasticSearch and Weaviate, focuses on efficiency at inference, while our work focuses on the scientific evaluation of retrieval performance which requires not efficient serving, but efficient experimentation of various factors of the datastore. Even with the state-of-the-art nearest neighbor search used in production, studying the effect of various factors of the datastore like the size, quality filters, and decontamination methods remains expensive. This is because such experiments require rebuilding search indices for every combination of the factors, orthogonal to what search algorithm is being used. Our study focuses on removing the needs for such repetitive rebuilding of indices. The outcome is the comprehensive study of the impact of various datastore factors, as we demonstrated as part of the results, which is the novel part of our work.

---

Weakness 3. This paper lacks an analysis of retriever. Contriever-msmarco is used, whose training data may differ from MedQA and MMLU. Considering the limited generalizability of dense retrievers (as shown by the BEIR paper), stronger retrievers like those in MTEB could be used.

A3. We agree with the reviewer that we indeed only investigated one base retriever due to being limited by academic computational resources, and we studied the effect of improved retrieval by applying a cross-encoder reranker which is more computationally efficient in the paper. In the table below, we run the evaluation using 2 more base retrievers of similar sizes that perform better than Contriever on the MTEB benchmark using 10% randomly sampled MassiveDS. Our results show that, interestingly, the performance of the base retrievers of similar sizes on our general web data does not necessarily align with the ranking on MTEB. We hypothesize that this low correlation is because the domain compositions of MassiveDS and MTEB are different. Due to limited computational resources, we defer the study with larger retrievers, such as GRIT-7B, to future work. But we note that such larger embedding models are often prohibitively expensive to scale up to a trillion-token datastore (Wang et al., 2024).

---

Question 1. What is the recall rate of the retriever for different datastore sizes?

A4. In our evaluation setup, both upstream language modeling and downstream evaluation sets do not provide gold documents for the questions. Therefore, we can only report the end-to-end performance on these tasks.

We thank the reviewer for highlighting our open-source contribution and acknowledge the potential impact on the community! We would like to address the reviewer’s concerns below:

Weakness 1. This work does not study the impact of the scaling of the datastore on inference time.

A1. Our main focus is to see the effect of dataset scaling on upstream and downstream model performance. While we agree that a larger datastore can introduce additional latency, there's active research on improving the speed of nearest neighbor search; we believe our contributions are orthogonal. In addition, we show a small LM with retrieval augmentation could match or even outperform larger LM-only models, as supplemented in Figure R1 in the general reply, which indicates a potential reduction of inference time by using a small LM augmented with MassiveDS. Given the complexity of inference speed optimization and its loose relationship with our focus, we defer the study of inference-time efficiency-performance tradeoff to future work.

---

Question 1. The codes are not available by the time of submission.

A2. We thank the reviewer for reminding us to provide anonymous codes with the submission. We uploaded our code in an anonymous github repository and sent it to the AC according to the rebuttal rules. We hope the reviewer could easily get access to the code.

---

Question 2. Why don’t the authors use the term “retrieval augmented generation (RAG)”?

A3. We define retrieval-based LMs as a general family of LMs that leverage large-scale text data at inference time [1], and RAG [2] is one such model. We believe that our finding is generally applicable to other types of retrieval-based LMs such as kNN-LM style models, and therefore decided to use the term. More specifically, we focus on retrieval-in-context LMs as mentioned in Section 2, which is often used interchangeably with RAG [3] but more specifically describes off-the-shelf LMs that use retrieved context at inference time.

[1] ACL 2023 Tutorial: Retrieval-based Language Models and Applications https://acl2023-retrieval-lm.github.io/

[2] Lewis, Patrick, et al. "Retrieval-augmented generation for knowledge-intensive nlp tasks."

[3] Min, Sewon, et al. "Silo language models: Isolating legal risk in a nonparametric datastore."

---

Question 3. Can more details be provided about data decontamination/deduplication?

A4. The motivation of deduplication for a datastore is to remove documents that are exactly or approximately the same such that the retrieved top-k documents won’t contain repetitive information. To achieve this, a similarity score is computed between every two documents in the datastore, and one document is removed from every duplicate pair. While decontamination aims at removing the same or near-duplicate documents in the datastore that directly contain the information about the test samples. The motivation is to avoid having test samples directly in the datastore such that the model can achieve high evaluation performance by cheating. The test sample is compared against every document in the datastore, and the documents with high similarity scores are removed from the datastore for this certain task. In our paper, the 13-gram Jaccard similarity score is one metric we use to measure the overlapping rate between two documents (in deduplication, it’s 2 documents from the datastore; in decontamination, one document is the test sample and another document is from the datastore). We set the threshold to 80% for both deduplication and decontamination that we remove a document if it has a similarity score that is higher than 80% with another document. Since decontamination plays an important role in evaluation, we combined the 13-gram Jaccard similarity metric with another metric for decontamination to make our datastore less likely to be contaminated. More details can be found in Appendix B.1.

---

Question 4. Is there a reason why the baselines LM-only on the subplots (d), (e), and (f) are not extended with a dashed line?

A5. We thank the reviewer for pointing out this detail! We forgot to add a dashed line for these 3 subfigures and we will add them back in the next revised version!

We thank the reviewer for acknowledging the importance of our study and the solidness of our work! We would like to address the only concern by the reviewer on the choice of retriever below.

Q1. The retriever and reranker used in this paper may be somewhat outdated. Recent models on the MTEB leaderboard might offer better performance. The benefits of scaling datastore sizes could be more significant with a more powerful retriever. However, this is a very minor point. Considering the computational cost of indexing all the data, trying other retrievers is not needed.

A1. We thank the reviewer for understanding our computational constraints, and we agree with the reviewer that the base retriever we used in this paper isn’t new. We chose Contriever-MSMACRO because it was used in many previous works cited in L178-179. We agree that benefits of scaling datastore sizes could be more significant with a more powerful retriever, which is consistent with our conclusion in Section 6.2 that improved retrieval could enhance scaling trends. In the table below, we tried 2 more base retrievers that outperform Contriever on the MTEB leaderboard using 10% randomly sampled MassiveDS. Our results show that, interestingly, the performance of the base retrievers of similar sizes on our general web data does not necessarily align with the ranking on MTEB. We hypothesize that this low correlation is because the domain compositions of MassiveDS and MTEB are different. Due to limited computational resources, we defer the study with larger retrievers, such as GRIT-7B, to future work. But we note that such larger embedding models are often prohibitively expensive to scale up to a trillion-token datastore (Wang et al., 2024).

---

Q2. Typo: L223 "Datascore" should be "Datastore".

A2. We thank the reviewer for pointing out the typo for us. We will fix it in our next version. We are happy to address more questions from the reviewer during the rebuttal period.

We thank the reviewer for acknowledging the novelty and solidness of our work! We would like to address the concerns and questions below:

Weakness 1. The evaluation is primarily focused on QA and knowledge-intensive tasks where one can expect the retrieval based LMs to work well. Inclusion of more diverse tasks, such as commonsense reasoning or open-ended text generation, would provide a more comprehensive assessment of the datastore's scalability.

A1. We agree that more diverse tasks would make our study more comprehensive. Our study takes a step in this direction as it is the first analysis, to the best of our knowledge, of datascore scaling on downstream tasks. Previous studies (RETRO, Borgeaud, et al.) examine retrieval-based language models with datastores exceeding 1 trillion tokens have focused solely on language modeling, and it remains unclear how datastore scaling and different datastore factors impact the performance of downstream tasks. In addition to knowledge-intensive tasks, we also include tasks that involve reasoning on top of knowledge. In particular, we included MMLU, which is widely recognized as a standard metric for assessing the performance of modern LMs, in our evaluation. Given our resource constraints, it was challenging to cover more downstream tasks, as we discussed in our limitation sections. We will emphasize this point in the next revision, and we hope that by open-sourcing our datastore and pipeline, we will enable others to build upon our work and evaluate large-scale retrieval-based models on a more diverse set of downstream tasks.

---

Weakness 2. For each of the analysis dimensions, the techniques are basic: E.g., a) quality filtering: basic filtering from DOLMA but leave the higher-quality filters b) Decontamination: 13-gram models for decontamination and leave out some techniques such as Min-K% Prob c) Retrieval Model. This might be due to the fact that a major portion of the paper was focused designing the scalable experimental set-up for this analysis.

A2. We thank the reviewer for acknowledging that our main focus is on the design of scalable experimental setup, as it is challenging to cover all possible datastore configurations given our computational constraints. Nevertheless, our analyses have covered several key aspects that people care about when designing the datastores, and our findings are inspiring.

Here are some discussions on the points mentioned by the reviewer:

(1) We agree with the reviewers that the DOLMA filters are basic and that higher-quality filters, such as fasttext filtering [1], are not covered in our study. We focus on standard filters that have been applied to both traditional and modern pretraining corpora such as RedPajama, C4, Dolma, etc. We believe these analyses are interesting to a broad community. In addition, we show that it is easy to test the effect of different quality filters using our pipeline and defer the study on more recent filters to future works.

(2) For data decontamination, we considered 13-gram Jaccard similarity-based decontamination because it has been widely adopted by many works, such as GPT-3, RETRO, Dolma, etc. Beside this, we also applied another decontamination method, longest-string decontamination, which enables us to experiment with different strictness levels of decontamination, as shown in Figure 4. By tuning the hyper-parameters of the decontamination methods, we show datastore scaling can benefit the model’s performance across different decontamination levels (Figure 4), so our conclusions hold regardless of the strictness of the decontamination method. Separately, Min-K% prob is designed for detecting if a given LM has been trained on certain data, and isn’t intended to be a data decontamination method.

(3) We indeed only used a single base retriever model in the submission, so we add a new ablation on the base retriever using 10% of MassiveDS and supplement the results in the table below. We chose DRAGON and GTR-Base because they rank higher than Contriever on the MTEB benchmark and are similar in size. Interestingly, the results indicate that these base retrievers exhibit similar performance on MassiveDS. We hypothesize that this is due to the differing domain compositions of MassiveDS compared to the datasets used in MTEB.

Due to limited compute, we leave it to future work to explore the performance with a larger and more capable retriever such as GRIT-7B. But we note that such larger embedding models are often prohibitively expensive to scale up to a trillion-token datastore (Wang et al., 2024). However, we are optimistic that significantly improving the retriever will lead to better scaling trends. This expectation is supported by the evidence in Figure 5, which shows that enhancing retrieval quality with a cross-encoder reranker improves performance.

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

---

Q3. Previous studies (https://arxiv.org/html/2307.07164v2) have shown that lack of diversity adversely impacts the results. What was the diversity of results retrieved by the underlying retrieval?

A3. As shown in Figure 3, the retriever retrieves from diverse sources in MassiveDS, and it tends to retrieve from relevant domains more frequently than other domains. Our results also indicate that increasing diversity helps improve the performance—MassiveDS outperforms single-domain datastores on the evaluated tasks, as shown in Table 3. We will discuss this in our next version.

We thank the reviewer for pointing out the comprehensiveness of our work! We would like to address the concerns and questions below, and we will edite the paper accordingly.

Weakness 1. Large-scale embedding-based search is not a new topic but has been studied for years. Arguably it is a mature technology ready in production [1, 2]. Information like hardware latency and memory/CPU footprint were not reported; comparisons with popular techniques like DiskANN [2] are not presented. This could be confusing and/or misleading to people reading the paper: should they adopt the search pipeline introduced in the paper, or just use a more developed technology? Also, this makes it difficult to assess the contribution of the entire paper.

A1. The work listed by the reviewer focuses on the search efficiency at inference with a fixed datastore configuration, but we would like to clarify that this is orthogonal to the focus of the current paper. Our main goal is to conduct a thorough analysis of the effect of scaling the datastore for retrieval-based language models and the impact of various datastore factors, where the inference speed isn’t a key bottleneck—even with the state-of-the-art nearest neighbor search used in productions such as DiskANN, studying the effect of various factors of the datastore like the size, quality filters, and decontamination methods remains expensive. This is because such experiments require rebuilding search indices for every combination of the factors, unrelated to which search algorithm is being used. Therefore, we design a new pipeline for efficient experimentation of various factors of the datastore by removing the need for repetitive datastore construction and repetitive large-scale retrieval, which are orthogonal and can be combined with prior work on efficient search algorithms.

---

Weakness 2. While the paper shows a general improving trend in RAG tasks and with relatively capable LMs, the models in the paper still seem to underperform the popular 2-year-old RAG model ATLAS [3] by a significant margin. Specifically, the design of the RAG system used in the paper can be somewhat basic. How much the results transfer to stronger system/system designs remains relatively unclear.

A2. Despite active research on advancing RAG designs, there have been few scientific studies of the scaling properties of datastores with more than a few billion tokens. As an initial open-source study on datastore scaling, we picked the most basic but widely adopted RAG design instead of any particularly optimized system. In Figure 5, we presented an analysis of the impact of improved retrieval on the datastore scaling performance, which shows that adopting a more advanced retrieval technique, i.e., a reranker, can further improve the scaling performance, indicating a positive sign that our conclusions could potentially be extended to more recent RAG designs. Overall, we focus on conducting a comprehensive analysis of datastore scaling and sharing our thoughts on how future works could further improve the scaling curves, rather than achieving state-of-the-art scores on specific datasets with a particularly optimized RAG design.

The purpose of Atlas is to finetune the retriever for best performance. Atlas finetunes the retriever and the LM to adapt to every specific task using the full task data (>10k) for their best-reported scores, while we evaluate the model in a training-free fashion by prepending only 5 examples in context. The purpose of our work is to disentangle the factors that affect performance and study them in detail. As such, Altas is orthogonal work. The insights from our work can be combined with Atlas, but here we ask more basic questions about RAG scaling.

---

Question 1. It was discussed in the paper that MASSIVEDS can outperform in-domain datastores on two eval sets. I wonder how this happens. Are there pieces of information easier to retrieve and/or read from out-of-domain?

A3. MassiveDS is a hybrid multi-domain datastore that includes both “in-domain datastores” as subsets and general web data to enhance its knowledge base. The key finding here is that the model is robust to other out-of-domain data in the same datastore and the retriever still preferentially retrieves from the right domain (Figure 3), so it eliminates the need to develop domain-specific datastores for each task, which is often a costly and complex process. Consequently, users can establish a single, general-purpose datastore that is effective across various tasks, even in cases where creating a task-specific in-domain datastore is unfeasible. This flexibility greatly simplifies data management and expands the utility of MassiveDS in diverse applications.

---

Question 2. The authors claim to observe an empirical log-linear scaling trend in the RAG system. It would be interesting to see some discussion on why this is the case and the implications of this. How does this happen? Does this mean retrieval datastore scaling has an extremely diminishing gain?

A4. Many previous scaling law studies have shown significant gains by scaling the pretraining data and the model size, and they all present performance gains with the x-axis (the number of tokens for pretraining or the number of parameters) in the log scale [1,2]. We find similar performance gains by scaling the datastore to the gains by scaling pretraining data. Therefore, our results indicate that the datastore could be another dimension to scale in addition to the pretraining data and the model size.

[1] Gadre, Samir Yitzhak, et al. "Language models scale reliably with over-training and on downstream tasks."

[2] Kaplan, Jared, et al. "Scaling laws for neural language models."

---

Question 3. Can the authors provide some statistics based on chunks/passages?

A5. We have 4B passages in our datastore, which average around 360 tokens. We will report detailed statistics in the next version of the paper.

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!