MoTE: Mixture of Task Experts for Embedding Models
Keywords: Embedding Models, Representation Learning, Mixture of Experts, Retrieval Augmented Generation, RAG, Search, Clustering, Classification
TL;DR: This paper builds on instruction embeddings literature introducing a new architecture with a dedicated set of parameters for each of the tasks and a task-aware training.
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.
Primary Area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 6960
Loading