Source-linked AI summary
Fine-Tuning LLaMA for Multi-Stage Text Retrieval
Xueguang Ma, Liang Wang, Nan Yang, Furu Wei, Jimmy Lin
TL;DR
Existing multi-stage retrieval research largely predates recent LLM advances, motivating evidence about whether fine-tuned LLMs improve retrieval and reranking. The study fine-tunes LLaMA as RepLLaMA and RankLLaMA, finding stronger effectiveness than smaller models and strong zero-shot performance, while supporting holistic document representation.
Problem
Existing multi-stage retrieval studies largely use models that predate recent LLM advances, leaving the effectiveness of state-of-the-art LLMs in this setting underexplored.
Method
The study fine-tunes LLaMA-2 as the dense retriever RepLLaMA and pointwise reranker RankLLaMA for passage and document retrieval on MS MARCO.
Results
RepLLaMA–RankLLaMA achieves state-of-the-art multi-stage retrieval effectiveness, outperforms smaller models, and demonstrates strong zero-shot effectiveness on BEIR.
Takeaways & Limitations
Long-context LLMs can represent entire documents and reduce reliance on traditional segmentation and pooling strategies for document retrieval.
Takeaways & Limitations
Longer input lengths have diminishing returns while substantially increasing training and inference costs.
Abstract
from arXiv · showhide
The effectiveness of multi-stage text retrieval has been solidly demonstrated since before the era of pre-trained language models. However, most existing studies utilize models that predate recent advances in large language models (LLMs). This study seeks to explore potential improvements that state-of-the-art LLMs can bring. We conduct a comprehensive study, fine-tuning the latest LLaMA model both as a dense retriever (RepLLaMA) and as a pointwise reranker (RankLLaMA) for both passage retrieval and document retrieval using the MS MARCO datasets. Our findings demonstrate that the effectiveness of large language models indeed surpasses that of smaller models. Additionally, since LLMs can inherently handle longer contexts, they can represent entire documents holistically, obviating the need for traditional segmenting and pooling strategies. Furthermore, evaluations on BEIR demonstrate that our RepLLaMA-RankLLaMA pipeline exhibits strong zero-shot effectiveness. Model checkpoints from this study are available on HuggingFace.
1 Introduction
Multi-stage retrieval combines efficient candidate retrieval with reranking, but recent LLM-based rerankers leave corpus-scale retrieval, supervised fine-tuning, and efficient parallel scoring unresolved. This study investigates fine-tuning LLaMA-2 as both retriever and reranker to address these gaps.
- Multi-stage retrieval uses a retriever to find top-k texts and a reranker to refine their ordering.
- Recent LLM retrieval methods explore zero-shot reranking through pairwise or listwise prompting.
- Existing prompted rerankers do not address corpus-scale retrieval, supervised labels, or efficient parallel scoring.
- The study asks how state-of-the-art LLMs perform when fine-tuned specifically for multi-stage text retrieval.
- Fine-tuned LLaMA-2 models, RepLLaMA and RankLLaMA, target retrieval and reranking across MS MARCO and BEIR.
2 Method
The method fine-tunes LLaMA-based bi-encoder and pointwise architectures for complementary retrieval stages. RepLLaMA uses dense similarity and contrastive negatives, while RankLLaMA scores and reorders retrieved query–document pairs.
- Retriever: RepLLaMA adapts the DPR bi-encoder architecture by initializing its backbone with LLaMA.
- Retriever: Because LLaMA is unidirectional, each query or document ends with an end-of-sequence token whose representation serves as its embedding.
- Retriever: RepLLaMA computes relevance as the dot product between query and document embeddings.
- Retriever: Retriever training uses InfoNCE with human-labeled positives, hard negatives from top-ranked results, and in-batch negatives.
- Reranker: RankLLaMA jointly inputs a query and candidate document, generates a relevance score, and reorders candidates by descending score.
- Reranker: RankLLaMA uses a linear projection of the end-of-sequence representation and trains with contrastive loss without in-batch negatives.
- Reranker: RankLLaMA hard negatives are sampled from RepLLaMA’s top-ranking results to match the reranker’s intended pipeline inputs.
3 Experiments
Experiments evaluate RepLLaMA and RankLLaMA for passage and document retrieval on MS MARCO, with additional zero-shot evaluation on BEIR. The models achieve strong effectiveness, while document retrieval uses long-context inputs without segmentation.
- Experimental setup: Experiments evaluate the RepLLaMA–RankLLaMA multi-stage pipeline on MS MARCO passage and document retrieval tasks.The passage experiments use approximately 500k training examples, while document retrieval uses 300k training examples.
- Training: RepLLaMA training uses hard negatives and in-batch negatives, with a blend of BM25 and CoCondenser negatives for the MS MARCO passage experiments.Both retriever and reranker training use batches of 128 and 15 hard-negative passages per query.
- Passage retrieval: RepLLaMA achieves the highest passage-retrieval effectiveness among competing methods, exceeding bi-SimLM by 2 points MRR@10 on development queries.RepLLaMA also achieves higher MRR@10 and Recall@1k than GTR-XXL across development and TREC DL19/DL20 evaluations.
- Passage retrieval: The complete RepLLaMA–RankLLaMA pipeline beats RankT5 by 1.5 points MRR@10, while RankLLaMA-13B exceeds the 7B model by 0.3 points MRR@10.The larger model also achieves slightly higher nDCG@10 on both DL19 and DL20.
- Passage retrieval: The pipeline outperforms RankGPT4 by 0.4 and 7.3 nDCG@10 points on DL19 and DL20, respectively.RankLLaMA’s pointwise scoring permits parallel candidate reranking, unlike RankGPT’s sequential sliding-window strategy.
- Zero-shot evaluation: On BEIR, RepLLaMA and RankLLaMA demonstrate superior zero-shot effectiveness, with RepLLaMA exceeding SGPT and Ada2 by 3 points and GTR-XXL by approximately 6 points.RepLLaMA uses the base pretrained model as initialization without additional contrastive pretraining; RankLLaMA-7B further improves retrieval by an average of 1.5 points nDCG@10.
- Document retrieval: For full-document retrieval, RepLLaMA exceeds CoCondenser-MaxP by approximately 3 points MRR@100, while RankLLaMA exceeds MORES+ by 1 point.The models use the first 2048 tokens without segmentation, covering about 77% of documents entirely.
4 Ablation Study and Analysis
The ablations examine parameter-efficient fine-tuning and input length, showing that LoRA generalizes better in some evaluations while longer reranking inputs improve effectiveness with diminishing returns.
- Full Fine-Tuning vs. LoRA: Full fine-tuning exceeds LoRA by approximately 6 MRR@10 points on the training set, but improves development effectiveness by only 0.4 points.On TREC DL19/DL20, LoRA performs better, suggesting full fine-tuning may overfit the training distribution.
- Full Fine-Tuning vs. LoRA: The main experiments use LoRA instead of full fine-tuning because it generalizes better on independent human-judgment datasets.The authors attribute this choice to the possibility that full fine-tuning overfits the training-set distribution.
- Input Sequence Length: RankLLaMA’s document-ranking MRR@100 rises from 48.5 to 50.3 as maximum training length increases from 512 to 2048 tokens.The comparison varies both training and inference maximum input lengths.
- Input Sequence Length: Increasing reranking input length to 4096 raises MRR@100 to 50.6, demonstrating gains from exploiting longer sequences.The model is evaluated on document reranking with different maximum input lengths.
- Input Sequence Length: Beyond a certain length, gains plateau: the 4096-token model is only 0.3 MRR@100 points above the 2048-token model at matched lengths.Training the 4096-token model takes about 8 days versus 4 days for the 2048-token model, with similar relative inference-latency costs.
5 Related Work
Related work spans Transformer-based retrievers and rerankers, dense-retrieval improvements, prompt-based LLM reranking, and billion-parameter retrievers. The paper positions RepLLaMA and RankLLaMA as open-source models optimized for the complete multi-stage pipeline.
- Pre-trained Language Models: Transformer retrieval research includes encoder-only, encoder–decoder, and decoder-only architectures, with scaling and instruction tuning improving general-purpose LLM capabilities.Examples include BERT, T5, GPT/GPT-2, InstructGPT, and GPT-4.
- Multi-Stage Text Retrieval: Multi-stage retrieval commonly combines a retriever with reranking stages, including BM25 followed by BERT-based rerankers in earlier work.The monoBERT pipeline demonstrated the effectiveness of pre-trained language models for reranking.
- Dense Retrievers and Rerankers: Dense-retrieval studies improve effectiveness through hard-negative mining, continuous retrieval-oriented pre-training, and related training-data strategies.Representative methods include ANCE, RocketQA, CoCondenser, RetroMAE, and SimLM.
- Dense Retrievers and Rerankers: Reranker research explores custom pre-training, contrastive loss, and variations in model architecture and loss functions.Examples include monoT5, monoELECTRA, and analyses using T5 rerankers.
- Prompt-based LLM Reranking: Prompt-based LLM rerankers generate reordered candidate lists or compare passages pairwise, but the paper instead fine-tunes models for retrieval.Examples include LRL, RankGPT, RankVicuna, and PRP.
- Multi-Stage Text Retrieval: RepLLaMA and RankLLaMA extend billion-parameter retrieval work by providing open-source models optimized for multi-stage retrieval and evaluated in-domain and out-of-domain.The paper contrasts them with GPT-XXL, SGPT, and the non-open-source cpt-text model.
6 Conclusion
The study shows that LLaMA can be fine-tuned as both a dense retriever and a pointwise reranker for an effective multi-stage retrieval system. It reports state-of-the-art performance, greater optimization and efficient-inference potential than generative prompting approaches, and continued potential for LLM-based retrieval.
- Conclusion: LLaMA can be fine-tuned as a dense retriever, RepLLaMA, and a pointwise reranker, RankLLaMA, within a multi-stage retrieval system.The resulting system is reported as effective and state of the art.
- Conclusion: The RepLLaMA–RankLLaMA system outperforms smaller models built on the same basic design.The conclusion frames this as evidence for using large language models in retrieval tasks.
- Conclusion: Fine-tuned LLM retrieval offers greater optimization and efficient-inference potential than recent generative prompting methods for text reranking.This comparison concerns methods that prompt LLMs to rerank text generatively.