summaries: This paper proposes a new class of surrogates called PiRank for ranking metrics optimization. In particular, it uses temperature-controlled relaxation to the permutation matrix to derive a continuous and differentiable sorting operator. Due to the large size of the ranking lists in practical applications, this paper also proposes a divide-and-conquer method to reduce the complexity of metric computation. Experiment results show that the proposed methods can achieve promising results on two popular benchmarks.   This paper proposes PiRank, a new neural learning-to-rank (LTR) approach based on differentiable sorting operators. The authors present empirical results against two widely used LTR data sets (MSLR-WEB30K and Yahoo! C14). 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. The authors provide a new scalable listwise loss for ranking called PiRank.  The loss is defined as 1 - a differentiable relaxation of NDCG metric with temperature-controlled inspired by NeuralSort algorithm. In this setting, the difference with the NDCG is negligeable the temperature tends to 0.  Scalability is ensured by a divide-and-conquer strategy where the sorting relaxation is applied to sub-lists of smaller size and then propagate only on the top items from each sub-list to merge the sort.
abstractText: A key challenge with machine learning approaches for ranking is the gap between the performance metrics of interest and the surrogate loss functions that can be optimized with gradient-based methods. This gap arises because ranking metrics typically involve a sorting operation which is not differentiable w.r.t. the model parameters. Prior works have proposed surrogates that are loosely related to ranking metrics or simple smoothed versions thereof, and often fail to scale to real-world applications. We propose PiRank, a new class of differentiable surrogates for ranking, which employ a continuous, temperature-controlled relaxation to the sorting operator based on NeuralSort [1]. We show that PiRank exactly recovers the desired metrics in the limit of zero temperature and further propose a divideand-conquer extension that scales favorably to large list sizes, both in theory and practice. Empirically, we demonstrate the role of larger list sizes during training and show that PiRank significantly improves over comparable approaches on publicly available internet-scale learning-to-rank benchmarks.
authors.name: Robin Swezey
authors.name: Aditya Grover
authors.name: Bruno Charron
authors.name: Stefano Ermon
references.authors: Aditya Grover Eric Wang Aaron Zweig Stefano Ermon
references.title: Stochastic optimization of sorting networks via continuous relaxations
references.venue: In ICLR,
references.authors: Tie-Yan Liu
references.title: Learning to rank for information retrieval
references.venue: Foundations and Trends R in Information Retrieval,
references.authors: David Cossock Tong Zhang
references.title: Subset ranking using regression
references.venue: In COLT,
references.authors: Ping Li Qiang Wu Christopher J Burges
references.title: Mcrank: Learning to rank using multiple classification and gradient boosting
references.venue: In NeurIPS,
references.authors: Koby Crammer Yoram Singer
references.title: Pranking with ranking
references.venue: In NeurIPS,
references.authors: Amnon Shashua Anat Levin
references.title: Ranking with large margin principle: Two approaches
references.venue: In NeurIPS,
references.authors: Ralf Herbrich
references.title: Large margin rank boundaries for ordinal regression
references.venue: Advances in large margin classifiers,
references.authors: Yoav Freund Raj Iyer Robert E Schapire Yoram Singer
references.title: An efficient boosting algorithm for combining preferences
references.venue: JMLR, 4(Nov):933–969,
references.authors: Chris Burges Tal Shaked Erin Renshaw Ari Lazier Matt Deeds Nicole Hamilton Greg Hullender
references.title: Learning to rank using gradient descent
references.venue: In ICML,
references.authors: Christopher JC Burges
references.title: From ranknet to lambdarank to lambdamart: An overview
references.authors: Zhaohui Zheng Hongyuan Zha Tong Zhang Olivier Chapelle Keke Chen Gordon Sun
references.title: A general boosting method and its application to learning ranking functions for web search
references.venue: NeurIPS,
references.authors: Yunbo Cao Jun Xu Tie-Yan Liu Hang Li Yalou Huang Hsiao-Wuen Hon
references.title: Adapting ranking svm to document retrieval
references.venue: In SIGIR,
references.authors: Christopher J.C. Burges Robert Ragno Quoc Viet Le
references.title: Learning to rank with nonsmooth cost functions
references.venue: In NeurIPS,
references.authors: Mingrui Wu Yi Chang Zhaohui Zheng Hongyuan Zha
references.title: Smoothing DCG for learning to rank: A novel approach using smoothed hinge functions
references.venue: In CIKM,
references.authors: Zhe Cao Tao Qin Tie-Yan Liu Ming-Feng Tsai Hang Li
references.title: Learning to rank: from pairwise approach to listwise approach
references.venue: In ICML,
references.authors: Fen Xia Tie-Yan Liu Jue Wang Wensheng Zhang Hang Li
references.title: Listwise approach to learning to rank: theory and algorithm
references.venue: In ICML,
references.authors: Jun Xu Hang Li
references.title: Adarank: a boosting algorithm for information retrieval
references.venue: In SIGIR,
references.authors: Yisong Yue Thomas Finley Filip Radlinski Thorsten Joachims
references.title: A support vector method for optimizing average precision
references.venue: In SIGIR,
references.authors: Michael Taylor John Guiver Stephen Robertson Tom Minka
references.title: Softrank: optimizing non-smooth rank metrics
references.venue: In WSDM,
references.authors: Tao Qin Tie-Yan Liu
references.title: Introducing LETOR 4.0 Datasets
references.authors: Rama Kumar Pasumarthi Sebastian Bruch Xuanhui Wang Cheng Li Michael Bendersky
references.title: Tf-ranking: Scalable tensorflow library for learning-to-rank
references.authors: Tie-Yan Liu
references.title: Learning to rank for information retrieval
references.venue: Springer Science & Business Media,
references.authors: Mu Zhu
references.title: Recall, precision and average precision
references.venue: Department of Statistics and Actuarial Science, University of Waterloo, Waterloo,
references.authors: Kalervo Järvelin Jaana Kekäläinen
references.title: Cumulated gain-based evaluation of ir techniques
references.venue: ACM Transactions on Information Systems,
references.authors: Tao Qin Tie-Yan Liu Hang Li
references.title: A general approximation framework for direct optimization of information retrieval measures
references.venue: Information retrieval,
references.authors: Mathieu Blondel Olivier Teboul Quentin Berthet Josip Djolonga
references.title: Fast differentiable sorting and ranking
references.venue: In ICML,
references.authors: Olivier Chapelle Mingrui Wu
references.title: Gradient descent optimization of smoothed information retrieval metrics
references.venue: Information retrieval,
references.authors: Ryan Prescott Adams Richard S Zemel
references.title: Ranking via sinkhorn propagation
references.venue: arXiv preprint arXiv:1106.1925,
references.authors: Gonzalo Mena David Belanger Scott Linderman Jasper Snoek
references.title: Learning Latent Permutations with Gumbel-Sinkhorn Networks
references.venue: In ICLR,
references.authors: Marco Cuturi Olivier Teboul Jean-Philippe Vert
references.title: Differentiable Ranks and Sorting using Optimal Transport
references.venue: In NeurIPS,
references.authors: Chris Burges
references.title: From RankNet to LambdaRank to LambdaMART
references.venue: An Overview. JMLR,
references.authors: Sashank J. Reddi Rama Kumar Pasumarthi Aditya Krishna Menon Ankit Singh Rawat Felix X. Yu Seungyeon Kim Andreas Veit Sanjiv Kumar
references.title: Rankdistil: Knowledge distillation for ranking
references.venue: In AISTATS,
references.authors: Yining Wang Liwei Wang Yuanzhi Li Di He Wei Chen Tie Yan Liu
references.title: A theoretical analysis of NDCG ranking
references.venue: measures. JMLR,
references.authors: Sebastian Prillo Julian Eisenschlos
references.title: Softsort: A continuous relaxation for the argsort operator
references.venue: In ICML,
references.authors: Yujia Xie Hanjun Dai Minshuo Chen Bo Dai Tuo Zhao
references.title: Differentiable top-k operator with optimal transport
references.venue: arXiv preprint arXiv:2002.06504,
references.authors: Leonardo Rigutini Tiziano Papini Marco Maggini Franco Scarselli
references.title: Sortnet: Learning to rank by a neural preference function
references.venue: IEEE transactions on neural networks,
references.authors: Ziqiang Cao Furu Wei Li Dong Sujian Li Ming Zhou
references.title: Ranking with recursive neural networks and its application to multi-document summarization
references.venue: In AAAI,
references.authors: Chenyan Xiong Zhuyun Dai Jamie Callan Zhiyuan Liu Russell Power
references.title: End-to-end neural ad-hoc ranking with kernel pooling
references.venue: In SIGIR,
references.authors: Robin L Plackett
references.title: The analysis of permutations
references.venue: Applied Statistics, pages 193–202,
references.authors: R Duncan Luce
references.title: Individual choice behavior: A theoretical analysis
references.venue: Courier Corporation,
references.authors: Yoshua Bengio Nicholas Léonard Aaron Courville
references.title: Estimating or propagating gradients through stochastic neurons for conditional computation
references.venue: arXiv preprint arXiv:1308.3432,
references.authors: O. Tange
references.title: Gnu parallel - the command-line power tool
references.venue: ;login: The USENIX Magazine,
sections.heading: 1 Introduction
sections.text: 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 cannot 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].
sections.heading: 2 Background and Related Work
sections.text: 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.
sections.heading: 2.1 Surrogate Objectives for LTR
sections.text: 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.
sections.heading: 2.2 Ranking Metrics
sections.text: 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.
sections.heading: 3.1 Relaxed Ranking Metrics
sections.text: We denote an LTR model by f✓ (e.g., deep neural network) with parameters ✓. The model outputs a vector of L scores ŷ = f✓(q,x1, . . . ,xL) for a query q and L candidate elements {xi}Li=1. We first consider the NDCG target metric. In Eq. 6, the numerator DCG(y, ⇡̂) involves computing ⇡̂ = sort(ŷ) which is non-differentiable w.r.t. ✓. Let g denote the column vector of graded relevance scores such that gj = 2yj 1. We can then rewrite DCG(y, ⇡̂) as:
  DCG(y, ⇡̂) = LX
  j=1
  g⇡̂j
  log2(1 + j) =
  LX
  j=1
  [P⇡̂g]j log2(1 + j) . (7)
  To obtain the DCG@k objective, one can replace L with k in the sum. We omit the suffix @k in the following, assuming that k has been defined, potentially equal to L which would yield the full metric.
  Let bPsort(s)(⌧) denote a relaxation to the permutation matrix Psort(s) that can be used for differentiable sorting of an input score vector s, for some temperature parameter ⌧ > 0 such that the true matrix is recovered as ⌧ ! 0+. Since ⇡̂ = sort(ŷ), we can obtain a differentiable relaxation to DCG(y, ⇡̂):
  [DCG(y, ŷ, ⌧) = kX
  j=1
  [ bPsort(ŷ)(⌧)g]j log2(1 + j) . (8)
  Substituting this in the expression for NDCG in Eq. 6, we obtain the following relaxation for NDCG:
  \NDCG(y, ŷ, ⌧) = [DCG(y, ŷ, ⌧) DCG(y,⇡⇤)
  (9)
  where the normalization in the denominator does not depend on ✓ and can be computed exactly via regular sorting. Finally, we define the PiRank surrogate loss for NDCG as follows:
  `PiRank NDCG = 1 \NDCG(y, ŷ, ⌧) (10) which is bounded between 0 and 1 as is NDCG, and whose difference with the actual (1 NDCG) gets negligible as ⌧ ! 0+. Figure 2 illustrates the model architecture for the above objective. Similarly, we can derive a surrogate loss for the ARP metric in Eq. 4 as:
  ˆ̀ PiRank ARP(y, ŷ, ⌧) =
  Pk j=1[
  bPsort(ŷ)(⌧)y]jj Pk
  j=1 yj
  . (11)
sections.heading: 3.2 Example: Differentiability via NeuralSort
sections.text: Typically, relaxations to permutation matrices consider the Birkhoff polytope of doubly stochastic matrices. A doubly-stochastic matrix is a square matrix with entries in [0, 1] where every row and column sum to 1. In contrast, NeuralSort [1] is a recently proposed relaxation of permutation matrices in the space of unimodal row-stochastic matrices. A unimodal matrix is a square matrix with entries in [0, 1] such that the entries in every row sum to 1 (i.e. row-stochastic), but additionally enforce the constraint that the maximizing entry in every row should have a unique column index. See Figure 3 for an example of each type. Note that a unimodal matrix is not necessarily doubly-stochastic and vice versa. Permutation matrices are both doubly-stochastic and unimodal.
  In NeuralSort [1], a unimodal relaxation of the permutation matrix Psort(ŷ) can be defined as follows. Let Aŷ denote the matrix of absolute pairwise score differences with i, j-th entry given as [Aŷ]ij = |ŷi ŷj |. Then, the i-th row of the relaxed permutation matrix is:
  bP (NS)sort(ŷ)(⌧)i,· = softmax [((L+ 1 2i)ŷ Aŷ1)/⌧ ] (12)
  where 1 is the vector with all components equal to 1. Its unimodal property makes it particularly well-suited to extracting top-k items because, as seen in Figure 3, taking the maximizing elements of the first k rows yields exactly k items but may yield less in the case of a doubly-stochastic relaxation. However, the complexity to obtain the top-k rows in this formulation, even for k as low as 1, is quadratic in L as the full computation of Aŷ1 is required for the softmax operation in Eq. 12. This is prohibitive when L k, a common scenario, and motivates the introduction of a new relaxation with a more favorable complexity for top-k ranking.
sections.heading: 3.3 Scaling via Divide-And-Conquer
sections.text: Our PiRank losses only require the first k rows of the relaxed permutation matrix bPsort(ŷ). This is specific to the LTR setting in which only the top-ranked items are of interest, in contrast to the full sorting problem that requires the full matrix. In PiRank, we leverage this insight to construct a divide-and-conquer variant of differentiable sorting relaxations such as NeuralSort to reduce the complexity of the metric computation. Our proposed construction can be viewed as a relaxed and truncated multi-way merge sort algorithm with differentiable sorting relaxations as building blocks. In the following discussion, we use NeuralSort as our running example while noting that the analysis extends more generally to other differentiable relaxations as well.
  Data Structure Construction. Let L = b1b2 · · · bd be a factorization of the list size L into d positive integers. Using this factorization, we will construct a tree of depth d with branching factor bj at height j. Next, we split the L-dimensional score vector ŷ into its L constituent scalar values. We set these values as the leaves of the tree. See Figure 4 for an example. At every other level of the tree, we will merge values from the level below into equi-sized lists. Let {kj}dj=0 be sizes for the intermediate results at level j, such that k0 = 1 (leaves) and min(k, kj 1bj)  kj  kj 1bj for j 1 (explained below). Then, in an iterative manner for levels j = 1, . . . , d, the value of a node at height j are the top-kj scores given by the application of the NeuralSort operator on the concatenation of the values of its children. With kd = k, the root value thus obtained is a relaxation of the top-k scores in ŷ. The top-k rows of the relaxed permutation matrix bPsort(ŷ) yielding these scores are constructed by compounding the operations at each iteration.
  Computational Complexity. The intuition behind the favorable scaling is as follows. At step j, NeuralSort is applied on blocks of size kj 1bj as it merges bj sub-blocks of size kj 1. Obtaining the full sorted list of scores would require to keep all intermediate scores during the process, i.e., k
  (max) j = b1 · · · bj = kj 1bj for j 1. In the last step, the NeuralSort operator is applied on a list of size kd 1bd, equal to L in this case, so the overall complexity would be at least quadratic in L as explained previously. However, since only the top-k scores are desired, intermediate outputs can be truncated if larger than k. Full truncation corresponds to k(min)j = min(k, kj 1bj). Any choice k
  (min) j  kj  k (max) j is acceptable to recover the top-k scores, with larger values allowing more information to flow at the expense of a higher complexity. Choosing bj ⇡ L1/d and kj minimal, the list sizes bjkj 1 on which NeuralSort is applied at each step can thus be of the order of L1/dk, much smaller than L in the d > 1 and k ⌧ L scenario. Formally, let ⌧1, ⌧2, . . . , ⌧d be the relaxation temperatures at each height, with ⌧d = ⌧ and ⌧j  ⌧j+1 for j 2 {1, . . . , d 1}. Define the tensor Ŷ (0) by reshaping ŷ to shape (k0, b1, b2, . . . , bd), yielding components
  Ŷ (0) 1,i1,i2,...,id = ŷ1+Pdj=1(ij 1) Qj 1 l=1 bl , (13)
  with ij 2 {1, . . . , bj} and the first index is always 1 as k0 = 1. With the tree representation, the first tensor index is the position in the node value vector and the rest of the indices identify the node by the index of each branching starting from the root. For j 2 {1, 2, . . . , d}, recursively define the tensors Q̂(j), Ŷ (j) and P̂ (j) of respective shapes (kj , kj 1, bj , . . . , bd), (kj , bj+1, . . . , bd) and
  (kj , b1, . . . , bd) with components
  Q̂ (j) l,m,ij ,...,id
  = softmax h⇣ (kj 1bj + 1 2l)Ŷ (j 1)m,ij ,...,id kj 1X
  p=1
  bjX
  q=1
  Ŷ (j 1)m,ij ,ij+1,...,id Ŷ (j 1) p,q,ij+1,...,id ⌘ /⌧j i ,
  (14)
  Ŷ (j) l,ij+1,...,id =
  kj 1X
  p=1
  bjX
  q=1
  Q̂ (j) l,p,q,ij+1,...,id Ŷ (j 1) p,q,ij+1,...,id , (15)
  P̂ (j) l,i1,...,id =
  kj 1X
  m=1
  Q̂ (j) l,m,ij ,...,id P̂ (j 1) m,i1,i2,...,id , (16)
  with P̂ (0) = 1. Intuitively, Ŷ (j) holds the relaxed top-kj scores at height j and Ŷ (d) is the desired top-k score vector. The interpretation of the indices in the tree structure is as for Ŷ (0), illustrated in Figure 4. More importantly, we keep track of the relaxed sorting operation that yielded this output. Q̂
  (j) is the relaxed permutation matrix obtained by applying NeuralSort in Eq. 12 to Ŷ (j), while P̂ (j) compounds the relaxed permutation matrices obtained so far so it always maps from the initial list size. Finally, define the k ⇥ L matrix P̂ by reshaping the tensor P̂ (d), yielding components
  P̂l,1+ Pd j=1(ij 1) Qj 1 l=1 bl = P̂ (d)l,i1,...,id , (17)
  for ij 2 {1, . . . , bj}. The k rows of P̂ are used as the top-k rows of the relaxed sorting operator bPsort(ŷ)(⌧). This approach is equivalent to NeuralSort, yielding Eq. 12 for d = 1. Proof of convergence for ⌧ ! 0+ of this relaxation in the general case d 1 is presented in Appendix B. In the simple case where L = bd and we set bj = b, kj = min(k, bj) for all j 2 {1, . . . , d}, the complexity to compute P̂ and thus the PiRank losses is then O(L1+1/d + (d 1)k2L), which scales favorably in L if d > 1 and k = O(1). In the general case, the score list can be padded, e.g. to the power of 2 following L, such that the previous complexity holds for b = 2 and d = dlog2 Le, but other factorizations may yield lower complexity depending on L.
sections.heading: 4 Experiments
sections.text: We present two sets of experiments in this section: (a) benchmark evaluation comparing PiRank with other ranking based approaches on publicly available large-scale benchmark LTR datasets, and (b) ablation experiments for the design choices in PiRank.
sections.heading: 4.1 Benchmark Evaluation via TF-Ranking
sections.text: Datasets. To empirically test PiRank, we consider two of the largest open-source benchmarks for LTR: the MSLR-WEB30K3 and the Yahoo! LTR dataset C144. Both datasets have relevance scores on a 5-point scale of 0 to 4, with 0 denoting complete irrelevance and 4 denoting perfect relevance. We give extensive details on the datasets and experimental protocol in Appendix C.
  Baselines. We focus on neural network-based approaches and use the open-source TensorFlow Ranking (TFR) framework [21]. TFR consists of high-quality GPU-friendly implementations of several LTR approaches, common evaluation metrics, and standard data loading formats. We compare our proposed loss, PiRank-NCDG, with the following baselines provided by TensorFlow Ranking: Approximate NDCG Loss [25], Pairwise Logistic Loss (RankNet, Eq. 2), Pairwise Logistic Loss with lambda-NDCG weights (LambdaRank, Eq. 3), and the Softmax Loss. We also include NeuralSort, whose loss is the cross-entropy of the predicted permutation matrix. Of these methods, the Pairwise Logistic Loss (RankNet) is a pairwise approach while the others are listwise. While our scope is on differentiable ranking surrogate losses for training neural networks, other methods such as the tree-based LambdaMART [31] could potentially yield better results.
  Setup. All approaches use the same 3-layer fully connected network architecture with ReLU activations to compute the scores ŷ for all (query, item) pairs, trained on 100,000 iterations. The maximum list size for each group of items to score and rank is fixed to 200, for both training and testing. Further experimental details are deferred to Appendix C. We evaluate Ordered Pair Accuracy (OPA), Average Relevance Position (ARP), Mean Reciprocal Rank (MRR), and NDCG@k with k 2 {1, 3, 5, 10, 15}. We determine significance similarly to [32]. For each metric and dataset, the best performing method is determined, then a one-sided paired t-test at a 95% significance level is performed on query-level metrics on the test set to compare the best method with every other method.
  Results. Our results are shown in Table 1. Overall, PiRank shows similar or better performance than the baselines on 13 out of 16 metrics. It significantly outperform all baselines on NDCG@k for k 5 while Approx. NDCG is competitive on NDCG@k for k  3 and MRR. PiRank is also significantly superior on OPA and ARP metrics on Yahoo! C14 while NeuralSort is superior for MSLR-WEB30K.
sections.heading: 4.2 Ablation Experiments
sections.text: Temperature. The temperature hyperparameter ⌧ is used in PiRank to control the degree of relaxation. We experiment on several values (⌧ 2 {0.1, 1, 5, 10}) using the MSLR-WEB30K dataset and the experimental settings for ablation provided in Appendix C. Figure 5 demonstrates the importance of correctly tuning ⌧ . High values (⌧ > 1) speed up training, especially in the early regime, while low values induce large gradient norms which are unsuitable for training and lead to the loss stalling or even diverging. Another observation is that the relaxed metric 1 ˆ̀PiRank NDCG@k
  3https://www.microsoft.com/en-us/research/project/mslr/ 4https://webscope.sandbox.yahoo.com
  closely follows the value of the downstream metric NDCG@k as expected. We further experimented with a very high temperature value and an exponentially decreasing annealing schedule in Appendix D.
  Training List Size. The training list size parameter Ltrain determines the number of items to rank for a query during training. The setting is the same as for the temperature ablation experiment, but with training list sizes Ltrain 2 {10, 20, 40, 100} which we then evaluate on testing list sizes in the same range of values Ltest 2 {10, 20, 40, 100}. The dataset is again MSLR-WEB30K. Figure 6 exposes four patterns for NDCG@k. First, for a fixed Ltest and k, a larger Ltrain is always better. Second, for a fixed Ltest, we observe diminishing returns along k, as relative improvements decrease for all Ltrain. This observation is confounded by NDCG@k values growing larger with k, but the metric is always able to distinguish between ranking functions [33]. Third, for a fixed k, our returns along Ltest increase with Ltrain (except for Ltrain = 20 and k = 1). This means that the need for a larger Ltrain is more pronounced for larger values of Ltest. Fourth and last, the returns increase most dramatically with Ltrain when Ltest k (top left), a common industrial setting. Values for NDCG@k, MRR, OPA, ARP are provided in Appendix D. For MRR, using a larger Ltrain is always beneficial regardless of Ltest, but not always for OPA and ARP.
  Depth. A main advantage of the PiRank loss is how it can scale to very large training list sizes Ltrain. This setting is difficult to come across with traditional LTR datasets, which are manually annotated, but occurs frequently in practice. One example is when the relevance labels are obtained from implicit signals such as clicks or purchases in recommendation systems. In this case, an LTR model is used to re-rank a list of candidates generated by another, simpler, model choosing among all possible items those potentially relevant to a query or context. An LTR model capable of handling very large lists can reduce the impact of errors made by the simpler candidate generation step, moving to the top an item lowly ranked at first that would have been cut off from a smaller list. To test the complexity shown in Section 3.2 in extreme conditions, we create a synthetic dataset as described in Appendix E. Figure 7 shows how the training time for depth d = 3 scales much more favorably than for d = 1, following their respective time complexities of O(L1+1/3) and O(L2).
sections.heading: 5 Summary and Limitations
sections.text: We proposed PiRank, a novel class of surrogate loss functions for Learning-To-Rank (LTR) that leverages a continuous, temperature-controlled relaxation to the sorting operator [1] as a building block. This allows us to recover exact expressions of the commonly used non-differentiable ranking metrics in the limit of zero temperature, which we proved in particular for the NDCG metric. Crucially, we proposed a construction inspired by the merge-sort algorithm that permits PiRank to scale to very large lists.
  In our experiments on the largest publicly available LTR datasets, we observed that PiRank has superior or similar performance with competing methods on the MSLR-WEB30K and Yahoo! C14 benchmarks on 13/16 ranking metrics and their variants.
  As future work, we would like to explore other recent relaxations of the sorting operator [26, 34, 35] as a building block for the PiRank framework. Further, as ranking is a core component of modern day technology and influences everyday decision making pipelines involving vulnerable populations, care needs to be taken that our proposed systems are extended to account for biases and fairness criteria when deployed in real world settings.
sections.heading: 6 Acknowledgements
sections.text: Robin Swezey and Bruno Charron were supported by Rakuten, Inc. Stefano Ermon is supported in part from NSF (#1651565, #1522054, #1733686), ONR (N0001419-1-2145), AFOSR (FA9550-19-1-0024) and Bloomberg.
title: PiRank: Scalable Learning To Rank via Differentiable Sorting