Learning-to-Rank with BERT in TF-Ranking

TL;DR

The paper tackles the problem of improving document and passage ranking by combining pretrained language models with learning-to-rank losses. It introduces TFR-BERT, which encodes query-document pairs with BERT and trains an LTR model via TF-Ranking using pointwise, pairwise, and listwise losses, with extensions to RoBERTa and ELECTRA. Key contributions include a practical encoding scheme, TPU-based training with list-wise losses, and extensive ensembling across losses and model families that achieves state-of-the-art re-ranking on MS MARCO and competitive full-ranking results. The results demonstrate that integrating ranking losses with strong LM representations yields robust, reproducible improvements without relying on additional indexing features. This approach offers a scalable, transferable framework for improving search relevance in real-world systems, validated by multiple leaderboard submissions and ablations.

Abstract

This paper describes a machine learning algorithm for document (re)ranking, in which queries and documents are firstly encoded using BERT [1], and on top of that a learning-to-rank (LTR) model constructed with TF-Ranking (TFR) [2] is applied to further optimize the ranking performance. This approach is proved to be effective in a public MS MARCO benchmark [3]. Our first two submissions achieve the best performance for the passage re-ranking task [4], and the second best performance for the passage full-ranking task as of April 10, 2020 [5]. To leverage the lately development of pre-trained language models, we recently integrate RoBERTa [6] and ELECTRA [7]. Our latest submissions improve our previously state-of-the-art re-ranking performance by 4.3% [8], and achieve the third best performance for the full-ranking task [9] as of June 8, 2020. Both of them demonstrate the effectiveness of combining ranking losses with BERT representations for document ranking.

Figures (1)

• Figure 1: An illustration of the TFR-BERT framework, in which a Learning-to-Rank model is constructed on top of the BERT representations of query-document pairs.