MoTE: Mixture of Task Experts for Embedding Models

Authors: Miguel Romero Calvo, Shuoyang Ding, Corey D Barrett, Georgiana Dinu, George Karypis
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: Embedding Models, Representation Learning, Mixture of Experts, Retrieval Augmented Generation, RAG, Search, Clustering, Classification

Abstract:
Dense embeddings are essential for Retrieval-Augmented Generation (RAG), search, classification, and clustering systems. Recent methods improve dense embeddings by enriching them with instructions that incorporate downstream task information, enabling a single model to generate task-specific embedding spaces. However, we empirically show that requiring all tasks to share the same model parameters imposes significant representational limitations. To address these challenges, we introduce Mixture of Task Experts (MoTE), a novel transformer block designed for embedding architectures. MoTE employs dedicated parameter sets tailored to the unique requirements of each task and is paired with a task-aware training framework to improve representation quality. Experiments on 56 datasets spanning $7$ tasks demonstrate that MoTE outperforms instruction-conditioned models, achieving, on average, $1.62$ higher NDCG@10 on retrieval datasets, $1.54$ higher MAP on re-ranking datasets, and a $0.65$ improvement in overall performance. Notably, these gains are achieved without altering inference-time information, training data, inference speed, or number of active parameters.

Decision: reject
Reviewer scores: 5, 6, 3, 5

Reviews and discussion:
<thread>
creation date:
1730253289375

thread_type:
official_review

<thread_object>
creator: reviewer_mksg

<content>
summary:
This paper introduces a novel architecture, MoTE, which allocates a dedicated set of parameters for each task to enhance dense text embeddings. MoTE offers a robust approach to improving dense text embeddings by employing a mixture of experts. This architecture, combined with a dynamic training strategy and WAVE, achieves better task specialization without intensively increasing latency or memory overhead. Experiments were conducted on the MTEB benchmark to compare MoTE with previous instruction-based embedding model training approaches and analyze the contributions of its techniques.

strengths:
1. This paper explores the novel Mixture of Task Experts architecture for embedding models, which has the potential to effectively address the challenges posed by various retrieval tasks with differing intents.
2. It conducts numerous analytical experiments and investigates the impact of training configurations, yielding insights that may aid further research on MoTE.

weaknesses:
1. The experiments on MTEB do not compare MoTE against mainstream retrieval models, such as E5 and BGE, nor does it provide comparisons by using training data of mainstream retrieval models, which weakens the reliability of the experimental results.
2. Many analytical experiments only present results without in-depth analysis, such as in sections 4.5 and 4.6.
3. The writing is relatively poor. For example, the two paragraphs in section 3.3 appear to be repetitive.

questions:
1. Is the WAVE method introduced in section 4.7 too simplistic, merely averaging the expert layers post-training but can maintain the performance?
2. As many current MoE architectures are based on LLMs, could this approach be more suitable for retrieval models grounded on LLMs?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730353256220

thread_type:
official_review

<thread_object>
creator: reviewer_yimc

<content>
summary:
This paper tries to introduce the MOE architecture into the general embedding models and provide different hyperparameters for each task expert to improve performance. Meanwhile, the Weight Average of Vector Experts is obtained by merging the trained experts, which reduces the overall number of parameters and keeps the performance almost equal to that before merging.

strengths:
1. The improvements presented in this paper are straightforward, with expected gains from the MOE architecture and Weight Averaging.
2. The ablation experiments are sufficient, covering performance variations across different settings.
3. The final model maintains the backbone parameter count, ensuring a fair comparison.

weaknesses:
1. The work lacks innovation. It primarily enhances model performance by applying two established strategies (MOE and Weight Averaging) to generic representations without offering new insights.
2. The paper centres on technical implementation but (1) does not commit to open-sourcing the code and (2) provides a very brief methodology, omitting details such as the implementation of the Task-Aware Training Strategy.
3. The paper appears incomplete, falling short of the 10-page limit and lacking a thorough methodology section. Grammar and citation issues also affect readability.
4. Weight Averaging is a model merging method, yet no background on model merging is provided, making the work hard to follow for readers unfamiliar with this area.

questions:
1. The performance difference between "TEM-Inst." and MoTE is minimal (less than 1% in Tables 1 and 2). However, as "TEM-Inst." results don’t appear to be from a public source and the hyperparameter search process isn’t clearly detailed, could you clarify the hyperparameter tuning efforts made to ensure a fair comparison?

2. The statement that “existing methods, ..., limits the degree of task-specialization of the final vector representation” is unclear. In Tables 1 and 2, "TEM-Inst." clearly outperforms STEM, suggesting that multi-task joint learning generally enhances performance across tasks. Thus, setting aside the results in MTEB:

(1) What evidence shows that the trained representation space is task-specialized?

(2) To what extent does the representation space differ between tasks?

(3) Why can’t instructions alone resolve this issue?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730926695238

thread_type:
official_review

<thread_object>
creator: reviewer_161q

<content>
summary:
This paper presents a novel approach for task-specific fine-tuning of embedding models in a mixture-of-experts framework, termed Mixture of Task Experts (MoTE). The MoTE model generates task-conditioned embeddings through dedicated parameter sets for each task, allowing for specialized representations tailored to specific downstream requirements. To address the increased memory demands associated with this approach, the authors introduce Weight Average of Vector Experts (WAVE), which mitigates the computational overhead of MoTE by efficiently managing memory without sacrificing performance. This dual approach overcomes the limitations of single-expert models applied across diverse tasks.

To evaluate its effectiveness, the proposed method was tested on seven in-scope and out-of-scope tasks across 56 datasets, demonstrating notable performance gains due to the MoTE architecture.

strengths:
- This work presents innovative methods for generating task-aware embeddings, building on and extending previous research.
- It introduces a novel approach that leverages task-specific hyperparameters to create specialized embeddings while effectively managing the memory consumption challenges of task-conditioned models.
- The method demonstrates practical applicability by addressing the limitations of using a single expert across multiple downstream tasks, showing promise for real-world implementations.
- The paper is clearly written, with well-chosen examples and motivating illustrations that enhance comprehension.
- It includes a robust experimental setup, testing across seven tasks (both in-scope and out-of-scope) with a total of 56 datasets.
- A thorough ablation study and in-depth analysis strengthen the insights derived from the experimental results.

weaknesses:
- This work represents an incremental advancement in mixture-of-expert models, focusing on generating task-aware embeddings for downstream tasks and addressing limitations of prior approaches.
- The approach shows only marginal gains for out-of-scope tasks, such as in pair similarity. For task-aware applications like Classification, Clustering, and Retrieval, in general, they share semantic or latent spaces across them. However, the applicability of the proposed method does not generalize effectively to out-of-scope tasks and domain?
For example, in Table 2, the proposed approach yields only slight improvements on semantic textual similarity (STS) tasks, highlighting limited performance gains in out-of-scope settings.

questions:
NA
</content>
</thread_object>

</thread>
<thread>
creation date:
1731220300986

thread_type:
official_review

<thread_object>
creator: reviewer_ezqd

<content>
summary:
The paper introduces a novel approach for embedding models called Mixture of Task Experts (MoTE) that incorporates task-specific MLP blocks within the transformer architecture. During training, different tasks are directed to task-specific blocks, allow the model to generate different embeddings for the same query based on the task context.
This adds new parameters (the task-specific blocks) corresponding to the number of tasks during training. To reduce the parameter count, they apply WAVE (Weight Average of Vector Experts) where they average the weights of the task specific blocks.

strengths:
The embeddings of two different queries may need to be close or far apart depending on the task. Their approach to generating task-specific embeddings enables the optimization of distinct parameters for each task, promoting better task specialization.

weaknesses:
1. Adding a new ML block after every other transformer layer will add significant memory requirement. It would be beneficial to explore optimization strategies to reduce this, such as adding ML blocks only in the topmost layers.
2. The authors show that using the WAVE approach to reduce the gpu requirement retains 98.2% average performance. However, the more important metric is how much of the improvement is retained. Table 6 indicates that WAVE does not provide significant performance improvement.
3. There should be an analysis of GPU memory increment as the number of tasks during training increases.
4. From Table 3, task-aware training does not seem to improvement performance for most tasks.

Minor:
1. There are duplicate lines in Section 3.3.
2. Page 5 line 216: texty-type should be text-type

questions:
N/A
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a Mixture of Task Experts (MoTE) architecture for embedding models that allocates dedicated parameters for different tasks to improve dense text embeddings. The approach combines task-specific experts with a dynamic training strategy and weight averaging to achieve better task specialization without significantly increasing latency or memory overhead. Experiments on 56 datasets spanning multiple tasks demonstrate modest improvements over instruction-conditioned baselines.

The reviewers acknowledged some strengths of the work, including its well-conducted experiments (reviewers 161q and MkSg) and practical applicability to real-world scenarios (reviewer 161q). However, there were several major concerns. For example, reviewer Yimc pointed out that the work lacks innovation, primarily applying established techniques (MOE and Weight Averaging) without offering new insights. Multiple reviewers (ezqd and Yimc) noted that the performance gains were minimal - less than 1% improvement over the baseline in many cases. Reviewer ezqd raised concerns about memory requirements and the effectiveness of the weight averaging approach in retaining improvements.

Given these considerations and the average rating of 4.75, the paper is recommended for rejection.