Investigating Mixture of Experts in Dense Retrieval

TL;DR

Dense Retrieval Models often struggle with generalizability. This paper introduces SB-MoE, a single Mixture-of-Experts block appended to the final embedding of a bi-encoder, guided by a gating function and pooled via Top-1 or softmax-weighted strategies, trained end-to-end with a contrastive loss. Across four benchmarks (BEIR and Multi-Domain) and three DRMs (TinyBERT, BERT, Contriever), SB-MoE yields strong gains for lightweight models but shows limited improvements for larger models, with results depending on the number of experts. The key takeaway is that SB-MoE can enhance retrieval when model size is constrained, but expert-count tuning is essential for domain-specific performance and practical deployment.

Abstract

While Dense Retrieval Models (DRMs) have advanced Information Retrieval (IR), one limitation of these neural models is their narrow generalizability and robustness. To cope with this issue, one can leverage the Mixture-of-Experts (MoE) architecture. While previous IR studies have incorporated MoE architectures within the Transformer layers of DRMs, our work investigates an architecture that integrates a single MoE block (SB-MoE) after the output of the final Transformer layer. Our empirical evaluation investigates how SB-MoE compares, in terms of retrieval effectiveness, to standard fine-tuning. In detail, we fine-tune three DRMs (TinyBERT, BERT, and Contriever) across four benchmark collections with and without adding the MoE block. Moreover, since MoE showcases performance variations with respect to its parameters (i.e., the number of experts), we conduct additional experiments to investigate this aspect further. The findings show the effectiveness of SB-MoE especially for DRMs with a low number of parameters (i.e., TinyBERT), as it consistently outperforms the fine-tuned underlying model on all four benchmarks. For DRMs with a higher number of parameters (i.e., BERT and Contriever), SB-MoE requires larger numbers of training samples to yield better retrieval performance.

Figures (2)

• Figure 1: Overview of the SB-MoE architecture, highlighting its three main parts.
• Figure 2: SB-MoE$_{\text{TOP-1}}$ on TinyBERT with 3, 6, 9, and 12 experts.