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: In information retrieval problems, there are two kinds of models: dual-encoder models and cross-attention models. Cross-attention models generally outperform dual encoder models by a large margin. Due to the latency requirement of real search systems, cross-attention models might not meet the requirements. So one ideal solution is to transfer the knowledge from cross-attention models to dual encoder models. This paper follows this line of research ideas by reducing the gap between the two models.   Firstly, the paper theoretically analyzes the effects with sufficiently large embedding dimensions. Secondly, the paper empirically verifies the theoretical results. Thirdly, the paper proposes a new method that further bridges the gap between the two methods. The writing is clear and well-structured. I like the idea and logic of the paper. I have some questions about the experiment part.