RankMamba: Benchmarking Mamba's Document Ranking Performance in the Era of Transformers

📄 arXiv: 2403.18276v3 📥 PDF

作者: Zhichao Xu

分类: cs.IR, cs.CL

发布日期: 2024-03-27 (更新: 2026-01-05)

🔗 代码/项目: GITHUB

💡 一句话要点

评估Mamba模型在文档排序任务中的表现

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 文档排序 信息检索 Transformer 状态空间模型 模型评估 机器学习

📋 核心要点

  1. 现有Transformer模型在处理长文本时面临训练时间复杂度高和推理效率低的问题。
  2. 本文提出Mamba模型，通过状态空间模型替代传统的注意力机制，旨在提高文档排序任务的性能。
  3. 实验结果显示，Mamba模型在文档排序任务中与Transformer模型表现相当，但训练效率较低，提供了新的研究方向。

📝 摘要（中文）

Transformer结构在自然语言处理、计算机视觉和信息检索等多个应用领域取得了显著成功。其核心机制——注意力机制在训练时的时间复杂度为O(n^2)，推理时为O(n)。为提高注意力机制的可扩展性，提出了多种改进方法。Mamba模型基于状态空间模型，在多个序列建模任务中表现出与Transformer相当的性能。本文通过文档排序这一经典信息检索任务，评估Mamba模型的有效性，发现Mamba模型在与Transformer模型相同训练条件下表现竞争力，但训练吞吐量低于高效的Transformer实现，如Flash Attention。希望本研究能为Mamba模型在其他经典信息检索任务中的探索提供起点。

🔬 方法详解

问题定义：本文旨在解决文档排序任务中现有Transformer模型在处理长文本时的效率问题，尤其是训练时间复杂度高和推理效率低的挑战。

核心思路：Mamba模型通过引入状态空间模型替代传统的注意力机制，旨在在保证性能的同时提高模型的可扩展性和效率。

技术框架：整体架构包括输入查询和文档的重排序模型，模型通过预测相关性得分来评估查询与文档之间的互动。主要模块包括特征提取、相关性评分和模型训练。

关键创新：Mamba模型的核心创新在于其基于状态空间的设计，能够在不依赖传统注意力机制的情况下实现与Transformer相当的性能，显著降低了训练复杂度。

关键设计：在模型设计中，Mamba使用了特定的参数设置和损失函数，确保模型在处理长文本时能够有效捕捉查询与文档之间的交互信息。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，Mamba模型在文档排序任务中与同类Transformer模型的性能相当，尽管其训练吞吐量低于高效的Transformer实现（如Flash Attention），但为未来的研究提供了新的方向和可能性。

🎯 应用场景

该研究的潜在应用领域包括信息检索、搜索引擎优化和自然语言处理等。Mamba模型的设计理念可以为处理长文本的任务提供新的思路，未来可能在文档排序和相关性评估等领域产生重要影响。

📄 摘要（原文）

Transformer structure has achieved great success in multiple applied machine learning communities, such as natural language processing (NLP), computer vision (CV) and information retrieval (IR). Transformer architecture's core mechanism\, -- \,attention requires $O(n^2)$ time complexity in training and $O(n)$ time complexity in inference. Many works have been proposed to improve the attention mechanism's scalability, such as Flash Attention and Multi-query Attention. A different line of work aims to design new mechanisms to replace attention. Recently, a notable model structure Mamba, which is based on state space models, has achieved transformer-equivalent performance in multiple sequence modeling tasks. In this work, we examine Mamba's efficacy through the lens of a classical IR task\, -- \,document ranking. A reranker model takes a query and a document as input, and predicts a scalar relevance score. This task demands the language model's ability to comprehend lengthy contextual inputs and to capture the interaction between query and document tokens. We find that \textbf{(1) Mamba models achieve competitive performance compared to transformer-based models with the same training recipe; (2) but also have a lower training throughput in comparison to efficient transformer implementations such as flash attention.} We hope this study can serve as a starting point to explore \mamba models in other classical IR tasks. Our \href{https://github.com/zhichaoxu-shufe/RankMamba}{code implementation} is made public to facilitate reproducibility. Refer to~\cite{xu-etal-2025-state} for more comprehensive experiments and results, including passage ranking.