In defense of dual-encoders for neural ranking

Transformer-based models such as BERT have proven successful in information retrieval problem , which seek to identify relevant documents for a given query . There are two broad flavours of such models : cross-attention ( CA ) models , which learn a joint embedding for the query and document , and dual-encoder ( DE ) models , which learn separate embeddings for the query and document . Empirically , CA models are often more accurate , which has motivated a series of works seeking to bridge this gap . However , a fundamental question remains less explored : does this performance gap reflect a limitation in the capacity of DE models , or in the training of such models ? In this paper , we study this question , with three contributions . First , we establish theoretically that with a sufficiently large encoder size , DE models can capture a broad class of scores without cross-attention . Second , we show empirically that on real-world problems , the gap between CA and DE models may be due to the latter overfitting to the training set . To mitigate this behaviour , we propose a distillation strategy that focuses on preserving the ordering amongst documents , and confirm its efficacy on benchmark neural re-ranking datasets . 1 Introduction : transformer-based neural ranking . Information retrieval ( Mitra & Craswell , 2018 ) is the classic problem of identifying relevant documents for a given query . Typically , such retrieval is performed in a two-step manner ( Matveeva et al. , 2006 ) : one uses an initial model to efficiently retrieve a candidate set of documents for a query , and then uses a second model to re-rank these candidates . The retrieval stage is generally implemented by a model with low inference cost , for which candidate document generation is tractable ; if this candidate set is reasonably small , it is feasible to use more complex models in the re-ranking stage . Transformer-based models such as BERT ( Devlin et al. , 2019 ) have proven successful in both the retrieval and re-ranking stages . Such neural ranking models have two flavours : dual-encoder ( DE ) models ( Lee et al. , 2019 ; Chang et al. , 2020 ; Karpukhin et al. , 2020 ; Xiong et al. , 2021 ; Luan et al. , 2021 ) , which learn separate ( factorised ) embeddings for the query and document ; and cross-attention ( CA ) models ( Nogueira & Cho , 2019 ; Dai & Callan , 2019 ; Yilmaz et al. , 2019 ; MacAvaney et al. , 2019 ; Gao et al. , 2020 ) , which learn a joint embedding for the query and document . Only DE models are applicable for retrieval , as they admit efficient nearest neighbour search ( Guo et al. , 2020 ; Johnson et al. , 2021 ) to identify candidate documents . Both CA and DE models are applicable for re-ranking ; however , empirically , CA models have proven more accurate ( Hofstätter et al. , 2020a ) . Does the re-ranking performance gap between CA and DE models reflect a limitation in the inherent capacity of the latter ’ s factorised representation , or in the training of DE models ? These questions are of conceptual interest , but are also practically relevant : DE models allow for more efficient inference owing to the ability to pre-compute document embeddings , which enables fast similarity search when presented with a candidate query ( Khattab & Zaharia , 2020 ; Hofstätter et al. , 2020b ) . Several works have explored means of improving DE models , including by changing1 the scoring layer ( Khattab & Zaharia , 2020 ; MacAvaney et al. , 2020 ; Hofstätter et al. , 2020b ) , and by distilling predictions from a CA model ( Lu et al. , 2020 ; Izacard & Grave , 2020 ; Yang & Seo , 2020 ; Hofstätter et al. , 2020a ; Miech et al. , 2021 ) . However , the root cause of the gap between CA and DE models remains elusive . In this paper , we study these questions , with the following contributions : 1Unless otherwise stated , we use “ DE ” to mean a dual-encoder with dot-product scoring per Equation 2 . ( i ) we establish theoretically that with a sufficiently large embedding dimension ( and under mild assumptions ) , any continuous ground-truth scores can be modelled by generic DE models ( Proposition 1 ) , and in particular by sufficiently deep transformer-based DEmodels ( Proposition 2 ) . Thus , in principle , there is no fundamental restriction in using DE versus CAmodels for re-ranking . ( ii ) we show empirically ( Figure 1 ) that on real-world problems , CA and DE models can achieve similar training performance , but DE models may do worse on test data due to overfitting ( §3.2 ) . Thus , the DE models may suffer due to poorer generalisation ability rather than capacity . ( iii ) to mitigate the above , we propose a distillation strategy focussing on mimicking the teacher ’ s ordering amongst documents ( §4 ) . This generalises the recent margin MSE loss ( Hofstätter et al. , 2020a ) ( §4.1 ) , and justifies the utility of softmax cross-entropy based distillation for re-ranking ( Proposition 3 , §4.2 ) . We confirm the efficacy of this strategy on the re-reranking benchmarks MSMARCO-Passage ( Nguyen et al. , 2016 ) and Natural Questions ( Kwiatkowski et al. , 2019 ) . Overall , our results aim to shed light on the fundamental differences between CA and DE models , and give a simple yet effective distillation strategy to boost the performance of the latter . 2 Background and notation . Retrieval problems . Suppose we have query space Q and document space D. For a given query q ∈ Q , the goal of information retrieval is to identify a set of relevant documents Drel ( q ) ⊆ D ( Mitra & Craswell , 2018 ) . This is typically achieved by learning a scorer s : Q × D → R that predicts the relevance of a query and document , and simply selecting the top-k highest scoring documents for a given query . Typically , this scorer is itself implemented in two phases ( Matveeva et al. , 2006 ; Nogueira & Cho , 2019 ; Chang et al. , 2020 ; Miech et al. , 2021 ; Hofstätter et al. , 2021 ) . In the retrieval phase , one identifies an initial candidate set of documents via an initial scorer sret . In the re-ranking phase , one re-scores only these candidate documents via a distinct scorer srrk to obtain Drel ( q ) . We may thus regard s : Q×D→ R as a suitable composition of both sret and srrk . To learn a scorer , suppose we have supervision { ( qn , dn , yn ) } Nn=1 , where for each query qn we have a list of K associated documents dn . = ( dn1 , . . . , dnK ) ∈ DK with ground-truth relevance labels yn . = ( yn1 , . . . , ynK ) ∈ { 0 , 1 } K . Typically , dn comprises a small subset of D ( i.e. , K |D| ) , since it is rarely feasible to collect relevance labels for all documents . Consider a parameterised class of scorers S .= { s ( · , · ; θ ) : Q × D → R | θ ∈ Θ } . For any given θ ∈ Θ , let sn ( θ ) . = ( s ( qn , dn1 ; θ ) , . . . , s ( qn , dnK ; θ ) ) ∈ RK denote the vector of model scores on the provided documents . One may then learn a scorer s via minimising R ( θ ) .= 1N ∑ n∈ [ N ] ` ( yn , sn ( θ ) ) for loss ` : RK × RK → R+ , such as the mean square error ` mse ( yn , sn ( θ ) ) . = ∑ k∈ [ K ] ( ynk − snk ( θ ) ) 2 ; and the softmax cross-entropy ( assuming for simplicity 1 > y = 1 ) : for temperature τ > 0 , ` ce ( yn , sn ( θ ) ) . = − ∑ k∈ [ K ] ynk · log pnk ( θ ) where pnk ( θ ) . = exp ( τ−1 · snk ( θ ) ) ∑ k′∈ [ K ] exp ( τ −1 · snk′ ( θ ) ) . ( 1 ) To instantiate either objective , one missing ingredient is the precise parameterisation of S. We now describe one possible parameterisation , given by a transformer model . Cross-attention and dual-encoders . Transformers ( Vaswani et al. , 2017 ) are sequence-to-sequence models with good empirical performance on language tasks . Suppose we have a sequence ( x1 , . . . , xL ) of L tokens ( e.g. , words in a language ) , with X .= ( x1 , x2 , . . . , xL ) ∈ RD×L being a corresponding sequence of token embeddings ( e.g. , word embeddings ) . A transformer encoder maps this sequence to Z .= ( z1 , z2 , . . . , zL ) ∈ RD×L , via a composition of attention and feedforward layers . The former is a mapping attn : RD×L → RD×L , which in its simplest ( “ single head ” ) form is attn ( X ) = X + WVXA , where A . = σ ( X > WQWKX/ √ D ) for softmax σ ( · ) and weights WQ , WK , WV . Intuitively , this embeds a token via a weighted sum of all tokens ’ embeddings . The weighting is given by the attention matrix A ∈ RL×L , which is a learned similarity measure between pairs of tokens . BERT ( Devlin et al. , 2019 ) is a canonical instantiation of the transform encoder , coupled with effective pre-training tasks shown to improve downstream performance in small-sample regimes . Transformers can readily score ( query , document ) relevance : given a pair ( q , d ) ∈ Q×D of tokenised queries and documents , one may pass their concatenation through a transformer encoder T . The resulting embedding may be pooled ( e.g. , using the [ CLS ] token ( Devlin et al. , 2019 ) ) to yield s ( q , d ) = w > pool ( T ( concat ( q , d ) ) ) for weights w ∈ RD . Such cross-attention ( CA ) models apply attention layers on the queries and documents jointly , which intuitively allow for rich interactions . A distinct strategy is to separately embed queries and documents ( e.g. , via BERT encoders ) , and score them with their dot-product ( Lee et al. , 2019 ; Chang et al. , 2020 ; Karpukhin et al. , 2020 ; Zhan et al. , 2020 ; Ma et al. , 2021 ; Xiong et al. , 2021 ; Luan et al. , 2021 ; Qu et al. , 2021 ; Zhan et al. , 2021a ) : s ( q , d ) = pool ( T ( q ) ) > pool ( T ( d ) ) . ( 2 ) Such dual-encoder ( DE ) models have their roots in Siamese networks for image processing ( Bromley et al. , 1993 ) , and have received recent interest in a range of language problems ( Reimers & Gurevych , 2019 ; Gillick et al. , 2019 ; Karpukhin et al. , 2020 ) . A salient feature of Equation 2 is that one can index the query and document embeddings . This implies that DE models can have significantly lower the inference cost compared to CA models , which makes them appealing to use for both the retrieval and re-ranking stages . Indeed , DE models are widely used for retrieval , owing to their amenability to fast approximate nearest neighbour search ( Guo et al. , 2020 ; Johnson et al. , 2021 ) . Unfortunately , for re-ranking , one typically observes a large gap between DE and CA performance ( see §3.2 ) . Cross-attention to dual-encoder distillation . One strategy to bridge the performance gap between CA and DE models is distillation ( Ba & Caruana , 2014 ; Hinton et al. , 2015 ) . Suppose we have supervision { ( qn , dn , tn , yn ) } Nn=1 , where tn ∈ RK are the scores from a “ teacher ” ( e.g. , CA ) model . Then , we may minimise Rdist ( θ ) . = 1N ∑ n∈ [ N ] ` ( tn , sn ( θ ) ) , noting that we use the teacher supervision tn rather than “ one-hot ” labels yn ; one may also combine the teacher and one-hot labels . Two canonical instantiations are logit matching ( Ba & Caruana , 2014 ) , which uses the the mean square error ` mse ( tn , sn ( θ ) ) between the teacher and student scores , and probability matching ( Hinton et al. , 2015 ) , which uses the softmax cross-entropy between the teacher and student softmax probabilities , i.e. , ` ( tn , sn ) = ` ce ( p teach n , sn ) where pteachnk . = exp ( τ −1·tnk ) ∑ k′∈ [ K ] exp ( τ −1·tnk′ ) for temperature τ > 0 . Empirically , distillation has been shown to at least partially bridge the chasm between CA and DE performance ( Lu et al. , 2020 ; Izacard & Grave , 2020 ; Yang & Seo , 2020 ; Hofstätter et al. , 2020a ; Miech et al. , 2021 ) .

Summary: This paper aims to narrow the performance gap between cross-attention BERT and dual-encoder BERT for re-ranking task. The authors empirically and theoretically analyze the underlying reasons of the performance gap. The gap could be mitigated by the proposed knowledge distillation method, where the original cross-attention BERT model acts as the “teacher” and the more efficient dual-encoder BERT model is used as student. Comprehensive experiments confirmed the effectiveness of the proposed KD method for various re-ranking tasks.