Beyond instruction-conditioning, M o TE : Mixture of Task Experts for Multi-task Embedding Models (2025)

Abstract: Dense embeddings are fundamental to modern machine learning systems, powering Retrieval-Augmented Generation (RAG), information retrieval, and representation learning. While instruction-conditioning has become the dominant approach for embedding specialization, its direct application to low-capacity models imposes fundamental representational constraints that limit the performance gains derived from specialization. In this paper, we analyze these limitations and introduce the Mixture of Task Experts (MoTE) transformer block, which leverages task-specialized parameters trained with Task-Aware Contrastive Learning () to enhance the model’s ability to generate specialized embeddings. Empirical results show that MoTE achieves 64% higher performance gains in retrieval datasets ( +3.27\rightarrow +5.21 ) and 43% higher performance gains across all datasets ( +1.81\rightarrow 2.60 ). Critically, these gains are achieved without altering instructions, training data, inference time, or number of active parameters.

## introduction semantic text representations are a key component to many real-world applications such as search, recommendation systems, and spam classification. these representations are commonly obtained using embedding models that map unstructured text into a dense n-dimensional vector referred to as dense embedding. contemporary embedding models are trained by simultaneously optimizing the model towards a wide range of downstream tasks including retrieval augmented generation (rag) @xcite , search @xcite , semantic text similarity (sts) @xcite , classification @xcite , and clus-tering @xcite . as a result, a single embedding model is able to generate a dense text representation that can be used in all these downstream tasks. historically, dense embedding models have been trained using a multi-task curriculum, where a single model is optimized across multiple datasets, each corresponding to a specific task such as clustering or retrieval. in this framework, a shared embedding space is used across all tasks, allowing a single representation to serve multiple downstream applications. however, different tasks impose distinct semantic similarity requirements, often leading to conflicting objectives. for instance, in a semantic textual similarity (sts) task, the questions "who was isaac newton?" and "who was the father of calculus?" should be close in the embedding space since they refer to the same person. in contrast, a retrieval task requires these queries to be far apart, as they can not answer each other @xcite . this presents a fundamental challenge in embedding model development, where a single embedding space must be carefully tuned to support multiple tasks without compromising performance on any of them @xcite . to address this, modern embedding models have shifted foot_0 toward generating task-specific embeddings, allowing representations to be adapted based on the downstream application. the consistent performance gains observed from embedding specialization, coupled with its widespread adoption in state-of-the-art models @xcite have provide strong evidence of the limitations of generic embeddings in downstream tasks. the predominant approach for embedding specialization is instruction-conditioning (ic) @xcite , where the input text is enriched with task-specific instructions to steer the generation of specialized embeddings. for instance, @xcite employs the instruction "query:" for queries and "passage:" for passages in retrieval tasks, providing the model with information to distinguish between different roles in the retrieval tasks. with this approach, the specialization signal is entirely encoded in the instruction tokens, which propagate through the model to produce the final embeddings. however, theoretical properties of neural networks, such as lipschitz continuity, suggest that small changes in input-like adding short instruction tokens-produce only minor shifts in output embeddings @xcite , limiting a model's ability to create disentangled, task-specific representations. this constraint is amplified in smaller models and with short prompts, where lowdimensional embeddings face capacity bottlenecks @xcite , forcing trade-offs between preserving semantic and task-relevant information. these limitations are compounded during multitask contrastive training, where gradient interference @xcite and conflicting objectives @xcite saturate model capacity, and where hyper-parameters like batching strategy and contrastive temperature must be carefully tuned to align with downstream task semantics, further complicating effective task-specific embedding specialization. in this paper, we study these limitations in the context of multi-task embedding models and propose alternative methods to enhance embedding specialization. our key contributions include: experiments conducted on 56 datasets across 7 tasks @xcite demonstrate that leveraging mote and ta-cl enables more effective utilization of task information for embedding specialization (figure 1 ). this approach leads to 50% higher performance gains in critical tasks such as retrieval and 31% higher gains in other tasks compared to instruction-conditioned specialization. critically, these improvements are achieved with identical input information, training data, latency, and number of active parameters. furthermore, exploiting the observation that embedding models are typically used for one downstream task at a time (e.g., retrieval-indexing, retrieval-querying, classification), mote can maintain the same gpu memory footprint as instruction-conditioned models by offloading inactive task experts to cpu or disk. this ensures computational efficiency without sacrificing performance. ## method while instruction-conditioning enables the generation of specialized embeddings f i θ (i, s), relying solely on the input instruction to steer the specialization of the embedding while the other parameters are shared can limit the models ability to generate specialized embeddings @xcite . to overcome this constraints we propose to introduce task specialized parameters to increase the model's capacity to specialize embeddings while maintaining the overall number of active parameters. formally, we propose a function ψ i [θ 0 |θ i ] such that: consistent with ic, ψ i [θ 0 |θ i ] accepts tuples (i, s) where @xmath0.g., "query: ", "passage: ") and s denotes the input sequence. however, unlike standard instruction-conditioning, ψ i [θ 0 |θ i ] dynamically selects the appropriate set of specialized parameters θ i to route the input sequence [i|s] and enhancing task specialization. ## experimental methodology we evaluate performance using the massive text embedding benchmark (mteb) @xcite , which spans 7 tasks and 56 datasets (see we evaluate the embedding specialization benefits derived from instruction-conditioning (ic) and mote on nussbaum et al. (2024)'s model foot_2 . for more information about the metrics please refer to appendix c. both candidates are initialized using the same pretrained checkpoint foot_3 , and trained on the same classification, clustering and retrieval datasets foot_4 during contrastive learning foot_5 . all candidates are trained for a single epoch using the infonce objective the ic and mote candidates leverage the same instructions: "classification: ", "clustering: ", "search query: ", and "search document: ". additionally, the mote candidate is implemented by using mote blocks in every transformer block and four experts in each transformer block, each of which is associated with one of the above instructions. further implementation details are provided in appendix d. additionally, to assess the performance benefits of multi-task embedding models with task-specific experts compared to specialized single-task models, we train and evaluate three specialized embedding models (sem). each sem uses the same underlying transformer architecture as the em and iem candidates but is trained independently on one of the classification, clustering, or retrieval subsets of the training data. performance is then reported for each sem on its corresponding downstream task. 5 results ## ablation: moe architecture design contemporary moe literature introduces two strategies for integrating moe blocks into transformer architectures, each offering a different balance between local and global parametrizations. specifically, moe blocks can be incorporated into every transformer block (eb) @xcite or at every other block (eob) @xcite . in this experiment, we evaluate the performance impact of these design choices when extending embedding architectures with mote blocks. ## conclusion in this paper, we identified the limitations of instruction conditioning for embedding specialization and introduced mote and ta-cl as alternative approaches to overcome these constraints. our results demonstrate that mote significantly improves overall performance-particularly in critical tasks such as retrieval-without increasing computational costs, inference latency, or the number of active parameters. by leveraging task-aware contrastive learning and task-based routing, mote enhances embedding specialization while maintaining computational efficiency. additionally, our findings suggest that a hybrid approach-combining mote with instruction conditioning-could be more effective than either method alone. while mote improves task disentanglement, we observed that in certain cases, preserving synergies between related tasks is equally important for performance. for example,