Source-linked AI summary
SetRank: Learning a Permutation-Invariant Ranking Model for Information Retrieval
Liang Pang, Jun Xu, Qingyao Ai, Yanyan Lan, Xueqi Cheng, Jirong Wen
TL;DR
Learning-to-rank needs models that capture cross-document context without making rankings depend on input order. SetRank uses self-attention to jointly encode document sets into permutation-invariant rankings, and it significantly outperformed traditional and neural baselines on three large-scale benchmarks.
Problem
Learning-to-rank requires a model that captures cross-document interactions while remaining permutation-invariant to document input order.
Method
SetRank jointly encodes document sets with stacked multi-head self-attention blocks and maps them to permutations.
Results
SetRank significantly outperformed traditional learning-to-rank and state-of-the-art neural IR baselines on three large-scale benchmarks.
Takeaways & Limitations
SetRank provides a ranking model that captures document interactions while remaining robust to input order and document-set size.
Takeaways & Limitations
The standard multi-head self-attention block is sensitive to input-set size and may not work when trained and tested with different document counts.
Abstract
from arXiv · showhide
In learning-to-rank for information retrieval, a ranking model is automatically learned from the data and then utilized to rank the sets of retrieved documents. Therefore, an ideal ranking model would be a mapping from a document set to a permutation on the set, and should satisfy two critical requirements: (1)~it should have the ability to model cross-document interactions so as to capture local context information in a query; (2)~it should be permutation-invariant, which means that any permutation of the inputted documents would not change the output ranking. Previous studies on learning-to-rank either design uni-variate scoring functions that score each document separately, and thus failed to model the cross-document interactions; or construct multivariate scoring functions that score documents sequentially, which inevitably sacrifice the permutation invariance requirement. In this paper, we propose a neural learning-to-rank model called SetRank which directly learns a permutation-invariant ranking model defined on document sets of any size. SetRank employs a stack of (induced) multi-head self attention blocks as its key component for learning the embeddings for all of the retrieved documents jointly. The self-attention mechanism not only helps SetRank to capture the local context information from cross-document interactions, but also to learn permutation-equivariant representations for the inputted documents, which therefore achieving a permutation-invariant ranking model. Experimental results on three large scale benchmarks showed that the SetRank significantly outperformed the baselines include the traditional learning-to-rank models and state-of-the-art Neural IR models.
1 INTRODUCTION
Learning-to-rank models must capture interactions among retrieved documents while remaining insensitive to their input order. SetRank addresses both requirements with self-attention over document sets and outperforms traditional and neural baselines on three benchmarks.
- Motivation: Traditional independent scoring models cannot capture cross-document interactions or local query context.They assign document scores separately under the probability ranking principle.
- Motivation: Sequential multivariate models capture context but are sensitive to input order and can degrade when the initial ranking is poor or disturbed.Their sequence-based architectures violate permutation invariance.
- SetRank: SetRank maps document sets of any size to permutations using stacked multi-head self-attention blocks.The model jointly learns document embeddings before a row-wise fully connected network generates ranking scores.
- SetRank: Self-attention lets SetRank model document interrelationships while learning a permutation-invariant ranking function.The model treats the retrieved documents as a whole rather than ranking them through a fixed sequence.
- Results: SetRank significantly outperformed traditional learning-to-rank and state-of-the-art neural IR baselines on three large-scale public benchmarks.Analyses also found robustness to input orders and document-set sizes.
2 RELATED WORK
Prior work spans conventional learning-to-rank, deep ranking models, and machine learning for sets. SetRank extends multivariate ranking with a self-attention architecture that operates directly on document sets without requiring an initial order.
- Learning-to-Rank: Learning-to-rank applies machine-learning techniques to ranking problems using query-document feature vectors and predicted document scores.Existing algorithms are commonly categorized as pointwise, pairwise, or listwise according to their loss functions.
- Learning-to-Rank: Earlier multivariate approaches jointly rank document lists by encoding local context or predicting rankings together.Examples use recurrent networks or slate optimization frameworks.
- Learning-to-Rank: SetRank differs by being permutation-invariant and applicable directly to document sets with or without preprocessing.This contrasts with methods that rerank documents according to an initial order.
- Deep Learning for Ranking: Deep IR models often use univariate scoring, whereas SetRank is a listwise multivariate ranking model built from self-attention networks.The framework assumes query-document feature representations are computed in advance but can also support joint feature and ranking learning.
- Machine Learning Defined on Sets: Set-based machine learning studies permutation-invariant inputs, and SetRank uses Transformer self-attention to capture interactions among documents.The approach builds on work characterizing permutation-invariant functions and modeling interactions within sets.
3 RANKING MODEL DEFINED ON SETS
The paper formulates ranking as a scoring function over document sets that produces a permutation, requiring both cross-document interaction modeling and permutation invariance. It motivates multivariate scoring because independent scoring misses document dependencies, while prior multivariate approaches can depend on input order.
- Problem formulation: Learning-to-rank seeks a scoring function over retrieved document sets whose sorted scores produce the target document permutation.The function is trained on labeled query-document sets by minimizing empirical loss.
- Requirements: A valid set-based ranking model must characterize inter-document relationships while remaining unchanged when input documents are permuted.Permutation invariance ensures the same set receives the same ranking regardless of input order.
- Limitations of univariate scoring: Univariate scoring assigns documents independently, following the probability ranking principle but failing to model cross-document interactions.The PRP assumption treats each document’s relevance as independent of other documents, which may not hold in realistic settings.
- Permutation properties: Permutation-equivariant document scores combined with sorting provide a route to permutation-invariant ranking outputs.Under a permutation-equivariant function, permuting inputs correspondingly permutes scores, so sorting yields the same document ranking.
4 OUR APPROACH: SETRANK
SetRank is a neural ranking model for document sets that uses self-attention to jointly encode documents while supporting arbitrary set sizes. Its pipeline represents documents, encodes their interactions, then scores and sorts them.
- Model overview: SetRank is a multivariate ranking model defined directly on document sets and designed to satisfy cross-document interaction and permutation-invariance requirements.The model uses multi-head self-attention and its modifications without relying on an initial document order.
- Pipeline: The representation layer separately generates query-document representations before the encoding layer jointly processes the retrieved documents.The encoding layer uses stacked MSAB or IMSAB sub-layers to produce internal document representations.
- Pipeline: The ranking layer applies a row-wise feed-forward function to encoded document vectors, generates relevance scores, and sorts the documents.This converts the jointly learned document representations into the final ranked list.
4.2 Document Representation
SetRank represents each retrieved document with a feature vector built from learning-to-rank features and optionally ordinal embeddings from initial rankings. These vectors are combined into the input feature matrix.
- Feature representation: Each retrieved document is represented as a feature vector extracted from query-document and document-level information.The experiments use benchmark-provided features, including signals such as TF-IDF, BM25, and PageRank.
- Ordinal embeddings: SetRank can optionally add ordinal embeddings that encode a document’s absolute position in an initial ranking.Initial rankings may come from models such as BM25 or LambdaMART.
- Input construction: The feature and ordinal-embedding vectors are added to form an N × E matrix for the retrieved documents.The resulting matrix is X = [d1 + p1, d2 + p2, · · · + dN + pN]^T.
- Multiple rankings: Ordinal embeddings from multiple initial rankings can be summed, allowing SetRank to incorporate several ranking sources simultaneously.This extends the input representation beyond a single preliminary ranking.
4.3 Document Encoding with (Induced) Multi-head Self Attention Block
SetRank’s encoding component stacks self-attention blocks to jointly model document interactions, while induced blocks use a fixed-size intermediate representation to handle varying set sizes. Standard MSAB can be sensitive to larger test-time sets, motivating IMSAB.
- Encoding component: The encoding component stacks MSAB or IMSAB blocks that jointly transform N document vectors into internal document representations.Each block applies attention, layer normalization, a row-wise feed-forward network, and another normalization step.
- MSAB: Multi-head attention projects inputs into multiple representation spaces before applying attention and returns an output with the query shape.The MAB combines multi-head attention with row-wise feed-forward processing and layer normalization.
- MSAB: MSAB uses self-attention with Q = K = V, allowing each document representation to incorporate information from all documents in the input set.The attention mechanism computes similarity-based weights and forms weighted sums of value vectors.
- Set-size limitation: MSAB trained with N0 labeled documents may not work well at test time when queries contain N1 > N0 documents.This set-size sensitivity reflects the limited number of labeled documents typically available during training.
- IMSAB: IMSAB first maps N documents to M induced cluster centers, then uses those centers to produce N contextualized document representations.Its two-stage construction enables processing document sets of any size, typically with M < N.
4.4 Document Ranking
SetRank constructs scoring functions by stacking multi-head self-attention blocks and applying row-wise feed-forward networks. Each document representation becomes a scalar score, and sorting these scores produces the final ranking.
- SetRank has MSAB and IMSAB variants formed by stacking the corresponding attention blocks before row-wise feed-forward networks.
- A row-wise feed-forward network projects each document representation into one real-valued ranking score.
- The final document ranking is obtained by sorting documents according to their predicted scores.
4.5 Model Training
SetRank trains with an attention rank loss that compares score-derived attention distributions against relevance-derived distributions. It also addresses variable input sizes through relative ordinal embedding sampling when ordinal embeddings are used.
- SetRank uses an attention rank loss to measure the distance between attention distributions generated from ranking scores and relevance judgments.
- For each document, relevance labels define an optimal attention allocation used as the training target.
- Predicted ranking scores generate a second attention distribution over the documents.
- The loss is a list-wise cross-entropy objective constructed from the relevance-based and score-based attention distributions.
- Relative ordinal embedding sampling addresses larger input sets when ordinal embeddings for those sizes are unavailable.
4.6 Theoretical Analysis
SetRank’s attention components and their compositions are permutation equivariant, while sorting the resulting scores yields permutation-invariant rankings. This supports ranking document sets independently of their input order.
- SetRank targets permutation-invariant ranking on document sets, unlike independent scoring or sequence-based multivariate ranking functions.
- The SetRankMSAB and SetRankIMSAB multivariate scoring functions are permutation equivariant.
- The multi-head attention block is permutation equivariant.
- The induced multi-head self-attention block is also permutation equivariant.
- Because compositions of permutation-equivariant functions remain permutation equivariant, both SetRank scoring functions inherit this property.
- Sorting the equivariant scores makes the two SetRank rankings permutation invariant.
5 EXPERIMENTS
Experiments on three large-scale learning-to-rank benchmarks compare SetRank with traditional and neural baselines, testing performance, robustness to noisy input order, set-size changes, and initial rankings.
- Experimental setup: Experiments used Istella, MSLR30K, and Yahoo learning-to-rank datasets with established traditional and neural ranking baselines.Models were evaluated using NDCG at ranks 1, 3, 5, and 10.
- Overall results: Multivariate scoring models generally outperformed univariate models, indicating benefits from modeling contextual information and document interactions.The comparison included models with and without LambdaMART-generated initial rankings.
- Overall results: SetRank outperformed DLCM and GSF with and without initial rankings, improving NDCG@1 by about 0.05 and 0.02 on Istella in the respective settings.The reported improvements compare SetRankIMSAB with DLCMw/o init and SetRankinitIMSAB with DLCM.
- Permutation invariance: As initial rankings became noisier, both SetRankinitIMSAB and DLCM degraded, but DLCM showed substantially larger performance declines.SetRankIMSAB without ordinal embeddings remained unchanged as reversed document pairs increased.
- Set-size adaptation: When training on 40-document queries and testing on 240-document queries, NDCG@10 drops were 0.0334 for DLCMw/o init, 0.0164 for SetRankMSAB, and 0.0049 for SetRankIMSAB.The induced self-attention variant had the smallest reported sensitivity to test-set size.
- Initial rankings: Using four initial rankings produced only a small improvement over using one, while one initial ranking outperformed using none.The authors attribute the limited benefit of multiple rankings partly to LambdaMART already providing substantially better rankings than the alternatives.
6 CONCLUSION
SetRank is a multivariate learning-to-rank model that maps document sets to permutations while modeling cross-document interactions and preserving permutation invariance. Experiments on three large-scale datasets show improvements over traditional and state-of-the-art deep ranking models.
- SetRank maps a document set to a permutation while satisfying cross-document interaction and permutation-invariance requirements.Its scoring function is implemented with Transformer self-attention.
- Self-attention enables SetRank to capture local context information and incorporate multiple initial rankings, while supporting robustness to input noise.
- SetRank outperformed traditional learning-to-rank models and state-of-the-art deep ranking models on three large-scale datasets.Analyses attributed improved performance and robustness to the two design requirements.
A PROOF TO PROPOSITION 1
The proposition establishes that the self-attention components used in SetRank preserve permutation-equivariance. The proof reduces the argument to the self-attention function because the other listed operations are element-wise, then invokes permutation invariance of summation.
- The multi-head parameters, row-wise feedforward function, and layer normalization are element-wise and therefore preserve permutation-equivariance.
- The proof represents the input as a set X = {x1, x2, . . . , xN} and considers the self-attention function as the remaining component requiring analysis.
- For any permutation π ∈ ΠN, the proof uses the permutation invariance of summation in establishing the required property.