Source-linked AI summary
Unsupervised Corpus Aware Language Model Pre-training for Dense Passage Retrieval
Luyu Gao, Jamie Callan
TL;DR
Dense retrievers are effective but fragile to noisy training data and typically require large batches and heavily engineered fine-tuning. coCondenser uses Condenser architecture with an unsupervised corpus-level contrastive loss to pre-train noise resistance and a structured passage-embedding space. Across MS-MARCO, Natural Question, and TriviaQA, it matches or slightly improves on RocketQA with reduced computation, data engineering, and batch size.
Problem
Dense retrievers are difficult to train because they are sensitive to mislabeled data and need large batches to learn a stable embedding space.
Method
coCondenser combines Condenser pre-training with an unsupervised corpus-level contrastive loss to warm up passage embeddings.
Results
Across Natural Question, TriviaQA, and MS-MARCO, coCondenser performs on-par or better than RocketQA while using small-batch fine-tuning without data engineering.
Takeaways & Limitations
Corpus-aware pre-training is query agnostic, reusable for varied end-task queries, and reduces the computation and engineering required for effective dense retrieval.
Takeaways & Limitations
Checkpoint validation uses the last training checkpoint because full-corpus encoding makes validation costly under limited computational resources.
Abstract
from arXiv · showhide
Recent research demonstrates the effectiveness of using fine-tuned language models~(LM) for dense retrieval. However, dense retrievers are hard to train, typically requiring heavily engineered fine-tuning pipelines to realize their full potential. In this paper, we identify and address two underlying problems of dense retrievers: i)~fragility to training data noise and ii)~requiring large batches to robustly learn the embedding space. We use the recently proposed Condenser pre-training architecture, which learns to condense information into the dense vector through LM pre-training. On top of it, we propose coCondenser, which adds an unsupervised corpus-level contrastive loss to warm up the passage embedding space. Retrieval experiments on MS-MARCO, Natural Question, and Trivia QA datasets show that coCondenser removes the need for heavy data engineering such as augmentation, synthesis, or filtering, as well as the need for large batch training. It shows comparable performance to RocketQA, a state-of-the-art, heavily engineered system, using simple small batch fine-tuning.
1 Introduction
Dense retrieval is effective but difficult to fine-tune because robust systems rely on noise handling and large batches. coCondenser pre-trains these properties so models can be fine-tuned more simply, achieving performance comparable to heavily engineered methods.
- Dense retrievers require carefully designed fine-tuning techniques to realize their capacity.Prior approaches include iterative negative mining, multi-vector representations, and the optimized RocketQA pipeline.
- RocketQA combines denoising hard negatives with large-batch training, but its computation-heavy pipeline is inaccessible to resource-limited researchers.
- Typical language models are sensitive to mislabeled data, while their CLS vectors do not naturally form effective passage-embedding spaces.
- coCondenser uses corpus-aware unsupervised pre-training to make dense-retrieval models noise-resistant and give them a structured global embedding space.
- On Natural Question, TriviaQA, and MS-MARCO passage ranking, small-batch fine-tuning without data engineering performs on-par or better than RocketQA and other contemporary methods.
2 Related Work
Related work established dense retrieval, explored specialized pre-training and fine-tuning, and brought contrastive learning into language representation. coCondenser instead targets the full passage-embedding space through corpus-aware pre-training.
- Dense Retrieval Transformer LM: Transformer language models have advanced dense retrieval, from ICT pre-training and co-training to direct fine-tuning from BERT.
- Dense Retrieval Transformer LM: RocketQA further improved dense retrieval through a heavily optimized fine-tuning pipeline.
- Dense Retrieval Transformer LM: DPR-PAQ revisited dense-retrieval pre-training with domain-matched pre-training on 65 million synthetic question-answer pairs.
- Dense Retrieval Transformer LM: Condenser introduced specialized pre-training architecture rather than designing another pre-training task.
- Dense Retrieval Transformer LM: Dense retrieval supports efficient first-stage search through corpus pre-encoding and mature maximum-inner-product-search libraries such as FAISS.
- Contrastive Learning: Contrastive learning has been applied to NLP sentence representations, while dense-retrieval work here focuses on learning the full embedding space rather than isolated representations.
- Contrastive Learning: Large-batch contrastive learning remains a memory barrier for resource-limited setups, motivating gradient caching for nearly constant peak memory.
3 Method
The method combines Condenser architecture with corpus-aware contrastive pre-training to create a noise-resistant, structured passage embedding space without requiring the full RocketQA pipeline. Memory-efficient gradient caching enables large-batch pre-training, after which the backbone initializes query and passage encoders for supervised fine-tuning.
- 3.1 Condenser: Condenser divides Transformer blocks into early backbone, late backbone, and head layers, routing late-layer CLS and early-layer token representations through the head.The head’s outputs support masked language model training, forcing late layers to aggregate information into the CLS representation.
- 3.2 coCondenser: coCondenser augments Condenser’s masked language model loss with corpus-level contrastive learning over randomly sampled spans from documents.Random spans act as passage surrogates, with within-document spans encouraged to be similar and spans from different documents differentiated.
- 3.2 coCondenser: The contrastive objective warms up the full passage embedding space rather than training only individual representations, addressing the semantic weakness of Condenser CLS inner products.The method is query agnostic and uses the target search corpus instead of artificial query-passage pairs.
- 3.3 Memory Efficient Pre-training: Gradient caching decouples representation-gradient computation from encoder-gradient computation, allowing the full-batch gradient to accumulate over small GPU-resident subbatches.An extra no-gradient backbone pass produces cached CLS representations and their contrastive-loss gradients before model updates.
- 3.4 Fine-tuning: After pre-training, the Condenser head is discarded, and the backbone initializes query and passage encoders that are supervisedly fine-tuned on target-task data.The fine-tuning procedure uses supervised contrastive loss and two training rounds with BM25 and mined hard negatives.
4 Experiments
The experiments pre-train and fine-tune coCondenser for dense retrieval on MS-MARCO, Natural Question, and Trivia QA, comparing it with RocketQA and other systems. coCondenser performs on-par or better while using substantially less fine-tuning computation and data engineering.
- 4.2.1 Setup: The experiments evaluate coCondenser retrieval on MS-MARCO, Natural Question, and Trivia QA, using the corresponding official retrieval metrics.MS-MARCO uses MRR@10 and Recall@1000; Natural Question and Trivia QA use Recall at 5, 20, and 100.
- 4.2.1 Setup: The comparison includes RocketQA, lexical and dense baselines, DPR-PAQ, and a first-stage Condenser model.RocketQA is the main baseline, while DPR, ANCE, ME-BERT, BM25, and other systems provide additional comparisons.
- 4.2.2 Results: Across the three query sets, dense systems outperform sparse systems, with careful fine-tuning or pre-training producing a large margin over earlier dense systems.The result supports the reported capacity of low-dimensional embeddings for dense retrieval, while noting that this capacity is difficult to exploit naively.
- 4.2.2 Results: coCondenser shows small improvements over RocketQA while reducing MS-MARCO fine-tuning batch size from 4096 to 64 and greatly reducing computation and data engineering.The comparison is presented through the RocketQA and coCondenser training pipelines.
- 4.2.2 Results: DPR-PAQ performs best on Natural Question with RoBERTalarge, but coCondenser is a strong alternative when computational resources or large pre-training sets are limited.On MS-MARCO, DPR-PAQ is less effective than RocketQA and coCondenser when its pre-training supervision is distant from the task.
- 4.2.2 Results: Adding corpus-aware contrastive loss gives coCondenser a pre-structured embedding space that improves Recall across all datasets and supports robust fine-tuning.The strongest coCondenser system combines the corpus-aware loss with hard negatives; Condenser systems achieve their performance without denoising.
5 Analysis of Training Stages
The analysis shows that RocketQA’s stages progressively improve performance but require the full heavy pipeline, whereas coCondenser’s corpus-aware pre-training recovers strong retrieval with small batches and without denoising.
- Each RocketQA stage is critical, with performance improving steadily as stages are added.
- Condenser with hard negatives performs very close to the full RocketQA system.
- Condenser with hard negatives achieves better MRR than coCondenser without hard negatives, indicating stronger local retrieval but global misses from a cold-start embedding space.
- CoCondenser without hard negatives achieves Recall very close to full RocketQA using only a size 64 batch.
- Adding hard negatives produces the strongest coCondenser system, effective both locally and globally.
- Condenser systems achieve their performance without denoising, and the approach avoids data augmentation while using small training batches.
6 Conclusion
coCondenser is a corpus-aware, unsupervised pre-training method for dense retrieval that improves robustness and structures the embedding space while reducing fine-tuning effort. It achieves strong retrieval performance with limited resources, though future work could combine it with additional established techniques.
- coCondenser combines the Condenser architecture with a corpus-aware contrastive loss to produce noise-resistant, structured embeddings for dense retrieval.Pre-training is performed once per search corpus, is query agnostic, and can support varied end-task queries.
- coCondenser reduces dense-retriever fine-tuning costs while improving retrieval performance.
- Its performance is close or similar to models that are several times larger and require semi-supervised pretraining.
- The method offers hands-off pre-training and removes the need to design and test pre-training or fine-tuning techniques.
- Practitioners can use the pre-trained weights and limited resources to train dense retrieval systems with state-of-the-art performance.
- Future work may integrate additional well-tested pre-training and fine-tuning methods to further improve performance.