[Paper Review] Multi-Stage Document Ranking with BERT

Rodrigo Nogueira, Wei Yang+2 more|arXiv (Cornell University)|Oct 31, 2019

Topic Modeling42 references199 citations

TL;DR

The paper proposes monoBERT (pointwise) and duoBERT (pairwise) within a three-stage multi-stage ranking architecture to balance effectiveness and latency for document ranking, demonstrated on MS MARCO and TREC CAR.

ABSTRACT

The advent of deep neural networks pre-trained via language modeling tasks has spurred a number of successful applications in natural language processing. This work explores one such popular model, BERT, in the context of document ranking. We propose two variants, called monoBERT and duoBERT, that formulate the ranking problem as pointwise and pairwise classification, respectively. These two models are arranged in a multi-stage ranking architecture to form an end-to-end search system. One major advantage of this design is the ability to trade off quality against latency by controlling the admission of candidates into each pipeline stage, and by doing so, we are able to find operating points that offer a good balance between these two competing metrics. On two large-scale datasets, MS MARCO and TREC CAR, experiments show that our model produces results that are either at or comparable to the state of the art. Ablation studies show the contributions of each component and characterize the latency/quality tradeoff space.

Motivation & Objective

• Motivate and enable end-to-end document ranking with BERT-based re-rankers in a multi-stage pipeline.
• Introduce monoBERT (pointwise) and duoBERT (pairwise) as stages H1 and H2.
• Show how candidate admission and stage-wise processing trade quality for latency in a production-style system.
• Demonstrate improvements over BM25 baselines on MS MARCO and TREC CAR.
• Analyze component contributions via ablations and characterize latency–quality tradeoffs.

Proposed method

• Three-stage ranking architecture with H0 (BM25 bag-of-words retrieval), H1 (monoBERT binary relevance classifier), and H2 (duoBERT pairwise re-ranker).
• monoBERT uses CLS embedding from query and candidate to output relevance probability s_i; top-k1 candidates are kept.
• duoBERT uses query, candidate i, and candidate j as input to BERT to output pairwise relevance p_{i,j}; aggregate p_{i,j} with Sum/Binary/Min/Max/Sample to get final s_i.
• Pre-training: initial general BERT pre-training, then target-corpus pre-training (TCP) on task corpus (MS MARCO or Wikipedia subset for CAR).
• Training: TPU v3, balanced batches of relevant/non-relevant passages, cross-entropy loss for monoBERT and a pairwise loss for duoBERT; fine-tuning on MS MARCO and CAR datasets.

Experimental results

Research questions

• RQ1 Can monoBERT and duoBERT improve ranking effectiveness over BM25 baselines in a multi-stage pipeline?
• RQ2 What is the latency–quality tradeoff when varying k0 (H0 input candidates) and k1 (H1 outputs) in monoBERT/duoBERT pipelines?
• RQ3 Does target-corpus pre-training (TCP) improve performance over out-of-domain pre-training for BERT in neural ranking?
• RQ4 How do different aggregation methods (Sum, Binary, Min, Max, Sample) affect duoBERT’s final ranking?
• RQ5 How do monoBERT and duoBERT perform on MS MARCO versus TREC CAR, and why do gains differ across datasets?

Key findings

• monoBERT significantly outperforms BM25 baselines on MS MARCO and CAR.
• Adding duoBERT (with Sum or Binary aggregation) yields additional gains over monoBERT alone on both datasets.
• The Max aggregation method underperforms and is discarded; Sum and Binary provide strongest gains in MS MARCO and CAR respectively.
• Target-corpus pre-training (TCP) further improves MRR@10 on MS MARCO by about 0.8 points.
• The approach achieves results at or near state of the art on both MS MARCO and CAR datasets, with clear latency–quality tradeoffs demonstrated.