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

📄 arXiv: 2501.16302v1 📥 PDF

作者: Zheng Liu, Chaofan Li, Shitao Xiao, Chaozhuo Li, Defu Lian, Yingxia Shao

分类: cs.CL

发布日期: 2025-01-27

备注: The Web Conference 2025

💡 一句话要点

提出 Matryoshka Re-Ranker，实现可配置深度和宽度的灵活重排序架构，提升 LLM 在实际场景的应用性。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 文本重排序 大型语言模型 模型压缩 自蒸馏 低秩适应 灵活架构 运行时配置

📋 核心要点

  1. 现有基于 LLM 的重排序模型计算开销大，难以在实际资源受限的场景中应用。
  2. Matryoshka Re-Ranker 允许运行时配置模型深度和序列长度，实现灵活的计算资源分配。
  3. 通过级联自蒸馏和分解补偿机制，有效缓解了模型压缩带来的精度损失，提升了性能。

📝 摘要（中文）

大型语言模型（LLMs）为执行细粒度的文本重排序提供了强大的基础。然而，由于计算带宽的限制，它们在现实中常常是难以实现的。在这项工作中，我们提出了一种 extbf{灵活的}架构，称为 extbf{Matryoshka Re-Ranker}，它旨在促进基于用户配置的，模型层和每层序列长度的 extbf{运行时定制}。因此，基于LLM的重排序器可以适用于各种现实场景。增加的灵活性可能会以精度损失为代价。为了解决这个问题，我们引入了一套技术来优化性能。首先，我们提出了 extbf{级联自蒸馏}，其中每个子架构学习从其超级组件中保持精确的重排序性能，其预测可以被用作平滑且信息丰富的教师信号。其次，我们设计了一个 extbf{分解补偿机制}，其中垂直和水平两个协作的低秩适应模块被共同用于补偿由层和序列压缩的任意组合导致的精度损失。我们基于MSMARCO的段落和文档检索数据集，以及来自BEIR基准的所有公共数据集进行了全面的实验。在我们的实验中，Matryoshka Re-Ranker大大优于现有方法，同时有效地保持了其在各种形式的压缩和不同应用场景中的卓越性能。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLMs）在文本重排序任务中计算成本过高，难以在实际应用中部署的问题。现有方法通常无法根据不同的计算资源限制进行灵活调整，导致性能受限或无法应用。

核心思路：论文的核心思路是设计一种灵活的重排序架构，允许在运行时根据用户的配置动态调整模型的深度（层数）和宽度（序列长度）。通过这种方式，可以在计算资源和模型性能之间进行权衡，使LLM重排序器能够适应各种实际场景。

技术框架：Matryoshka Re-Ranker 的整体架构包含一个基础的 LLM 重排序模型，以及两个关键的优化模块：级联自蒸馏和分解补偿机制。首先，通过级联自蒸馏，让模型的各个子架构（不同深度和宽度的配置）学习从其父架构中获取知识，保持重排序性能。其次，利用分解补偿机制，通过垂直和水平两个低秩适应（LoRA）模块，补偿由于层和序列压缩导致的精度损失。

关键创新：该论文的关键创新在于提出了一个灵活的、可配置的重排序架构，允许在运行时动态调整模型的深度和宽度。此外，级联自蒸馏和分解补偿机制的结合，有效地解决了模型压缩带来的精度损失问题，使得该架构在各种计算资源限制下都能保持良好的性能。与现有方法相比，Matryoshka Re-Ranker 提供了更高的灵活性和更好的性能。

关键设计：级联自蒸馏使用父架构的预测作为教师信号，指导子架构的学习。分解补偿机制采用两个 LoRA 模块，分别在垂直方向（层之间）和水平方向（序列之间）进行参数调整，以补偿压缩带来的信息损失。具体的参数设置和损失函数选择可能需要根据具体的 LLM 和数据集进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Matryoshka Re-Ranker 在 MSMARCO 和 BEIR 基准数据集上显著优于现有方法。该方法在各种压缩配置下都能保持良好的性能，证明了其灵活性和鲁棒性。具体的性能提升幅度取决于数据集和压缩配置，但总体而言，Matryoshka Re-Ranker 能够有效地在计算资源和模型性能之间进行权衡。

🎯 应用场景

Matryoshka Re-Ranker 可应用于各种需要文本重排序的场景，例如搜索引擎、推荐系统、问答系统等。该方法能够根据实际的计算资源限制，动态调整模型的深度和宽度，从而在保证性能的同时，降低计算成本，使其更易于部署在资源受限的设备或环境中。未来，该方法可以进一步扩展到其他自然语言处理任务中。

📄 摘要（原文）

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.