BERT-QE: Contextualized Query Expansion for Document Re-ranking

1 Introduction

• PRF-based expansion can pollute queries with non-relevant information, and existing methods are described as insufficiently equipped to assess the relevance of expansion information.The cited approaches select tokens, topical terms, concepts, or documents, while their selection and re-weighting can rely on scalar counts.Source: §1 · p. 1
• BERT's ability to identify relevant document chunks offers a more flexible expansion unit than tokens, short concepts, or whole documents.The motivation draws on prior work applying contextualized BERT models to ranking and identifying relevant chunks.Source: §1 · p. 1
• BERT-QE uses a contextualized BERT model to select relevant information from feedback documents for query expansion.The authors report significant improvements over BERT-Large on Robust04 and GOV2, including variants with different computational costs.Source: §1 · pp. 1–2

2.1 Overview

• A fine-tuned BERT model, initialized from an MS MARCO checkpoint, is used to re-rank documents and supply first-phase relevance scores.The model is fine-tuned on a target dataset using query-passage pairs selected from the MS MARCO-trained model's top-ranked passages.Source: §2.1 · p. 2

2.2 Selecting Chunks for Query Expansion

• Overlapping chunks are formed from top-ranked feedback documents, scored individually by BERT, and the top-k_c chunks are selected as expansion evidence.Neighboring chunks overlap by up to m/2 words; the selected chunks distill feedback information for the next phase.Source: §2.2 · pp. 2–3, Figure 1 · p. 3

2.3 Final Re-ranking using Selected Chunks

• For each candidate document, BERT evaluates the selected chunks and aggregates their scores using softmax weights based on query-chunk relevance.The aggregate chunk-based relevance is combined with original query-document relevance, with α controlling their relative importance.Source: §2.3 · p. 3
• Different BERT variants may be used in different phases, such as using BERT-Small for chunk selection and BERT-Large for the other phases.The experiments consider the impact of using different variants across phases.Source: §2.3 · p. 3

3 Experimental Setup

• Code, cross-validation partitions, initial-ranking result files, and trained models are available online.Source: §3 · p. 3

3.1 Dataset and Metrics

• Robust04 (249 queries) and GOV2 (150) are evaluated with P@20, NDCG@20, MAP@100, and MAP@1000; rankings interpolate DPH+KL scores.Source: §3.1 · pp. 3–4

3.2 Initial Ranking

• DPH+KL generates the initial ranking by applying DPH after Rocchio query expansion with Kullback-Leibler divergence.Source: §3.2 · p. 4

3.3 Models in Comparisons

• Comparisons include Rocchio and RM3 expansion, BERT MaxP, SNRM, CEDR-KNRM, and Birch(MS→MB).Source: §3.3 · p. 4

3.4 Variants of BERT

• BERT variants differ in hidden-layer count, embedding size, and attention heads; suffixes identify the model used in each phase.Source: §3.4 · pp. 4–5, Table 1 · p. 5

3.5 Implementation of BERT-QE

• Implementation uses 10 feedback documents, selects 10 chunks of 10 words, and re-ranks the top 1,000 documents in phases one and three.Source: §3.5 · p. 5

3.6 Training

• Training concatenates each query with its most relevant passage; fivefold cross-validation uses three partitions for training, one for validation, and one for testing.Models train for two epochs, and reported performance averages results across folds.Source: §3.6 · p. 5

3.7 Computation of FLOPs

• FLOPs measure model computational complexity and include computations across all three BERT-QE phases.Source: §3.7 · p. 5

4.1 Results for BERT-QE-LLL

• BERT-QE-LLL significantly outperforms Birch(MS→MB) on Robust04 at the 0.01 level, while the fine-tuned BERT-Large achieves comparable results to Birch(MS→MB).The comparison is for title queries; significance tests against other models are omitted because their result rankings were unavailable.Source: §4.1 · pp. 5–6, Table 2 · p. 6
• BERT-QE-LLL requires 11.19x BERT-Large’s computations per query: 2603T FLOPs versus 232.6T, largely because it uses BERT-Large in all three phases.The authors note that parallelizing phases may reduce inference time.Source: §4.1 · p. 6

4.2 Employing Smaller BERT Variants in BERT-QE

• On Robust04, smaller phase-two chunk selectors all significantly outperform BERT-Large at the 0.01 level; on GOV2, Tiny, Small, and Medium selectors still outperform it, with metrics decreasing in most cases.BERT-QE-LML can achieve slightly higher results than BERT-QE-LLL on Robust04; smaller selectors use lower FLOPs.Source: §4.2 · p. 6

4.1 Results for BERT-QE-LLL

• On Robust04, BERT-QE-LLL improves over BERT-Large by 2.5% on NDCG@20 and 3.3% on MAP@1K; on GOV2, it significantly improves on all reported metrics.The reported Robust04 gains cover shallow and deep pools.Source: §4.1 · p. 6, Table 3 · p. 7

4.2 Employing Smaller BERT Variants in BERT-QE

• For final re-ranking, smaller models reduce effectiveness on Robust04, but some variants still outperform BERT-Large; BERT-QE-LMT uses 3% more FLOPs for consistent significant shallow-metric gains across both datasets.BERT-QE-LLS significantly outperforms BERT-Large on shallow and deep metrics across both datasets with 30% more FLOPs.Source: §4.2 · p. 7

5.1 First-round Re-ranker Ablation Analyses

• The first-round ranker provides the document score used in final re-ranking and the top-kd documents used to select candidate chunks.The selected chunks are used to compute rel(C, d).Source: §5.1 · p. 7
• Removing the BERT-Large document score or selecting chunks from DPH+KL documents prevents BERT-QE from outperforming BERT-Large.The ablations indicate that both the first-round ranking score and its feedback documents matter.Source: §5.1 · p. 8, Table 5 · p. 8

5.2 Hyper-parameter study

• On Robust04, kc = 10 and 20 perform similarly, kc = 5 reduces results, and m = 10 performs best among tested settings.Because phase-three computational cost is proportional to kc, the small performance gap makes kc = 10 a reasonable, robust choice.Source: §5.2 · p. 8, Figure 3 · p. 8

BERT for IR

• BERT ranking work includes passage-level scoring, incorporating BERT representations into neural rankers, cross-domain sentence evidence, and multi-stage ranking.Examples include Dai and Callan (2019), MacAvaney et al. (2019), Yilmaz et al. (2019), and Nogueira et al. (2019a).Source: §6 · p. 8

Query expansion

• Query expansion research includes feedback-based keyword, concept, embedding, and neural approaches, while Doc2query-family methods expand documents for short-text retrieval.The passage notes memory consumption as a reason document expansion methods are limited to short-text retrieval tasks.Source: §6 · p. 9
• BERT-QE uses contextualized BERT to select and incorporate relevant information from feedback for query expansion.This contrasts with prior approaches described as using language models, topical terms, concepts, or feedback documents directly.Source: §6 · p. 9

7 Conclusion

• BERT-QE significantly outperforms BERT-Large on Robust04 and GOV2 with up to 30% extra computational cost.The conclusion characterizes the extra cost as relatively small.Source: §7 · p. 9
• BERT-QE is proposed to select relevant information for query expansion.The authors identify combining it with pre-computation techniques as future work to improve efficiency.Source: §7 · p. 9

A.1 Interpolation Parameters in BERT-QE

• The interpolation coefficients α and β are tuned on validation sets by grid search over (0,1) with stride 0.1, selecting the models with best nDCG@20.Table 6 reports validation results and the chosen parameters for BERT-QE-LLL.Source: §A.1 · p. 11, Table 6 · p. 11
• The final score interpolates the log re-ranking score M(q, d) with the initial-ranking score I(q, d), using β.The equation is final score = β · log(M(q, d)) + (1 −β) · I(q, d).Source: §A.1 · p. 11
• Table 7 lists the number of parameters in the BERT variants.Source: Table 7 · p. 11

• Table 7 lists the parameter counts of the different BERT variants used in BERT-QE.Source: §A.2 · p. 11