Mamba Retriever: Utilizing Mamba for Effective and Efficient Dense Retrieval (2408.08066v2)

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 LLMs (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.

An In-depth Analysis of "Mamba Retriever: Utilizing Mamba for Effective and Efficient Dense Retrieval"

The paper "Mamba Retriever: Utilizing Mamba for Effective and Efficient Dense Retrieval" presents an innovative approach to addressing the perennial challenge in Information Retrieval (IR) of balancing efficiency and effectiveness in dense retrieval (DR) models. The proposed solution, termed Mamba Retriever, leverages the Mamba architecture as an encoder for DR models. Recent literature and experimental results from this paper suggest that the Mamba architecture is not only competitive with Transformer-based pre-trained LLMs (PLMs) in terms of effectiveness but also superior in terms of computational efficiency, particularly with long-text retrieval tasks.

Core Contributions

The key contributions of this work can be delineated as follows:

1. Implementation of Mamba Retriever: The authors propose Mamba Retriever, a novel bi-encoder retrieval model based on the Mamba architecture. Unlike traditional Transformer-based models, the Mamba architecture employs selective state space models (SSMs) to achieve linear time scaling relative to sequence length, thereby circumventing the quadratic complexity inherent in Transformer-based models.
2. Effectiveness on Short-text Retrieval: The Mamba Retriever was fine-tuned on the MS MARCO passage ranking dataset and evaluated against both short-text (MS MARCO) and various BEIR benchmark datasets. Experimental results reveal that Mamba Retriever exhibits comparable or superior performance to well-established Transformer-based models such as BERT, RoBERTa, and OPT at different model sizes. Notably, the effectiveness scales positively with model size, with significant evidence provided by metrics such as MRR@10 and Recall@1k.
3. Effectiveness on Long-text Retrieval: The Mamba architecture's capacity to handle long-text retrieval was examined using the LoCoV0 dataset. The model demonstrated robust performance on long-text retrieval, maintaining or exceeding the effectiveness of other long-text retrieval models, including the M2-BERT model. Remarkably, the Mamba Retriever managed to extend its competence beyond its pre-trained length post fine-tuning, affirming its adaptability to longer sequences.
4. Inference Efficiency: One of the pivotal advantages of the Mamba-based model highlighted by this paper is its superior inference speed for long-text retrieval. When tested across a range of text lengths, Mamba Retriever consistently outperformed Transformer-based models by a substantial margin. The linearity in time complexity with respect to sequence length offers a compelling benefit in terms of scalability and practical application in environments where efficiency is paramount.

Implications and Future Directions

The implications of this research are multifaceted:

• Practical Implementations: The Mamba Retriever's efficiency and effectiveness make it highly suitable for practical deployment in various IR applications, especially where long-text processing is a requirement. The demonstrated speed advantages can translate to significant reductions in computational overhead and improved responsiveness in real-time systems.
• Theoretical Contributions: The employment of SSMs and the introduction of selective state mechanisms in the Mamba architecture offer a notable expansion in the toolkit for modeling long-range dependencies in sequential data. This could spur further research into alternative architectures that prioritize both efficiency and effectiveness.
• Benchmarking and Comparisons: The rigorous benchmarking on both MS MARCO and LoCoV0 datasets provides a robust validation of Mamba Retriever's capabilities. Moreover, comparisons with contemporary models like Jina Embeddings v2 and fine-tuned M2-BERT model offer a comprehensive landscape of current methodologies and their relative benchmarks.

Speculation on Future Developments in AI

This research opens avenues for several interesting future developments in AI, particularly within the domain of IR:

• Hybrid Models: Future models may explore hybrid architectures that integrate selective state mechanisms with other efficient modeling techniques, potentially leading to further improvements in handling exceedingly long-text inputs.
• Adaptive and Dynamic Architectures: Leveraging the principles of selective state mechanisms, more adaptive and dynamic architectures could be developed to dynamically adjust the computation based on the input characteristics, optimizing both resource utilization and performance.
• Cross-domain Applications: The efficiency of the Mamba architecture indicates potential applications beyond IR, including but not limited to natural language understanding, machine translation, and large-scale text summarization tasks.

In conclusion, the Mamba Retriever represents a significant stride towards more efficient and effective dense retrieval, evidenced by comprehensive experimental results. This positions it as a highly practical model for modern IR tasks, particularly where long documents are involved. The research not only underscores the advantages of non-Transformer PLMs but also sets a foundation for future work exploring the vast potential of selective state space mechanisms in various AI applications.