Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Mixture of Experts Approaches in Dense Retrieval Tasks

Effrosyni Sokli, Pranav Kasela, Georgios Peikos, Gabriella Pasi

TL;DR

This paper tackles the generalization challenge of dense retrieval models by introducing SB-MoE, a single Mixture-of-Experts block appended after the final Transformer layer to enrich query and document embeddings without heavy parameter costs. The approach couples a gating function with multiple FFN experts and supports two pooling modes (Top-1 and ALL) to produce refined embeddings, trained with unsupervised gating. Across seven benchmarks and three IR tasks, including a zero-shot MSMARCO→BEIR evaluation, SB-MoE yields consistent gains for lightweight DRMs and demonstrates cross-domain generalization, while providing guidance on expert count and activation. The findings highlight the gating mechanism’s critical role, show favorable efficiency trade-offs, and offer a practical path toward more robust dense retrieval in resource-constrained settings.

Abstract

Dense Retrieval Models (DRMs) are a prominent development in Information Retrieval (IR). A key challenge with these neural Transformer-based models is that they often struggle to generalize beyond the specific tasks and domains they were trained on. To address this challenge, prior research in IR incorporated the Mixture-of-Experts (MoE) framework within each Transformer layer of a DRM, which, though effective, substantially increased the number of additional parameters. In this paper, we propose a more efficient design, which introduces a single MoE block (SB-MoE) after the final Transformer layer. To assess the retrieval effectiveness of SB-MoE, we perform an empirical evaluation across three IR tasks. Our experiments involve two evaluation setups, aiming to assess both in-domain effectiveness and the model's zero-shot generalizability. In the first setup, we fine-tune SB-MoE with four different underlying DRMs on seven IR benchmarks and evaluate them on their respective test sets. In the second setup, we fine-tune SB-MoE on MSMARCO and perform zero-shot evaluation on thirteen BEIR datasets. Additionally, we perform further experiments to analyze the model's dependency on its hyperparameters (i.e., the number of employed and activated experts) and investigate how this variation affects SB-MoE's performance. The obtained results show that SB-MoE is particularly effective for DRMs with lightweight base models, such as TinyBERT and BERT-Small, consistently exceeding standard model fine-tuning across benchmarks. For DRMs with more parameters, such as BERT-Base and Contriever, our model requires a larger number of training samples to achieve improved retrieval performance. Our code is available online at: https://github.com/FaySokli/SB-MoE.

Mixture of Experts Approaches in Dense Retrieval Tasks

TL;DR

This paper tackles the generalization challenge of dense retrieval models by introducing SB-MoE, a single Mixture-of-Experts block appended after the final Transformer layer to enrich query and document embeddings without heavy parameter costs. The approach couples a gating function with multiple FFN experts and supports two pooling modes (Top-1 and ALL) to produce refined embeddings, trained with unsupervised gating. Across seven benchmarks and three IR tasks, including a zero-shot MSMARCO→BEIR evaluation, SB-MoE yields consistent gains for lightweight DRMs and demonstrates cross-domain generalization, while providing guidance on expert count and activation. The findings highlight the gating mechanism’s critical role, show favorable efficiency trade-offs, and offer a practical path toward more robust dense retrieval in resource-constrained settings.

Abstract

Dense Retrieval Models (DRMs) are a prominent development in Information Retrieval (IR). A key challenge with these neural Transformer-based models is that they often struggle to generalize beyond the specific tasks and domains they were trained on. To address this challenge, prior research in IR incorporated the Mixture-of-Experts (MoE) framework within each Transformer layer of a DRM, which, though effective, substantially increased the number of additional parameters. In this paper, we propose a more efficient design, which introduces a single MoE block (SB-MoE) after the final Transformer layer. To assess the retrieval effectiveness of SB-MoE, we perform an empirical evaluation across three IR tasks. Our experiments involve two evaluation setups, aiming to assess both in-domain effectiveness and the model's zero-shot generalizability. In the first setup, we fine-tune SB-MoE with four different underlying DRMs on seven IR benchmarks and evaluate them on their respective test sets. In the second setup, we fine-tune SB-MoE on MSMARCO and perform zero-shot evaluation on thirteen BEIR datasets. Additionally, we perform further experiments to analyze the model's dependency on its hyperparameters (i.e., the number of employed and activated experts) and investigate how this variation affects SB-MoE's performance. The obtained results show that SB-MoE is particularly effective for DRMs with lightweight base models, such as TinyBERT and BERT-Small, consistently exceeding standard model fine-tuning across benchmarks. For DRMs with more parameters, such as BERT-Base and Contriever, our model requires a larger number of training samples to achieve improved retrieval performance. Our code is available online at: https://github.com/FaySokli/SB-MoE.
Paper Structure (10 sections, 3 equations, 4 figures, 3 tables)

This paper contains 10 sections, 3 equations, 4 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Methodology
  4. Experimental Setup
  5. Baselines and Metrics
  6. Training Hyperparameters
  7. Datasets
  8. Results
  9. Hyperparameter Analysis
  10. Conclusions

Figures (4)

  • Figure 1: SB-MoE with its three main parts: (i) the experts, (ii) the gating function, and (iii) the pooling module.
  • Figure 2: nDCG@10 and Recall@100 scores of SB-MoE$_{\text{ALL}}$ on TinyBERT with 3, 6, 9, and 12 experts.
  • Figure 3: 3D t-SNE visualizations of a query and its top 1000 retrieved documents, on the left embedded by the original DRM and on the right by SB-MoE. Color figure online.
  • Figure 4: Continuation of Figure \ref{['fig:tsne1']}. 3D t-SNE visualizations of a query and its top 1000 retrieved documents, on the left embedded by the original DRM and on the right by SB-MoE.