RankMamba：在Transformer时代基准测试Mamba的文档排序性能

Transformer结构在自然语言处理（NLP）、计算机视觉（CV）和信息检索（IR）等多个应用机器学习社区中取得了巨大成功。Transformer架构的核心机制——注意力，在训练时需要 $O(n^2)$ 的时间复杂度，在推理时需要 $O(n)$ 的时间复杂度。已有许多工作被提出以改善注意力机制的可扩展性，例如Flash Attention和Multi-query Attention。另一条研究路线旨在设计新机制来替代注意力。最近，一种基于状态空间模型的著名模型结构Mamba在多个序列建模任务中取得了与Transformer相当的性能。在这项工作中，我们通过经典IR任务——文档排序的视角来检验Mamba的有效性。重排序模型将查询和文档作为输入，并预测一个标量相关性分数。该任务要求语言模型具备理解长上下文输入以及捕获查询与文档token之间交互的能力。我们发现 \textbf{(1) 与使用相同训练配方的基于Transformer的模型相比，Mamba模型取得了具有竞争力的性能；(2) 但与Flash Attention等高效Transformer实现相比，其训练吞吐量较低。} 我们希望这项研究能作为在其他经典IR任务中探索Mamba模型的起点。我们的 \href{https://github.com/zhichaoxu-shufe/RankMamba}{代码实现} 已公开以促进复现。有关更全面的实验和结果（包括段落排序），请参阅~\cite{xu-etal-2025-state}。