MoTE: Mixture of Task Experts for Embedding Models

Thank you for your thoughtful and detailed review. We appreciate your recognition of the novel Mixture of Task Experts (MoTE) architecture and the insights provided by our experimental analysis.

Weaknesses

1. We appreciate your concern about the choice of baseline models. Our rationale for selecting Nussbaum et al. (https://arxiv.org/pdf/2406.18587v1) as our primary baseline is three folded. First, reproducibility - Nussbaum et al. provide intermediate checkpoints, hyper-parameters, architecture, pre-processed data, and publicly available code. Given that MoTE reuse most of these artifacts, this baseline eases the reproducibility of our results. Second, relevance to production settings - smaller models, like the one from Nussbaum et al., are more practical in production due to their lower latency, inference cost and vector storage cost. Demonstrating MoTE’s benefits in this context ensures a more tight alignment with real-world deployment scenarios. Third, performance ranking - after both contrastive training phases, the Nussbaum et al. model achieves a top-6 ranking in its size category, performing better than E5 but slightly below BGE. This ensures it remains a competitive baseline for our experiments.

2. Thank you for this suggestion. We have expanded Sections 4.5 and 4.6 (now part of the Appendix) to include a deeper analysis alongside memory analysis. Additionally, we have included similarity analysis across the task embedding space to provide readers more insight into both the current limitation of instruction-conditioned embedding models and how MoTE overcomes them.

3. We acknowledge the writing issue and appreciate your feedback. We have reviewed the entire manuscript to improve its presentations and readability.

Questions

1. We do not find the efficacy of the expert averaging (what we referred as WAVE in the initial submission) method surprising, as it aligns conceptually with established model averaging techniques, which are known to enhance performance and generalization capabilities (e.g., https://arxiv.org/pdf/2302.07027, https://arxiv.org/pdf/2205.09739). To provide additional context, we have added a background section on model averaging (Section 2.4) in the revised manuscript, as also suggested by another reviewer.

2. This is an excellent point, and we appreciate your insight. While our current experiments focus on encoder embedding models, we agree that exploring MoTE’s applicability to LLM-based retrieval models and hope that this paper encourages future research in this direction.

Thank you again for your constructive feedback. We hope these updates address your concerns and provide additional clarity. We are happy to further refine the paper based on additional suggestions.

Thank you for your thoughtful review, feedback and for highlighting the strengths of our paper.

Weaknesses

1. We would like to clarify that our goal is not to invent new MoE architectures, but to show that adaptations of MoE can be leveraged to address the limitations of instruction-conditioned embedding models, which leads to clear performance improvements as demonstrated in our results. Besides, while MoE has been widely adopted in LLMs, it is rare to see them applied to embedding models and has not been previously used to address the multi-task aspect of such models. Although this approach may not be feasible for LLM applications, it fits naturally into the training and evaluation paradigm of embedding models. By applying MoE in this context, we offer a fresh perspective on the interplay between MoE, instruction-conditioned embedding models and multi-task learning. However, we acknowledge that this contribution could have been more clearly articulated in the original submission. To mitigate this weakness we clarified the area of our contributions in the Introduction.

2. Regarding weaknesses (2) - (4): Thank you for the detailed feedback. We mitigated weakness (2) in the new review by expanding the Methods section to provide a detailed description of the implementation of task-aware training, up-cycling and MoTE. However, if any aspect remains unclear, we are happy to iterate further based on your suggestions. Regarding weakness (4) - We acknowledge the lack of foundational background on weight averaging and based on your suggestion have included a dedicated subsection in the revised manuscript. Lastly, the expanded methodology and background section, alongside a detailed analysis in task similarity using MoTE and traditional instruction-conditioned model have expanded the manuscript to the $10$ page limit with $3$ and $2$ additional pages for reference and appendix, respectively (addressing weakness (3)).

Questions

1. Thank you for pointing this out. We have clarified our hyper-parameter selection process in the Experimental Methodology section. Specifically, we use hyper-parameters from Nussbaum et al (https://arxiv.org/pdf/2406.18587v1) with a batch size of 6144 to ease reproduction. We only ablate the contrastive temperature and batching strategy parameters using their base transformer architecture with static training.

2. These are great questions. Thanks for asking. The specialized weights in MoTE are trained with data specific to each task, optimizing representations inherently for those tasks. We included Section 5.3 to compare to what extent the embedding spaces differ between tasks for both instruction-conditioned embedding models and MoTE. We observe that in most cases the task spaces produced by instruction-conditioned embedding models exhibit a statistically higher inter-task similarity compared to MoTE (Figure $4$). This demonstrates the existing limitations of optimizing a single set of weights for all tasks and highlights how MoTE achieves higher performance by relaxing this constraint while maintaining the same number of active parameters per task. To our surprise this was not the case for clustering and retrieval documents for which our analysis shows that those task spaces hold a higher similarity (Figure $4$ - bottom, right). Thus, we believe that MoTE could be further memory optimizations by using one expert and different instructions for retrieval documents and clustering. We leave this line of research for future work.

Thank you again for your constructive feedback, and we hope these updates and explanations address your concerns.

Thank you for your thorough review and constructive feedback. We appreciate your recognition of the innovative methods, practical applicability, and strong experimental setup in our paper.

Weaknesses

1. While we acknowledge that our work builds upon existing mixture-of-expert (MoE) architectures, we believe it represents a significant advancement in applying MoE to embedding models. Specifically, MoTE introduces task-specific embeddings with clear improvements in task-aware applications, addressing key limitations of prior approaches. However, we recognize that this contribution could have been more clearly articulated in the original submission. To address this, we have expanded our analysis in the revised draft to better contextualize the limitations of prior methods and demonstrate how MoTE mitigates these challenges. These additional studies further highlight MoTE's ability to effectively enhance embedding model expressiveness and performance across diverse downstream tasks.

2. We think there might have been some misunderstandings regarding Table 2 (now Table 3 in the revised draft). MoTE, our proposed method, in fact consistently demonstrates improvements across the majority of tasks, including task-aware applications such as classification, clustering, and retrieval. While the semantic textual similarity (STS) task is the one out-of-scope task where MoTE shows only marginal gains, it is important to note that MoTE outperforms prior methods, including TEM-Inst, on all other tasks. We hope this clarification resolves any confusion and highlights the broad applicability of our approach.

Please let us know if we understood your concerns correctly and if the above explanation and attached review addresses them. Thank you once again for your insightful feedback and for helping us improve the clarity and rigor of our work.

Thank you for your detailed and insightful review. We greatly appreciate your recognition of the strengths of our approach and your constructive feedback on areas that could be improved. We have carefully addressed the concerns you raised and made corresponding updates to the paper:

Weaknesses

1. You are absolutely correct that the inclusion of task-specialized experts introduces an inherent memory cost, even as we maintain low inference latency. We acknowledge that our initial draft could have explained this trade-off more clearly. To address this:

   1. We have clarified in the Introduction that Expert Average (EA) (previously referred to as WAVE) represents an initial step toward mitigating memory requirements.

   2. As you noted, the decision to upcycle experts in the EOB configuration was arbitrary. We have added a paragraph in the Methods section explaining that practitioners have the flexibility to customize what dense blocks to up-cycle into MoTE blocks —whether in all layers, every other layer, the topmost layer, etc.— based on their memory and performance requirements.

   3. While EA establishes a baseline for balancing memory and performance, we acknowledge that more advanced memory optimization techniques (e.g., reusing experts across tasks) are promising directions for future research. We have elaborated on this in Section 5.3.

2. We revised the discussion of EA to better communicate its role and impact. EA is primarily intended as a baseline for reducing memory overhead while retaining a significant portion of the performance gains achieved by MoTE. Specifically, EA achieves 58.60 average performance points compared to the 59.07 points of MoE with respect to the 58.43 points of the base instruction-conditioned model while maintaining the same memory footprint.

   1. In response to your feedback, we added a row to Table 5 showing the average performance per dataset for clarity and removed the percentage retention metric to avoid potential confusion.

3. As suggested, we have included a GPU memory consumption analysis as a function of the number of tasks (experts) during training. This analysis is provided in Appendix C, comparing both the EB and EOB configurations.

4. We think there might have been a misunderstanding here. While task-aware training yields an average improvement of +0.29 over the baseline, we recognize that this trend was not sufficiently highlighted in the original draft. We have updated Table 4 to include the average performance across datasets and expanded the results analysis paragraph with a more detailed explanation, making this improvement more evident.

Additionally, we also corrected the paper writing to improve clarity and comprehension. We hope that these revisions address your concerns and clarify the contributions and significance of our work. Thank you again for your valuable feedback, which has helped improve the quality of our paper.

As the discussion period nears its conclusion, we would like to follow up to ensure that our responses have satisfactorily addressed your feedback and clarified the concerns raised in your review. Based on your insightful comments, we have made several key updates and improvements to the manuscript:

1. We clarified in the Introduction that Expert Average (EA), formerly referred to as WAVE, is an initial strategy to balance memory and performance. We also expanded the Methods section to outline how practitioners can flexibly decide which dense blocks to up-cycle into task-specific experts, providing options to balance the performance and memory footprint trade-off.

2. As suggested, we conducted an analysis of GPU memory consumption as the number of tasks increases during training. This detailed analysis is now included in Appendix C, covering both EB (Every Blocks) and EOB (Every-Other Block) configurations.

3. We revised the discussion of EA’s role to highlight its utility in significantly reducing memory overhead while retaining substantial performance gains. To improve clarity, we added dataset-specific averages to Table 5 and removed the confusing metric that you kindly highlighted.

4. We addressed concerns regarding task-aware training by explicitly emphasizing its average improvement of +0.29 points over the baseline. Table 4 now includes average performance across datasets, and we expanded the accompanying discussion to make the trend more evident.

5. We corrected duplicated lines in Section 3.3, addressed minor errors (e.g., "texty-type"), and improved clarity throughout the paper.

Thank you for recognizing the strengths of our approach and for providing such constructive feedback. If there are any additional concerns or areas for further clarification, we would be more than happy to address them.

We sincerely appreciate your time and effort in helping us enhance the quality of our work. We look forward to hearing your final thoughts.