TLDR: A new research paper demonstrates that the information retrieval capabilities of Large Language Models (LLMs) predictably improve with increased model size, longer training durations, and higher pretraining computational effort (FLOPs). The study also found a strong correlation between an LLM’s in-context learning abilities and its retrieval performance, suggesting that smaller models, when trained extensively, can achieve retrieval performance comparable to larger models trained less. This research highlights the growing potential of open-source, decoder-style LLMs for advanced retrieval tasks.
A recent study from Databricks Mosaic Research delves into a fundamental question for the future of artificial intelligence: how does the ability of Large Language Models (LLMs) to retrieve information improve as they are trained with more computational power? The paper, titled Retrieval Capabilities of Large Language Models Scale with Pretraining FLOPs, provides compelling evidence that retrieval performance scales predictably with various factors, offering crucial insights for the development of next-generation AI.
Understanding the Foundation of Retrieval
The research, conducted by Jacob Portes, Connor Jennings, Erica Ji Yuen, Sasha Doubov, and Michael Carbin, explores how decoder-style LLMs perform on information retrieval tasks across different model sizes, training durations, and estimated FLOPs (Floating Point Operations – a measure of computational effort). The idea of converting pretrained LLMs into effective retrieval models isnos new, with earlier examples like SentenceBERT and GTR successfully adapting models like BERT and T5. However, these often required extensive further pretraining or fine-tuning with contrastive loss.
More recently, models like RepLlama have shown surprisingly good performance by simply fine-tuning a Llama 2 7B model on a relatively small dataset. This raised questions about what inherent properties of modern foundation models enable their strong retrieval capabilities. With models like Llama 2 7B being trained on trillions of tokens and possessing robust in-context learning abilities, it’s becoming clear that scale plays a significant role.
The Scaling Relationship
The researchers investigated this scaling by using pretrained checkpoints of MPT decoders, ranging from 125 million to 7 billion parameters. These models were trained on datasets varying from 1 billion to over 2 trillion tokens. Each model was then minimally fine-tuned on 500,000 samples from the MS MARCO dataset using InfoNCE loss and evaluated on the BEIR retrieval benchmark.
The study revealed several key findings:
- Retrieval performance consistently improves with increasing model size and longer pretraining durations for a fixed model size. This relationship is best described by accuracy-FLOPs curves, indicating that retrieval performance scales directly with the computational effort invested in pretraining.
- For most BEIR tasks, a smaller model trained on more data can achieve similar accuracy to a larger model trained on less data, up to a certain point. This means that models with similar total FLOPs can exhibit comparable retrieval performance, regardless of their exact size-vs-data split.
- A strong correlation exists between In-Context Learning (ICL) scores and retrieval scores across BEIR tasks. LLMs that demonstrate higher ICL capabilities almost invariably also achieve higher retrieval scores, suggesting a deep connection between these two abilities.
Also Read:
- Unpacking LLM Factual Stability: Introducing a New Robustness Score
- Large Language Models for Identifying Entities in Historical Documents
Implications for Future AI
These findings have significant implications for the future of dense retrieval models. The “Cambrian explosion” of open-source LLMs in the 1B to 13B parameter range, often trained on trillions of tokens and fine-tuned for nuanced language, makes them strong candidates for embedding models. This research suggests a continued shift away from older BERT-style models, which are seeing less development, towards these more powerful, decoder-based LLMs.
While the study focused on scaling properties with a baseline fine-tuning approach and did not aim for state-of-the-art performance, its results provide a strong foundation for understanding how to build more effective LLM-based retrievers. The consistent scaling trends observed across various model sizes and training durations highlight the importance of computational investment in developing highly capable information retrieval systems.
