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 a computationally accessible manner. Finally, we analyze the effect of improving the retriever, datastore quality filtering, and other design choices on our observed scaling trends. Overall, our results show that datastore size should be considered as an integral part of LM efficiency and performance trade-offs. To facilitate future research, we open-source our datastore and code at https://github.com/RulinShao/retrieval-scaling.
Paper
Similar papers
Peer review
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.
Strengths
1. The exploration of how datastore configurations impact retrieval-based language models is both intriguing and significant. Some of the findings provide valuable contributions to the advancement of these systems. 2. The paper is well-written and accessible, with each step of the implementation thoroughly explained and presented. 3. The proposed datastore holds potential for related research in this field.
Weaknesses
1. The paper attempts to draw critical and general conclusions, yet the scope and depth of their experiments lack the robustness required to fully support these assertions. The study focuses on the impact of various datastore configurations, such as size, filters, and decontamination strategies. Evaluations are carried out on OpenQA tasks using a retriever specifically trained for QA. However, several key questions remain unaddressed. For example, it is unclear whether the findings would remain consistent using a single data source, whether the conclusions apply to specific tasks that are distinct from broader tasks like QA, and how the interplay between different types of retrievers and language models might affect the results. Without exploring these areas, claiming a definitive trend is premature. 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
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.
Rating
6
Confidence
4
Soundness
3
Presentation
4
Contribution
3
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.
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.
Strengths
1. This paper studies an important research question: the scaling of datastores for retrieval-based LMs. 2. This paper open-sources the 1.4T token datastore for retrieval-based LMs, which consists of diverse domains and is carefully filtered. This datastore could be a valuable resource for future retrieval-based LM studies. 3. Evaluation on multiple tasks demonstrates the importance of datastore scaling and highlights some interesting questions, such as how to improve retrieval performance and how to filter datastores.
Weaknesses
1. The claim of proposing the 'largest' datastore could be reconsidered. For LM pre-training, there are larger datastores like RedPajama-v2 with 30T tokens, and for IR, there are corpora like ClueWeb22 with 16T tokens. The proposed MassiveDS is mainly sourced from the RedPajama, making it similar to other existing collections. 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
1. What is the recall rate of the retriever for different datastore sizes?
Rating
6
Confidence
4
Soundness
3
Presentation
4
Contribution
2
Limitations
The limitations section is included and adequately discussed.
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
Strengths
1) OPEN SOURCE Both the models and data used in the work are open source and when the code is made available, it will facilitate future work and additional research on retrieval-based LM at scale. 2) WRITING/PRESENTATION The paper is easy to read, the contributions and findings are explicitly stated, and the tables and figures are easy to read. 3) IMPACT Given the importance of retrieval augmented generation for making LM more trustworthy and adapted to a specific domain/environment, releasing a large-scale dataset and pipeline for retrieval-based LMs can facilitate future research in that domain.
Weaknesses
INFERENCE TIME ANALYSIS 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
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?
Rating
7
Confidence
3
Soundness
3
Presentation
4
Contribution
3
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).
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.
Strengths
(1) To my knowledge, this is the first work studying scaling laws regarding datastore size. The author introduces the largest and most diverse open-sourced datastore for retrieval-based language models, containing 1.4 trillion tokens. This large-scale datastore can serve as a promising testbed for developing new retrieval-based LMs. (2) The author promises to release not only the data but also the code to reproduce the experiments. This could greatly facilitate research efforts within the community. (3) The author proposes a simple yet effective data processing pipeline that enables efficient investigation of the impact of different datastore features. (4) Extensive experiments show that the performance of retrieval-based LMs benefits from scaling datastore sizes. The evaluation includes both language modeling and downstream question answering. The results on MMLU and MedQA also highlight the need for developing retrieval-augmented LMs that can excel in tasks requiring reasoning abilities. (5) Some ablation studies provide interesting and insightful findings, such as data deduplication being a crucial factor for enhancing language model performance when the datastore size is extremely large. Overall, I believe this is a good paper studying an important problem supported by solid experiments.
Weaknesses
Overall, I do not find any significant weaknesses in the paper. The only minor issues I noticed are: (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
Please see my comments in Weakness.
Rating
7
Confidence
3
Soundness
4
Presentation
4
Contribution
3
Limitations
The limitations are properly stated in the Conclusion and Limitations section.
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.
Strengths
- **Originality**: The paper tackles the important question of how scaling retrieval-based LMs by the size of the inference datastore impacts upstream and downstream task performance which is novel and adds a new dimension to LM scaling laws. - **Quality**: The research is rigorous, with well-designed experiments and robust methodologies. - **Clarity**: The paper is well-organized, and the key contributions are clearly articulated. - **Significance**: The open-sourced resources will facilitate further research in this direction. - **Experiment Design**: The design of the experimental set-up is novel and theoretically grounded which will be interest to the research community on how to efficiently experiment with multiple dataset configurations
Weaknesses
The following are some of the limitations of the paper (some of these have already been highlighted by the authors): - **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
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?
Rating
7
Confidence
4
Soundness
4
Presentation
4
Contribution
2
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.
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.
Strengths
- The authors present a substantially larger retrieval datastore for retrieval augmented LMs. - Extensive experiments have been conducted to study how datastore designs/properties can impact RAG results. - Detailed analysis has been conducted on several eval datasets and models.
Weaknesses
- Large-scale embedding-based search is not a new topic but has been studied for years. Arguably it is a mature technology ready in production [1, 2]. The authors have made a rather arbitrary design in the index design. Information like hardware latency and memory/CPU footprint were not reported; comparisons with popular techniques like DiskANN [2] are not presented. This could be confusing and/or misleading to people reading the paper: should they adopt the search pipeline introduced in the paper, or just use a more developed technology? Also, this makes it difficult to assess the contribution of the entire paper. - 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
- 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?
Rating
5
Confidence
4
Soundness
2
Presentation
3
Contribution
3
Limitations
n.a.
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!
Response to Authors
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?
Response to the follow-up questions
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!
Response to Authors
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.
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!
Response to Rebuttal from Authors
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.
Reviewer-Author Discussions
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.
Thank the authors for further explanations. I do not have other concerns and will keep my original positive rating.
Decision
Accept (poster)