PaperlayerView exact revision on arXiv

Source-linked AI summary

Mixture of Experts Approaches in Dense Retrieval Tasks

Effrosyni Sokli, Pranav Kasela, Georgios Peikos, Gabriella Pasi

arXiv:2510.15683v1cs.IRcs.AI

TL;DR

Dense retrieval models can struggle to generalize across tasks and domains without additional fine-tuning, while layer-wise MoE approaches add parameters. This paper adds a single MoE block after the final Transformer layer. It reports consistent gains for lightweight DRMs and less pronounced effects for larger models.

  • Problem

    Dense retrieval models can struggle to adapt robustly across tasks and domains without additional fine-tuning, despite requiring large labeled datasets for training.

    Source: §I · p. 1

  • Method

    The paper adds a single MoE block after the final Transformer layer and trains it with the underlying DRM to refine query and document representations.

    Source: §I · pp. 1–2

  • Results

    SB-MOE consistently improves nDCG@10 and Recall@100 for lightweight DRMs across benchmarks, while gains are less pronounced for larger models; on NQ, SB-MOETOP-1 exceeds FINE-TUNED by 5.58% in nDCG@10.

    Source: §V · p. 4, §VII · p. 6

  • Takeaways & Limitations

    The reported results identify lightweight DRMs as a setting where SB-MOE consistently improves retrieval, with effective expert routing important to performance gains.

    Source: §VII · pp. 6–7

    For larger DRMs, the impact is less pronounced, and improved retrieval may require more training samples.

    Source: §V · p. 4, §VII · p. 6

Abstract

from arXiv · show

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.

I. INTRODUCTION

Dense retrieval models can struggle to generalize across tasks and domains without additional fine-tuning, while earlier layer-wise MoE designs add substantial parameters. This work proposes a single MoE block after the final Transformer layer and evaluates it across retrieval tasks.

I. INTRODUCTION

  • Earlier IR approaches inserted MoE into every Transformer layer, improving performance but substantially increasing additional parameters.SB-MOE instead adds one block after the final Transformer layer.Source: §I · p. 1
  • SB-MOE uses paired Feed-Forward Networks as experts, with a learned gating function assigning importance weights and a pooling strategy aggregating their outputs.The gating function is trained in an unsupervised manner.Source: §I · p. 1
  • The study evaluates SB-MOE on passage retrieval, open-domain Q&A, and domain-specific academic search using in-domain and zero-shot setups.The setups use seven benchmarks for in-domain evaluation and thirteen BEIR datasets for zero-shot evaluation after fine-tuning on MSMARCO.Source: §I · p. 1
  • The contribution includes testing four underlying DRMs and examining how the numbers of employed and activated experts affect retrieval effectiveness.The block takes query and document embeddings from the underlying DRM’s final Transformer layer as input.Source: §I · pp. 1–2

II. RELATED WORK

Prior IR work uses MoE within Transformer layers or only on query embeddings, but these designs can add parameters or leave document representations unmodified. This paper adds one MoE block to both query and document representations.

II. RELATED WORK

  • The study uses Contriever, BERT-Base, BERT-Small, and TinyBERT to examine MoE effects across different base-model sizes and pretraining strategies.The listed parameter counts range from 14.5M for TinyBERT to 110M for Contriever and BERT-Base.Source: §II · p. 2
  • DRMs can have low generalizability without further domain- or task-specific fine-tuning, motivating investigation of MoE approaches for retrieval.The related-work discussion notes MoE applications in passage retrieval and question answering.Source: §II · p. 2
  • Earlier IR MoE approaches replace Transformer feed-forward blocks with multiple expert networks, often substantially increasing model parameters.Another approach applies MoE only to the query embedding, partially exploring its use in DRMs.Source: §II · p. 2
  • The paper adds a single MoE block rather than blocks in every Transformer layer, applying it to both query and document representations.It evaluates retrieval effectiveness on passage retrieval, open-domain Q&A, and domain-specific academic search.Source: §II · p. 2

III. METHODOLOGY

SB-MOE refines both query and document embeddings using experts placed after the base DRM’s final Transformer layer. A gating function selects or weights expert outputs, and pooling produces the final embeddings for similarity estimation.

III. METHODOLOGY

  • The DRM independently encodes queries and documents, and SB-MOE applies experts to both embeddings after the final Transformer layer.Pooling aggregates expert representations into the final embedding used for query-document similarity estimation.Source: §III · pp. 2–3
  • The gating function maps each input embedding to weights indicating how suitable the experts are, while the pooling module aggregates their outputs.Training uses noisy Top-1 gating to encourage exploration of every expert.Source: §III · p. 3
  • SB-MOETOP-1 selects only the output of the expert assigned the highest gating score for the input embedding.The selected expert’s output forms the refined representation.Source: §III · p. 3
  • SB-MOEALL uses softmax-normalized gating scores to form a weighted sum of the outputs from all employed experts.This pooling strategy combines all expert outputs into the final refined representation.Source: §III · p. 3

IV. EXPERIMENTAL SETUP

The experimental setup frames comparisons with fine-tuned underlying DRMs around retrieval effectiveness, expert routing, and cross-domain generalization. The questions cover whether gains persist across tasks without additional fine-tuning.

IV. EXPERIMENTAL SETUP

  • The setup asks whether SB-MOE improves retrieval effectiveness over fine-tuning the underlying DRM.The comparison uses baselines described in the experimental setup.Source: §IV · p. 3
  • The evaluation asks whether performance gains can be attributed to the learned gating function and experts’ ability to capture useful textual characteristics.Source: §IV · p. 3
  • The study tests whether SB-MOE generalizes across multiple domains and tasks without additional fine-tuning.Source: §IV · p. 3

A. Baselines and Metrics

The evaluation compares SB-MoE with standard fine-tuning and random expert weighting, then tests zero-shot generalization after MSMARCO fine-tuning.

A. Baselines and Metrics

  • SB-MoE is compared with the underlying DRM fine-tuned on the same data and hyperparameters.This FINE-TUNED baseline addresses RQ1.Source: A. Baselines and Metrics · p. 3
  • RANDOM-GATE assigns experts random weights to assess the learned gate and experts’ contribution to embedding refinement.This baseline addresses RQ2.Source: A. Baselines and Metrics · p. 3
  • For zero-shot evaluation, all four DRMs are fine-tuned on MSMARCO, with and without MoE, then tested on 13 BEIR datasets.This FTMSMARCO setup addresses RQ3; results use nDCG@10 and Recall@100.Source: A. Baselines and Metrics · p. 3

B. Training Hyperparameters

The experiments use six experts by default and vary expert count in additional trials, with training schedules and settings adapted across models and datasets.

B. Training Hyperparameters

  • Six experts are used for all models and datasets, with additional experiments varying the count from 3 to 12 in steps of 3.Expert counts were selected with reference to prior work on expert-count effects.Source: B. Training Hyperparameters · p. 3
  • Each expert uses a down-projection FFN that halves the input dimension, followed by an up-projection that restores it.The design also uses a skip connection for residual learning, and the gate has a hidden layer that halves vector dimension.Source: B. Training Hyperparameters · pp. 3–4
  • Training uses batch size 64 and learning rates 10^-6 for the DRM and 10^-4 for experts; TinyBERT trains for 30 epochs, while other models train for 20 or 10 on CS.Five percent of training data is used for validation; the lowest-validation-loss checkpoint is retained.Source: B. Training Hyperparameters · p. 4

C. Datasets

The evaluation covers passage retrieval, open-domain Q&A, and domain-specific academic search, using seven benchmarks and a separate 13-dataset zero-shot BEIR evaluation.

C. Datasets

  • MSMARCO and TREC Deep Learning 2019 and 2020 share a passage corpus but differ in relevance judgments and test-query counts.MSMARCO has approximately 8.8 million passages and 532k training queries; the TREC datasets have 43 and 54 test queries, respectively.Source: C. Datasets · p. 4
  • Natural Questions and HotpotQA provide open-domain Q&A evaluations over Wikipedia-based corpora, with HotpotQA using multihop questions.Their corpora contain 2.6 million and 5.2 million passages or documents, respectively.Source: C. Datasets · p. 4
  • Political Science and Computer Science benchmark domain-specific academic search, and SB-MoE is also evaluated zero-shot on 13 BEIR datasets after MSMARCO fine-tuning.The two academic-search collections contain 4.8 million documents each.Source: C. Datasets · p. 4

V. RESULTS

SB-MoE’s strongest reported gains are for lightweight DRMs, including in zero-shot BEIR evaluation; learned gating also outperforms random expert weighting.

RQ1

  • SB-MoE consistently improves nDCG@10 and Recall@100 for TinyBERT and BERT-Small across datasets, while BERT-Base is mostly on par with fine-tuning on MSMARCO and TREC DL 19 & 20.For BERT-Small on NQ, SB-MOETOP-1 exceeds FINE-TUNED by 5.58% in nDCG@10; added experts may be redundant in larger models.Source: §V · p. 4

RQ3

  • TinyBERT and BERT-Small improve across 10 and 9 of 13 BEIR datasets, respectively, while BERT-Base and Contriever match or exceed FTMSMARCO on 6 and 5 datasets.Table II reports the zero-shot comparison; gains for larger models are less pronounced.Source: Table II · p. 5, §V · p. 5

RQ2

  • RANDOM-GATE performs comparably to FINE-TUNED, whereas SB-MoE consistently outperforms RANDOM-GATE, underscoring the learned gate’s role in selecting relevant experts.The passage attributes the gains to learned gating and experts’ ability to capture useful textual characteristics.Source: §V · pp. 4–5

VI. HYPERPARAMETER ANALYSIS

The number of employed experts substantially affects retrieval performance, but no configuration is optimal across all metrics and datasets. Inference activates only a subset of experts, while visualizations show changes in query and document embeddings after adding experts.

VI. HYPERPARAMETER ANALYSIS

  • Figure 2 reports that, on NQ, SB-MOEALL2’s nDCG@10 retrieval performance varies considerably with the number of employed experts, compared with fine-tuned TinyBERT’s performance.Source: §VI · p. 5
  • On TREC DL 19, 3 experts maximize nDCG@10 while 12 maximize Recall@100; no expert count is consistently optimal across datasets and metrics.Increasing the number of experts does not guarantee improved performance.Source: §VI · p. 6
  • During inference, 2 to 7 experts activate per collection, and employing 3 experts activates all experts across every collection.For HotpotQA, 9 employed experts activate 7, while 12 employed experts activate 5.Source: §VI · p. 6, Table III · p. 6
  • With 3 experts, query and relevant-document positions shift, same-expert documents cluster, and query-document representations align better in the dense vector space.These visualizations cover seven benchmarks; the 3-expert configuration activates all experts across the evaluation collections.Source: §VI · p. 6

VII. CONCLUSIONS

SB-MoE consistently improves retrieval for lightweight DRMs, while larger models show only marginal gains. The authors also report that effective gating matters, and that the approach can generalize across domains and tasks with low computational overhead.

VII. CONCLUSIONS

  • SB-MoE consistently improves nDCG@10 and Recall@100 for lightweight DRMs across benchmarks, while larger models show only marginal gains.The reported lightweight models include TinyBERT and BERT-Small.Source: §VII · pp. 6–7
  • Retrieval gains rely on effective, optimized expert routing; adding experts alone is insufficient.The learned gating function plays a critical role in expert selection.Source: §VII · p. 7
  • SB-MoE can generalize across domains and tasks while maintaining low computational overhead, supporting its suitability for low-resource environments.Source: §VII · p. 7

Related papers

Loading 2510.15683v1…