PiRank: Scalable Learning To Rank via Differentiable Sorting

1 Introduction . The goal of Learning-To-Rank ( LTR ) models is to rank a set of candidate items for any given search query according to a preference criterion [ 2 ] . The preference over items is specified via relevance labels for each candidate . The fundamental difficulty in LTR is that the downstream metrics of interest such as normalized discounted cumulative gain ( NDCG ) and average relevance position ( ARP ) depend on the ranks induced by the model . These ranks are not differentiable with respect to the model parameters , so the metrics can not be optimized directly via gradient-based methods . To resolve the above challenge , a popular class of LTR approaches map items to real-valued scores and then define surrogate loss functions that operate directly on these scores . Surrogate loss functions , in turn , can belong to one of three types . LTR models optimized via pointwise surrogates [ 3–6 ] cast ranking as a regression/classification problem , wherein the labels of items are given by their individual relevance labels . Such approaches do not directly account for any inter-dependencies across item rankings . Pairwise surrogate losses [ 7–14 ] can be decomposed into terms that involve scores of pairs of items in a list and their relative ordering . Finally , listwise surrogate losses [ 15–19 ] are defined with respect to scores for an entire ranked list . For many prior surrogate losses , especially those used for listwise approaches , the functional form is inspired via downstream ranking metrics , such as NDCG . However , the connection is loose or heuristically driven . For instance , SoftRank [ 14 , 19 ] introduces a Gaussian distribution over scores , which in turn defines a distribution over ranks and the surrogate is the expected NDCG w.r.t . this rank distribution . ⇤This work was done prior to joining Amazon . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . We propose PiRank , a listwise approach where the scores are learned via deep neural networks and the surrogate loss is obtained via a differentiable relaxation to the sorting operator . In particular , we choose as building block the temperature-controlled NeuralSort [ 1 ] relaxation for sorting and specialize it for commonly used ranking metrics such as NDCG and ARP . The resulting training objective for PiRank reduces to the exact ranking metric optimization in the limit of zero temperature and trades off bias for lower variance in the gradient estimates when the temperature is high . Furthermore , PiRank scales to real-world industrial scenarios where the size of the item lists is very large but the ranking metrics of interest are determined by only a small set of top ranked items . Scaling is enabled by a novel divide-and-conquer strategy akin to merge sort , where we recursively apply the sorting relaxation to sub-lists of smaller size and propagate only the top items from each sub-list for further sorting . Empirically , we benchmark PiRank against 5 competing methods on two of the largest publicly available LTR datasets : MSLR-WEB30K [ 20 ] and Yahoo ! C14 . We find that PiRank is superior or competitive on 13 out of 16 ranking metrics and their variants , including 9 on which it is significantly superior to all baselines , and that it is able to scale to very large item lists . We also provide several ablation experiments to understand the impact of various factors on performance . To the best of our knowledge , this work is the first to analyze the importance of training list size on an LTR benchmark . Finally , we provide an open-source implementation2 based on TensorFlow Ranking [ 21 ] . 2 Background and Related Work . The LTR setting considers a finite dataset consisting of n triplets D = { qi , { xi , j } Lj=1 , { yi , j } Lj=1 } ni=1 . The i-th triplet consists of a query qi 2 Q , a list of L candidate items represented as feature vectors xi , j 2 X , and query-specific relevance labels yi , j for each item j . The relevance labels yi , j can be binary , ordinal or real-valued for more fine-grained relevance . For generality , we focus on the real-valued setting . Given a training dataset D , our goal is to learn a mapping from queries and itemsets to rankings . A ranking ⇡ is a list of unique indices from { 1 , 2 , . . . , L } , or equivalently a permutation , such that ⇡j is the index of the item ranked in j-th position . Without loss of generality , we assume lower ranks ( starting from 1 ) have higher relevance scores . This is typically achieved by learning a scoring function f : Q ⇥ XL ! RL that maps a query context and list of candidate items to L scores . At test time , the candidate items are ranked by sorting their predicted scores in descending order . The training of f itself can be done by a suitable differentiable surrogate objective , which we discuss next . 2.1 Surrogate Objectives for LTR . In this section , we briefly summarize prominent LTR approaches with a representative loss function for each category of pointwise , pairwise or listwise surrogate losses . We refer the reader to the excellent survey by [ 22 ] for a more extensive review . Omitting the triplet index , we denote the relevance labels vector as y 2 RL and an LTR model ’ s score vector obtained via the scoring function f as ŷ 2 RL . The simplest pointwise surrogate loss for ranking is the mean-squared error ( MSE ) between y and ŷ : ˆ̀ MSE ( y , ŷ ) = 1 L LX i=1 ( ŷi yi ) 2 . ( 1 ) As the example loss above shows , pointwise LTR approaches convert ranking into a regression problem over the relevance labels and do not account for the relationships between the candidate items . Pairwise approaches seek to remedy this by considering loss terms depending in the predicted scores of pairs of items . For example , the widely used RankNet [ 9 ] aims to minimize the number of inversions , or incorrect relative orderings between pairs of items in the predicted ranking . It does so by modeling the probability p̂i , i0 that the relevance of the i-th item is higher than that of the i0-th item as a logistic map of their score difference , for all candidate items i , i0 . The objective is then the 2https : //github.com/ermongroup/pirank cross-entropy : ˆ̀ RankNet ( y , ŷ ) = LX i=1 LX i0=1 ( yi > yi0 ) log p̂i , i0 ( 2 ) where ( · ) denotes the indicator function and p̂i , i0 is a function of ŷ . Pairwise approaches effectively model relationships between pairs of items and generally perform better than pointwise approaches , but still manifest limitations on downstream metrics which consider rankings in the full list and not just pairs . In fact , the larger the list of candidate items , the weaker these approaches tend to be : an error between the first and second item on a list is weighted the same in the RankNet loss as one between the last two items , despite the top items being of more importance in the LTR setting . Listwise approaches learn from errors on the complete list . LambdaRank [ 13 ] extends RankNet by assigning weights to every loss term from Eq . 2 : ˆ̀ LambdaRank ( y , ŷ ) = LX i=1 LX i0=1 ` NDCG ( i , i 0 ) log p̂i , i0 ( 3 ) with ` NDCG ( i , i0 ) the difference in the downstream metric NDCG ( defined below ) when swapping items i and i0 . 2.2 Ranking Metrics . Downstream metrics operate directly on the predicted ranking ⇡̂ ( obtained by sorting ŷ in descending order ) and the true relevance labels y . They differ from conventional metrics used for other supervised learning problems in explicitly weighting the loss for each item by a suitably choosen increasing function of its predicted rank . For example , relevance position ( RP ) [ 23 ] multiplies the relevance labels with linearly increasing weights , and normalizes by the total relevance score for the query : RP ( y , ⇡̂ ) = PL j=1 y⇡̂j jPL j=1 yj ( 4 ) Averaging RP across the predictions made for all the queries in the test set gives the average relevance position ( ARP ) metric . Lower ARP signifies better performance . A very common metric is the discounted cumulative gain ( DCG ) [ 24 ] . DCG computes the rescaled relevance of the j-th candidate by exponentiating its relevance label , and further divides it by the assigned log-ranking . This model incentivizes ranking models to focus on elements with higher graded relevance scores : DCG ( y , ⇡̂ ) = LX j=1 2y⇡̂j 1 log2 ( 1 + j ) ( 5 ) A more common variant NDCG normalizes DCG by the maximum possible DCG attained via the optimal ranking ⇡⇤ ( obtained by sorting y in descending order ) : NDCG ( y , ⇡̂ ) = DCG ( y , ⇡̂ ) DCG ( y , ⇡⇤ ) ( 6 ) Higher DCG and NDCG signify better performance . Their truncated versions DCG @ k and NDCG @ k are defined by replacing L with a cutoff k in Eq . 5 so metrics are computed on the top-k items . 3 Scalable and Differentiable Top-k Ranking via PiRank In PiRank , we seek to design a new class of surrogate objectives for ranking that address two key challenges with current LTR approaches . The first challenge is the gap between the downstream ranking metric of interest ( e.g. , NDCG , ARP ) that involve a non-differentiable sorting operator and the differentiable surrogate function being optimized . The second challenge concerns the scalability w.r.t . the size of the candidate list L for each query item . Larger list sizes are standard in industrial applications but present computational and memory challenges for current approaches during both training and test-time inference . Pairwise and listwise methods ( or hybrids ) typically scale quadratically in the list size L , the number of items to rank for each query . Combining surrogates for truncated metrics , such as LambdaRank in Eq . 3 with NDCG @ k has a reduced complexity of O ( kL ) but comes at the cost of vanishing gradient signal from relevant entries below k ( see Figure 1 for an illustration ) . Soft versions of the truncation metrics , such as Approximate NDCG Loss [ 25 ] , can learn from all items but again scale quadratically with L or do not take advantage of GPU acceleration [ 26 ] . As defined previously , a ranking ⇡ is a list of indices equivalent to a permutation of { 1 , 2 , . . . , L } . The set of possible rankings can thus be seen as the symmetric group SL , of size L ! . Every permutation ⇡ can be equivalently represented as a permutation matrix P⇡ , an L⇥ L matrix such that its ( i , ⇡i ) -th entry is 1 for all i 2 { 1 , 2 , · · · , L } and the remaining entries are all 0 . We define the sorting operator sort : RL ! SL as a map of an L-dimensional input vector to the permutation that corresponds to the descending ordering of the vector components . Prior work in relaxing the sorting operator is based on relaxation of its output , either in the form of rankings [ 19 , 26 , 27 ] , or permutation matrices [ 1 , 28–30 ] . PiRank first applies the latter kind of relaxations to the ranking problem by introducing a new class of relaxed LTR metrics , then introduces a new relaxation that is particularly suited to these metrics .

Summary: This paper tackles the problem of Learning To Rank. The idea is to relax ranking metrics (eg DCG), that are non differentiable with respect to the model parameters theta. Originally, these metrics measure how well the predicted ranking (induced by the scores predicted by the model) places high items with high true relevance labels. The relaxed metrics replace the permutation matrix in the original metrics by a relaxed, unimodal matrix, parametrized by a temperature.  This relaxation was proposed in Neuralsort, but is of quadratic complexity in the total number of items L at each time. In this paper, the authors propose a computationally efficient version of Neuralsort, truncated at the k first items, relying on a divide and conquer strategy. This approach is of much lower complexity, expressed in terms of L,k, and d the depth of the tree.  Finally, the method is evaluated on benchmark learning to rank datasets. It obtains similar or better performance than sota methods.