Matryoshka Re-Ranker: A Flexible Re-Ranking Architecture With Configurable Depth and Width

TL;DR

The paper introduces Matryoshka Re-Ranker, a flexible, configurable architecture for LLM-based re-ranking that adapts to various scenarios while maintaining high performance through innovative optimization techniques.

Contribution

It presents a novel flexible re-ranking architecture with runtime customization and introduces techniques like cascaded self-distillation and low-rank adaptation to mitigate precision loss.

Findings

Outperforms existing re-ranking methods on MSMARCO and BEIR datasets.

Maintains high accuracy across various compression levels and scenarios.

Demonstrates effective trade-off between flexibility and precision.

Abstract

Large language models (LLMs) provide powerful foundations to perform fine-grained text re-ranking. However, they are often prohibitive in reality due to constraints on computation bandwidth. In this work, we propose a \textbf{flexible} architecture called \textbf{Matroyshka Re-Ranker}, which is designed to facilitate \textbf{runtime customization} of model layers and sequence lengths at each layer based on users' configurations. Consequently, the LLM-based re-rankers can be made applicable across various real-world situations. The increased flexibility may come at the cost of precision loss. To address this problem, we introduce a suite of techniques to optimize the performance. First, we propose \textbf{cascaded self-distillation}, where each sub-architecture learns to preserve a precise re-ranking performance from its super components, whose predictions can be exploited as smooth and informative teacher signals. Second, we design a \textbf{factorized compensation mechanism}, where two collaborative Low-Rank Adaptation modules, vertical and horizontal, are jointly employed to compensate for the precision loss resulted from arbitrary combinations of layer and sequence compression. We perform comprehensive experiments based on the passage and document retrieval datasets from MSMARCO, along with all public datasets from BEIR benchmark. In our experiments, Matryoshka Re-Ranker substantially outperforms the existing methods, while effectively preserving its superior performance across various forms of compression and different application scenarios.