Investigating Mixture of Experts in Dense Retrieval

TL;DR

This paper explores the integration of a single Mixture-of-Experts (MoE) block after the final Transformer layer in Dense Retrieval Models to enhance their effectiveness and robustness across various benchmarks.

Contribution

It introduces and empirically evaluates a novel SB-MoE architecture applied after the last Transformer layer in DRMs, comparing its performance to standard fine-tuning.

Findings

SB-MoE improves retrieval effectiveness for small models like TinyBERT.

Larger models like BERT and Contriever need more training data for SB-MoE to outperform fine-tuning.

SB-MoE's performance varies with the number of experts and model size.

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.