Shallow Cross-Encoders for Low-Latency Retrieval

Abstract

Transformer-based Cross-Encoders achieve state-of-the-art effectivness in text retrieval. However, Cross-Encoders based on large transformer models (such as BERT or T5) are computationally expensive and allow for scoring only a small number of documents within a reasonably small latency window. However, kee** search latencies low is important for user satisfaction and energy usage. In this paper, we show that weaker shallow transformer models (i.e. transformers with a limited number of layers) actually perform better than full-scale models when constrained to these practical low-latency settings, since they can estimate the relevance of more documents in the same time budget. We further show that shallow transformers may benefit from the generalised Binary Cross-Entropy (gBCE) training scheme, which has recently demonstrated success for recommendation tasks. Our experiments with TREC Deep Learning passage ranking querysets demonstrate significant improvements in shallow and full-scale models in low-latency scenarios. For example, when the latency limit is 25 ms per query, MonoBERT-Large (a cross-encoder based on a full-scale BERT model) is only able to achieve NDCG@10 of 0.431 on TREC DL 2019, while TinyBERT-gBCE (a cross-encoder based on TinyBERT trained with gBCE) reaches NDCG@10 of 0.652, a +51% gain over MonoBERT-Large. We also show that shallow Cross-Encoders are effective even when used without a GPU (e.g., with CPU inference, NDCG@10 decreases only by 3% compared to GPU inference with 50 ms latency), which makes Cross-Encoders practical to run even without specialised hardware acceleration.

A Effect of gBCE Training Scheme on Tiny BERT-Based Cross-Encoder on the MS MARCO Dev Set

Table 3 reports the effectiveness of a Tiny BERT model on a 6,980 queries sub-set of the MS MARCO dev set (dataset irds:msmarco-passage/dev/small in PyTerrier). The evaluation follows the scheme described in Sect. 4.3, with the exception of using MRR@10, which is the official metric for this queryset, instead of NDCG@10. As we can see from the table, the overall trends follow the observations in Sect. 4.3. In particular, an increased number of negatives is more important than the loss function; gBCE loss improves results with a small number of negatives but has a moderate effect when the number of negatives increases. However, we observe that overall gBCE in this experiment is better than BCE loss in 5 out of 6 cases. With 1 negative, the improvement over BCE loss is statistically significant. Overall, the combination of gBCE loss and 128 number of negatives provides a significant improvement of MRR@10, from 0.2942 to 0.3200 (+8.76%), compared to the “standard” training scheme with 1 negative and BCE loss. Note that this result is lower compared to the larger models – e.g. Nogueira et al. [27] achieved MRR@10 of 0.36 on this queryset with a BERT-Large model. Lower effectiveness compared to the full-scale models is an expected result, as we do not control for latency in this experiment. When latency is limited, shallow Cross-Encoders are more effective (see Fig. 1).

Table 3. Effect of the loss function and the number of negatives training on Tiny BERT-based Cross-Encoder MRR@10 on the MS MARCO dev set. Bold indicates the best result, and * indicates a statistically significant difference ($pvalue < 0.05$) compared to the baseline (BCE loss, one negative).