PaperlayerView exact revision on arXiv

Source-linked AI summary

Large Dual Encoders Are Generalizable Retrievers

Jianmo Ni, Chen Qu, Jing Lu, Zhuyun Dai, Gustavo Hernández Ábrego, Ji Ma, Vincent Y. Zhao, Yi Luan, Keith B. Hall, Ming-Wei Chang, Yinfei Yang

arXiv:2112.07899v1cs.IRcs.CL

TL;DR

Dual encoders can have limited out-of-domain generalization because their final query–passage interaction is a simple dot product. GTR addresses this by scaling T5-based dual encoders while fixing the bottleneck embedding size and using multi-stage training, achieving strong BEIR generalization and high data efficiency.

  • Problem

    Dual encoders trained in one domain often fail to generalize, and their simple dot-product bottleneck may limit retrieval expressiveness.

  • Method

    GTR scales T5-based dual encoders while fixing the bottleneck embedding dimension, combining community question-answer pre-training with MS Marco fine-tuning.

  • Results

    GTR achieves state-of-the-art BEIR performance, with scaling improving out-of-domain retrieval and 10% of MS Marco data matching overall out-of-domain performance.

  • Takeaways & Limitations

    Scaling backbone encoders is a viable direction for improving single-vector dual encoders and their out-of-domain generalization.

  • Takeaways & Limitations

    Scaling increases inference latency from 17 ms for GTR-Base to 349 ms for GTR-XXL, and GTR-XXL underperforms on Trec-Covid while retrieving shorter documents.

Abstract

from arXiv · show

It has been shown that dual encoders trained on one domain often fail to generalize to other domains for retrieval tasks. One widespread belief is that the bottleneck layer of a dual encoder, where the final score is simply a dot-product between a query vector and a passage vector, is too limited to make dual encoders an effective retrieval model for out-of-domain generalization. In this paper, we challenge this belief by scaling up the size of the dual encoder model {\em while keeping the bottleneck embedding size fixed.} With multi-stage training, surprisingly, scaling up the model size brings significant improvement on a variety of retrieval tasks, especially for out-of-domain generalization. Experimental results show that our dual encoders, \textbf{G}eneralizable \textbf{T}5-based dense \textbf{R}etrievers (GTR), outperform %ColBERT~\cite{khattab2020colbert} and existing sparse and dense retrievers on the BEIR dataset~\cite{thakur2021beir} significantly. Most surprisingly, our ablation study finds that GTR is very data efficient, as it only needs 10\% of MS Marco supervised data to achieve the best out-of-domain performance. All the GTR models are released at https://tfhub.dev/google/collections/gtr/1.

1 Introduction

Dual encoders may struggle with out-of-domain retrieval because their query–passage interaction is compressed into a simple dot product. This paper tests whether scaling the encoder while fixing the bottleneck can improve generalization.

  • Dual encoders separately encode queries and documents into a shared fixed-dimensional space, then use approximate nearest-neighbor search for retrieval.
  • The bottleneck’s dot-product interaction may limit expressiveness and out-of-domain generalization.
  • The paper scales dual encoder models while keeping the bottleneck embedding dimension fixed.
  • GTR uses T5 encoders with up to 5 billion parameters and 768-dimensional bottleneck embeddings, trained with community question-answer pre-training and MS Marco fine-tuning.
  • On 18 BEIR tasks across 9 domains, scaling improves generalizability despite the fixed bottleneck, and 10% of MS Marco data matches overall out-of-domain performance.

2 Background

Dense retrieval uses learned vector representations to match queries with semantically related documents, but single-vector dot-product scoring may limit expressiveness. More interactive alternatives improve expressiveness while adding latency, motivating scaled single-vector retrievers.

  • Lexical methods such as BM25 rely on term overlap and can miss documents that are semantically related but have low lexical overlap.
  • Dense retrieval embeds queries and documents into low-dimensional vectors, enabling matching beyond lexical overlap.
  • Dual encoder performance may be bounded by the dot-product similarity function.
  • Multi-vector and token-level interaction models improve expressiveness but introduce substantial latency overhead.
  • The paper studies whether increasing model capacity can produce better fixed embeddings for single dot-product retrieval.

3 Generalizable T5 Retriever

GTR builds T5-based dual encoders using fixed-size mean-pooled embeddings and multi-stage training. The procedure combines web-mined question-answer pre-training with supervised retrieval fine-tuning.

  • T5 dual encoder: GTR initializes dual encoders from pretrained T5 models, whose encoder family spans millions to billions of parameters.
  • T5 dual encoder: The question and passage are encoded by the T5 encoder using mean pooling, with output embeddings fixed at size 768.
  • Training objective: Training uses an in-batch sampled softmax loss based on cosine similarity and a mini-batch temperature.
  • Training objective: The loss can include negative passages in its denominator and is trained bidirectionally for question-to-document and document-to-question matching.
  • Multi-stage training: Multi-stage training first uses web-mined semi-structured pairs, then human-annotated search datasets for retrieval fine-tuning.

4 Experimental setup

The experiments use T5-based dual encoders trained with community question-answer pre-training and supervised fine-tuning, evaluated across model sizes and retrieval baselines.

  • Training data: CommunityQA contributes 2 billion question-answer pairs for dual-encoder pre-training.The pairs are collected from online forums and question-answer websites.
  • Training data: MS Marco and Natural Questions provide human-annotated search data for fine-tuning.MS Marco contains 532K query-document pairs, while Natural Questions contains 130k query-passage pairs.
  • Baselines: The comparison includes BM25, DocT5Query, DPR, ANCE, TAS-B, and GenQ as sparse and dense retrieval baselines.The GTR experiments use four model sizes to investigate scaling effects across training stages.
  • Model configurations: GTR models use the encoder portion of T5 and vary across Base, Large, XL, and XXL configurations.The implementations use JAX and the encoder portion rather than the full T5 models.
  • Training configurations: The full GTR configuration combines pre-training and fine-tuning, while GTR-FT and GTR-PT isolate fine-tuning and pre-training, respectively.These settings test the effects of the two training stages separately.

5 Evaluation Results

Scaling GTR improves both in-domain and out-of-domain retrieval, with gains over established sparse and dense baselines and strong data efficiency in full models.

  • 5.1 Results on MS Marco: Increasing GTR model size consistently improves MS Marco NDCG@10 and other reported in-domain metrics.Similar improvements are reported for MRR@10 and Recall@1000.
  • 5.2 Results on BEIR generalization tasks: Increasing model size produces clear BEIR out-of-domain NDCG@10 gains, and GTR-XXL yields another performance jump.Similar improvements are reported for Recall@100, with scaling eventually outperforming all baselines across evaluation metrics on average.
  • 5.2 Results on BEIR generalization tasks: GTR-Base outperforms BM25 on 9 datasets, while GTR-XXL catches up to or surpasses BM25 on 5 additional datasets.GTR-XXL underperforms on 5 of the remaining tasks.
  • 5.2 Results on BEIR generalization tasks: GTR-Large already outperforms TAS-B and DocT5Query on BEIR, while larger models continue the ascending performance trend.The result supports scaling as a path toward better generalizability.
  • 5.2 Results on BEIR generalization tasks: 5.5%: GTR-XL outperforms BM25 on NFCorpus, reaching 0.343 versus 0.325 NDCG@10.The Base model underperforms BM25 on NFCorpus, whereas scaling improves retrieval performance on NFCorpus and BioASQ.
  • 5.3 Data efficiency for large retrievers: 10% of MS Marco data is sufficient for fine-tuning full GTR-Large, XL, and XXL models to achieve comparable or better out-of-domain performance.Using 10% reduces in-domain performance, while GTR-FT shows mixed out-of-domain results.

6 Ablation Study and Analysis

The analyses examine how scaling, pre-training, fine-tuning data, and retrieval-document length affect GTR’s zero-shot performance. Scaling improves generalization across training settings, while document-length effects vary by dataset.

  • Effect of scaling across training stages: GTR-FT XL reaches average zero-shot NDCG@10 of 0.433, exceeding TAS-B’s 0.415.
  • Effect of scaling across training stages: Scaling improves both in-domain and out-of-domain performance for fine-tuning-only models, and improves out-of-domain performance for pre-training-only models.
  • Effect of scaling across training stages: Combining generic pre-training with fine-tuning consistently improves GTR over GTR-FT models at every tested size.
  • Importance of the fine-tuning dataset: Fine-tuning on MS Marco creates a clear BEIR performance gap over pre-fine-tuning models, showing the importance of high-quality search data.
  • Importance of the fine-tuning dataset: Scaling improves zero-shot BEIR performance after fine-tuning on both NQ and MS Marco, with larger gains from weaker NQ data when scaling from Large to XL.
  • Document length versus model capacity: Scaling changes retrieved-document lengths: lengths increase on several datasets, while GTR-XXL retrieves much shorter documents on Trec-Covid.

7 Related Work

Related work frames neural retrieval as a progression from lexical matching to dense representations and stronger zero-shot methods. Prior work also highlights training techniques while leaving backbone capacity comparatively less studied.

  • Neural information retrieval: Neural retrieval extends beyond term matching and has achieved strong results across passage, question-answering, conversational, and bitext retrieval.
  • Dual encoders for neural retrieval: Dual encoders separately encode queries and documents into shared representations, with pre-trained backbones, negative mining, and large batches improving performance.
  • Dual encoders for neural retrieval: Few previous works had examined the effect of backbone model capacity on dual-encoder retrieval.
  • Zero-shot neural retrieval: Earlier work studied distillation and synthetic data for zero-shot dual-encoder retrieval, while multi-vector approaches increase query-document interactions.

8 Inference latency

Scaling GTR increases inference latency, creating a practical efficiency trade-off for larger dual encoders.

  • 17 ms, 34 ms, 96 ms, and 349 ms are the measured latencies for increasing GTR model sizes at batch size 1 and input length 128.
  • GTR-Base has latency close to TAS-B, whereas GTR-XXL has latency similar to re-ranking models.
  • The authors identify sparsity, distillation, and prompt-tuning as possible directions for reducing large-model inference time.

9 Conclusion

The paper introduces GTR, a scaled-up dual encoder with a fixed-size bottleneck, and evaluates its zero-shot retrieval performance on BEIR. Scaling improves retrieval broadly, especially out-of-domain generalization, with GTR-XXL achieving state-of-the-art BEIR performance.

  • GTR is a scaled-up dual encoder with a fixed-size bottleneck layer.
  • Scaling improves retrieval performance across the BEIR zero-shot benchmark, especially for out-of-domain generalization.
  • GTR-XXL achieves state-of-the-art performance on BEIR and outperforms many models using earlier query-document interactions.
  • The findings support continuing to improve single-vector representations through better backbone encoders.

A.1 Comparisons on MS Marco

The comparisons show that larger GTR dual encoders improve retrieval performance against established baselines, including on BEIR out-of-domain evaluation. GTR-XXL exceeds RocketQA and surpasses contrastive-pretrained models as model size increases.

  • Baseline comparisons: GTR-XXL outperforms RocketQA on both MRR and recall.The best RocketQA model uses additional augmented data beyond MS Marco, whereas the other compared models do not.
  • Comparison with contrastive pre-training: 0.49 average NDCG@10 is achieved by GTR-XXL, compared with 0.46 for CL-Pretrain.GTR models underperform CL-Pretrain with base size on over half the datasets, but GTR-Large and GTR-XXL show significant gains as model size increases.
  • Comparison setup: The comparison with Anonymous (2022) uses NDCG@10 across 15 reported BEIR datasets, contrasting CL with GTR models of different sizes.CL denotes contrastive learning pre-training on C4 and Wiki.
  • BEIR retrieval results: Scaling up dual encoders yields significant gains on the BEIR benchmark in Recall@100, similar to the NDCG@10 results.Table 8 marks the best result for each dataset in bold.
Loading 2112.07899v1…