Beyond instruction-conditioning, MoTE: Mixture of Task Experts for Multi-task Embedding Models

📄 arXiv: 2506.17781v1 📥 PDF

作者: Miguel Romero, Shuoyang Ding, Corey D. Barret, Georgiana Dinu, George Karypis

分类: cs.LG, cs.CL

发布日期: 2025-06-21

💡 一句话要点

提出MoTE：用于多任务嵌入模型的混合任务专家方法，提升低容量模型性能。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 多任务学习 嵌入模型 混合专家 指令条件化 对比学习 低容量模型 检索增强生成

📋 核心要点

  1. 指令条件化在低容量模型中存在表征约束，限制了嵌入专业化的性能提升。
  2. MoTE通过任务专用参数增强模型生成专业嵌入的能力，利用任务感知对比学习进行训练。
  3. 实验表明，MoTE在检索和所有数据集上均显著提升性能，且不增加额外开销。

📝 摘要（中文）

稠密嵌入是现代机器学习系统的基础，为检索增强生成（RAG）、信息检索和表征学习提供支持。虽然指令条件化已成为嵌入专业化的主要方法，但其在低容量模型上的直接应用施加了根本性的表征约束，限制了从专业化中获得的性能提升。本文分析了这些限制，并引入了混合任务专家（MoTE）Transformer块，该模块利用通过任务感知对比学习（TACL）训练的任务专用参数来增强模型生成专用嵌入的能力。实验结果表明，MoTE在检索数据集中的性能提升提高了64%（+3.27 → +5.21），在所有数据集中的性能提升提高了43%（+1.81 → +2.60）。关键是，这些提升是在不改变指令、训练数据、推理时间或激活参数数量的情况下实现的。

🔬 方法详解

问题定义：论文旨在解决低容量模型在多任务嵌入学习中，由于指令条件化的表征能力不足而导致的性能瓶颈。现有方法直接将指令作为条件输入模型，但对于参数量有限的模型，这种方式难以充分捕捉不同任务的特性，从而限制了嵌入的专业化程度。

核心思路：论文的核心思路是引入混合任务专家（MoTE）机制，为每个任务分配一组专门的参数（即专家），并通过门控机制动态地选择合适的专家组合来生成嵌入。这种方式允许模型针对不同任务学习到更具针对性的表征，从而提升整体性能。

技术框架：MoTE作为一个Transformer块，可以嵌入到现有的Transformer模型中。其主要包含以下几个模块：1) 多个任务专家（Task Experts）：每个专家都是一个独立的神经网络层，负责处理特定任务的信息。2) 门控网络（Gating Network）：根据输入指令，动态地为每个专家分配权重，决定哪些专家参与最终的嵌入生成。3) 任务感知对比学习（TACL）：用于训练任务专家的损失函数，鼓励模型学习到区分不同任务的表征。

关键创新：MoTE的关键创新在于其任务专用参数的设计和动态专家选择机制。与传统的指令条件化方法相比，MoTE能够更有效地利用模型参数，为每个任务学习到更具针对性的表征。此外，MoTE的门控机制允许模型在不同任务之间共享知识，从而提高模型的泛化能力。

关键设计：MoTE中的任务专家可以是任何类型的神经网络层，例如前馈网络或卷积层。门控网络通常是一个简单的线性层或多层感知机，其输入是指令嵌入，输出是每个专家的权重。TACL损失函数通常基于对比学习的思想，鼓励模型将同一任务的样本拉近，将不同任务的样本推远。具体的参数设置，如专家数量、门控网络的结构和TACL损失函数的具体形式，需要根据具体任务进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，MoTE在检索数据集上取得了显著的性能提升，相较于基线模型，性能提升高达64%（+3.27 → +5.21）。在所有数据集上的平均性能提升也达到了43%（+1.81 → +2.60）。值得注意的是，这些提升是在不改变指令、训练数据、推理时间或激活参数数量的情况下实现的，充分证明了MoTE的效率和有效性。

🎯 应用场景

MoTE技术可广泛应用于需要多任务嵌入的场景，如检索增强生成（RAG）、信息检索、推荐系统和自然语言理解等。通过提升低容量模型的性能，MoTE能够降低计算成本，提高部署效率，并为资源受限的设备提供更强大的AI能力。未来，MoTE有望成为构建高效、灵活的多任务嵌入模型的重要组成部分。

📄 摘要（原文）

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 (\tacl) to enhance the model 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.