Mamba Retriever: Utilizing Mamba for Effective and Efficient Dense Retrieval

TL;DR

The paper addresses the efficiency bottleneck of Transformer-based dense retrieval by introducing Mamba Retriever, a bi-encoder that leverages the selective state-space model of Mamba to achieve linear-time processing for long texts. It trains the model with a contrastive InfoNCE objective to produce embeddings for queries and passages, and evaluates on classic short-text (MS MARCO) and long-text (LoCoV0) datasets. Results show that Mamba Retriever matches or exceeds Transformer-based retrievers in effectiveness, with performance improving as model size grows, and demonstrates superior long-text retrieval capability and inference speed due to linear sequence scaling. The work suggests that Mamba Retriever is a practical, scalable solution for IR tasks, especially where long documents are prevalent, and provides an accessible implementation at the project repository.

Abstract

In the information retrieval (IR) area, dense retrieval (DR) models use deep learning techniques to encode queries and passages into embedding space to compute their semantic relations. It is important for DR models to balance both efficiency and effectiveness. Pre-trained language models (PLMs), especially Transformer-based PLMs, have been proven to be effective encoders of DR models. However, the self-attention component in Transformer-based PLM results in a computational complexity that grows quadratically with sequence length, and thus exhibits a slow inference speed for long-text retrieval. Some recently proposed non-Transformer PLMs, especially the Mamba architecture PLMs, have demonstrated not only comparable effectiveness to Transformer-based PLMs on generative language tasks but also better efficiency due to linear time scaling in sequence length. This paper implements the Mamba Retriever to explore whether Mamba can serve as an effective and efficient encoder of DR model for IR tasks. We fine-tune the Mamba Retriever on the classic short-text MS MARCO passage ranking dataset and the long-text LoCoV0 dataset. Experimental results show that (1) on the MS MARCO passage ranking dataset and BEIR, the Mamba Retriever achieves comparable or better effectiveness compared to Transformer-based retrieval models, and the effectiveness grows with the size of the Mamba model; (2) on the long-text LoCoV0 dataset, the Mamba Retriever can extend to longer text length than its pre-trained length after fine-tuning on retrieval task, and it has comparable or better effectiveness compared to other long-text retrieval models; (3) the Mamba Retriever has superior inference speed for long-text retrieval. In conclusion, Mamba Retriever is both effective and efficient, making it a practical model, especially for long-text retrieval.

Figures (1)

• Figure 1: The efficiency of Mamba Retriever compared to long-text retrieval models at different maximum text lengths.