PaperlayerView exact revision on arXiv

Source-linked AI summary

How to Train Your DRAGON: Diverse Augmentation Towards Generalizable Dense Retrieval

Sheng-Chieh Lin, Akari Asai, Minghan Li, Barlas Oguz, Jimmy Lin, Yashar Mehdad, Wen-tau Yih, Xilun Chen

arXiv:2302.07452v1cs.IRcs.CL

TL;DR

Dense retrievers often face a supervised-versus-zero-shot effectiveness tradeoff, and it was unclear whether limited model capacity caused this pattern. The paper studies dense-retrieval data augmentation and proposes diverse queries and progressive supervision, with DRAGON achieving state-of-the-art effectiveness in both settings without increasing model size.

  • Problem

    Existing dense retrievers often trade supervised effectiveness against zero-shot effectiveness, raising the question of whether limited model capacity causes this tradeoff.

  • Method

    The paper systematically studies data augmentation and trains DRAGON with diverse queries, multiple retrievers, and progressive relevance-label augmentation.

  • Results

    DRAGON achieves state-of-the-art effectiveness in supervised and zero-shot evaluations without increasing model size or retrieval complexity.

  • Takeaways & Limitations

    Diverse relevance labels and mixtures of cropped-sentence and synthetic queries provide a generalizable dense-retrieval training recipe.

  • Takeaways & Limitations

    More training iterations do not produce effective improvement in either supervised or zero-shot evaluations.

Abstract

from arXiv · show

Various techniques have been developed in recent years to improve dense retrieval (DR), such as unsupervised contrastive learning and pseudo-query generation. Existing DRs, however, often suffer from effectiveness tradeoffs between supervised and zero-shot retrieval, which some argue was due to the limited model capacity. We contradict this hypothesis and show that a generalizable DR can be trained to achieve high accuracy in both supervised and zero-shot retrieval without increasing model size. In particular, we systematically examine the contrastive learning of DRs, under the framework of Data Augmentation (DA). Our study shows that common DA practices such as query augmentation with generative models and pseudo-relevance label creation using a cross-encoder, are often inefficient and sub-optimal. We hence propose a new DA approach with diverse queries and sources of supervision to progressively train a generalizable DR. As a result, DRAGON, our dense retriever trained with diverse augmentation, is the first BERT-base-sized DR to achieve state-of-the-art effectiveness in both supervised and zero-shot evaluations and even competes with models using more complex late interaction (ColBERTv2 and SPLADE++).

1 Introduction

Dense retrieval faces a tradeoff between supervised and zero-shot effectiveness, but DRAGON shows that BERT-base-sized models can achieve state-of-the-art performance on both without increased model size. The approach uses diverse queries and progressively augmented relevance labels.

  • Motivation: Dense retrievers face a clear tradeoff between supervised and zero-shot effectiveness, while larger or more complex models add efficiency or latency costs.GTR-XXL breaks the tradeoff using 40 times more parameters, whereas SPLADE++ and ColBERTv2 add retrieval complexity and latency.
  • Approach: DRAGON progressively augments relevance labels using multiple retrievers rather than relying on a strong cross-encoder.The method uses diverse sparse, dense, and multi-vector retrievers to provide varied supervision.
  • Approach: DRAGON is trained on 28 million augmented queries comprising cropped sentences and synthetic queries.These queries are combined with progressive relevance-label augmentation.
  • Results: DRAGON achieves state-of-the-art effectiveness in both supervised MS MARCO and zero-shot BEIR and LoTTE evaluations without increasing model size or retrieval complexity.The paper presents it as the first dense retriever to break the supervised–zero-shot tradeoff under these constraints.

2 Background

The background frames dense retrieval as bi-encoder similarity search trained with contrastive learning, then recasts several improvement strategies as data augmentation. It also establishes the supervised and zero-shot evaluation settings used in the study.

  • Dense Retrieval: Dense retrieval encodes queries and documents as low-dimensional vectors and ranks documents using dot-product similarity.Top-k retrieval can be performed with approximate nearest-neighbor search libraries.
  • Contrastive Learning: Contrastive learning trains dense retrievers by increasing similarity for relevant query-document pairs and decreasing similarity for irrelevant documents.Because comparing against the full corpus is impractical, training uses hard negatives or cross-batch samples.
  • Data Augmentation: Knowledge distillation, contrastive pre-training, and pseudo-query generation can be interpreted as different data-augmentation recipes.These approaches augment training instances or supervision for dense retrieval.
  • Data Augmentation: Query augmentation often assumes that the original document is relevant to augmented queries, providing only one relevance-label view.Cross-encoder-based label augmentation supplies predicted positive passages but is treated as another augmentation strategy.
  • Experimental Setting: The experiments use MS MARCO as supervised data and BEIR datasets for zero-shot evaluation.MS MARCO evaluation reports MRR@10 and Recall@1000, while BEIR provides zero-shot datasets.

3 Pilot Studies on Data Augmentation

The pilot studies compare strategies for augmenting queries and relevance labels, finding that diverse supervision and progressive training improve dense retriever generalization. They also show that scaling query augmentation favors cropped sentences, while mixing cropped and generative queries can strengthen performance when data is limited.

  • Pilot Studies on Data Augmentation: The pilot studies frame query and label augmentation as complementary data-augmentation dimensions for training dense retrievers.The framework addresses how to create supervision for augmented queries and how to train from multiple teachers.
  • Training with Diverse Supervisions: Diverse supervision samples relevance labels from multiple retrievers rather than fusing them into one ranking list.The study uses sparse, dense, and multi-vector teachers, with each ranking list representing a distinct relevance perspective.
  • Training with Diverse Supervisions: Uniform supervision improves zero-shot retrieval over single or fused supervision, while progressive supervision improves generalization further.The results motivate using progressive supervision in subsequent experiments.
  • Empirical Studies: The trajectory of progressive supervision substantially affects zero-shot effectiveness but has only a minor impact on supervised evaluation.Switching two teachers decreases averaged nDCG@10 over BEIR-13 by 1 point, while reversing the full trajectory decreases it by more than 1.5 points.
  • Empirical Studies: Using only ColBERTv2 at the final iteration worsens zero-shot retrieval compared with diverse supervision from three teachers.This comparison supports retaining multiple supervision sources throughout progressive training.
  • Empirical Studies: Scaling cropped-sentence queries from 0.8M to 28M significantly improves supervised retrieval and produces stronger generalization than human-like generative queries.A mixture of cropped and generative queries is especially effective when the query set contains 0.8–8M examples.

4 Comparison with the State of The Art

DRAGON is evaluated against dense retrievers across supervised MS MARCO and TREC DL and zero-shot BEIR and LoTTE settings. It maintains strong supervised effectiveness while advancing zero-shot retrieval, without increasing model size or retrieval complexity.

  • Supervised Evaluations: Models performing well on MS MARCO dev do not necessarily transfer well to TREC DL, whose judgments contain fine-grained relevance labels.
  • Supervised Evaluations: DRAGON variants achieve RR@10 38.0 on MS MARCO dev and nDCG@10 70.0 on TREC DL, showing consistently strong supervised effectiveness.
  • Zero-Shot Evaluations: DRAGON+ reaches state-of-the-art zero-shot effectiveness on BEIR, while all DRAGON variants outperform dense retrievers by a large margin and compete with SPLADE++ and ColBERTv2 on LoTTE.
  • Overall Results: DRAGON advances state-of-the-art zero-shot effectiveness on BEIR and LoTTE while retaining strong MS MARCO supervised effectiveness.
  • Ablation Results: DRAGON-S achieves the highest BEIR effectiveness among the compared query variants while sacrificing only slightly on MS MARCO datasets.

5 Discussions

The discussions examine supervision, initialization, and query augmentation choices, finding benefits from masked auto-encoding and diverse supervision while cross-encoder labels and soft labels are inefficient or difficult to normalize.

  • Relevance-label augmentation: Cross-encoder relevance-label augmentation does not improve DRAGON-S and further training only on cross-encoder rankings worsens retrieval effectiveness.Training exclusively on cross-encoder rankings performs worse than training with ranking lists from three retrievers.
  • Relevance-label augmentation: Using all retrievers for relevance-label augmentation remains more efficient per query than using a cross encoder.
  • Initialization: Masked auto-encoding benefits the approach more than contrastive pre-training, while combining generative and contrastive pre-training remains future work.
  • Soft labels: Uniform soft-label supervision from sparse and dense retrievers causes supervised and zero-shot effectiveness drops because their relevance-score distributions are difficult to normalize.
  • Query augmentation: Cropped sentences provide more diverse query topics than generative queries, which tend to surround the original passage’s main topic.

6 Related Work

Prior work improves dense retrieval through knowledge distillation and label augmentation, but these methods commonly rely on human-labeled passages or limited supervision. DRAGON instead studies multiple supervision sources over large-scale augmented queries.

  • Knowledge Distillation: Knowledge-distillation methods use soft labels from ColBERT, cross-encoders, or ensembles while retaining human-labeled relevant passages.These approaches improve dense-retriever effectiveness but do not replace human relevance labels.
  • Label Augmentation: Cross-encoders are challenging to use for large-scale relevance-label augmentation because of their low efficiency.This motivates alternatives that can generate supervision at scale.
  • Label Augmentation: Earlier work used a single sparse retriever to augment labels, whereas this work explores multiple supervision sources on large-scale augmented queries.The comparison distinguishes single-teacher label augmentation from DRAGON’s diverse supervision.
  • Curriculum Learning: Combining masked autoencoding pre-training with supervised contrastive learning can further improve model generalization capability.This result connects curriculum-learning-related strategies with dense-retriever generalization.

7 Conclusion

The paper presents DRAGON as a dense retriever developed through a unified data-augmentation framework. Its empirical studies examine query and relevance-label augmentation to identify training strategies for generalization.

  • Conclusion: DRAGON is introduced as a dense retriever trained with diverse augmentation.The name denotes Dense Retriever trained with diverse AuGmentatiON.
  • Conclusion: The paper uses a unified data-augmentation framework to analyze recent progress in dense-retriever training.The framework covers both query augmentation and relevance-label augmentation.
  • Conclusion: The experiments identify insights that challenge the assumption that cross-encoders are the most effective teachers and human-like queries are required.These findings motivate examining alternative augmentation choices for dense-retriever training.

A.1 Impacts of Top-k Positive Sampling

The section examines how the choice of top-k positive passages affects progressive label augmentation, using cropped MS MARCO sentences as augmented queries. The ablation results are reported in Table 8.

  • Table 8 presents the ablation results for progressive label augmentation from top-k passages using cropped sentences as queries.
  • The sampling scheme treats the top 10 passages from each teacher ranking list as positives and ranks 45–50 as negatives.
  • The experiment uses sentences cropped from MS MARCO as augmented queries.
  • Progressive label augmentation assigns top-k passages as positive examples across the tested sampling schemes.

A.2 An Intuition Behind Uniform and Progressive Supervisions

Uniform supervision can provide accurate and diverse labels early, while progressive supervision improves generalization more over training iterations. The analysis estimates supervision quality from positive-sampling probabilities and diversity across supervision sources.

  • A.2 An Intuition Behind Uniform and Progressive Supervisions: Uniform supervision estimates positive-selection accuracy by computing each passage’s sampling probability and ranking positives accordingly.Passages retrieved by multiple teachers receive higher sampling probability under uniform supervision.
  • A.2 An Intuition Behind Uniform and Progressive Supervisions: Ties between positives are further ranked using the sum of reciprocal teacher ranks.For passage p_m, the score is 1/r_m1 + 1/r_m2 + 1/r_m3.
  • A.2 An Intuition Behind Uniform and Progressive Supervisions: Supervision diversity is estimated by counting positive passages in the union sets formed from different supervision sources.This measures how many distinct positive passages the sources collectively provide.
  • A.2 An Intuition Behind Uniform and Progressive Supervisions: Uniform supervision is accurate and diverse at the beginning, but progressive supervision achieves greater generalization improvement across iterations.The comparison is reported for models trained with cropped sentences as queries.

A.3 MS MARCO Dataset Statistics

The MS MARCO statistics include original training and test queries, together with augmented queries used to train DRAGON models with full relevance-label augmentation by five teachers.

  • A.3 MS MARCO Dataset Statistics: MS MARCO statistics cover original training queries, Dev, DL19, and DL20 test queries, plus augmented queries used with five teachers.The augmented queries support DRAGON training with full relevance-label augmentation.

A.4 Latency Measurement for Relevance Label Augmentation

Latency measurements compare relevance-label augmentation across cross-encoder, dense, sparse, and multi-vector retrieval systems under specified hardware and batch-size settings. The detailed LoTTE results also identify weaker DRAGON performance on science-and-technology Forum queries than SPLADE++ and ColBERTv2.

  • A.4 Latency Measurement for Relevance Label Augmentation: Latency is measured using batch retrieval on one NVIDIA A100 GPU and 60 Intel Xeon CPUs.The hardware configuration supports GPU and CPU search comparisons.
  • A.4 Latency Measurement for Relevance Label Augmentation: Cross-encoder label augmentation re-ranks text pairs with batch size 100.This setting is part of the latency comparison across augmentation systems.
  • A.4 Latency Measurement for Relevance Label Augmentation: Contriever, GTR-XXL, and sparse retrieval use Faiss-GPU or Lucene indexes with batch size 100, while GTR-XXL uses batch size 25 because of GPU memory constraints.The GTR-XXL memory limitation is identified as its main bottleneck for batch retrieval.
  • A.4 Latency Measurement for Relevance Label Augmentation: ColBERTv2 uses PLAID for multi-vector retrieval with batch size 1, described as its only available option.The comparison therefore uses distinct retrieval configurations across model families.
  • A.4 Latency Measurement for Relevance Label Augmentation: DRAGON variants perform poorly on science-and-technology Forum queries compared with SPLADE++ and ColBERTv2.The paper suggests combining science-corpus pre-training with DRAGON training as a possible response.
Loading 2302.07452v1…