RouterDC: Query-Based Router by Dual Contrastive Learning for Assembling Large Language Models

📄 arXiv: 2409.19886v1 📥 PDF

作者: Shuhao Chen, Weisen Jiang, Baijiong Lin, James T. Kwok, Yu Zhang

分类: cs.LG, cs.AI, cs.CL

发布日期: 2024-09-30

备注: Accepted by NeurIPS 2024

🔗 代码/项目: GITHUB

💡 一句话要点

RouterDC：通过双重对比学习的查询式路由，用于组装大型语言模型

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 模型组装 路由选择 对比学习 查询编码

📋 核心要点

  1. 现有路由模型在多个LLM对同一查询表现良好时效果不佳，无法有效利用LLM的互补能力。
  2. RouterDC通过双重对比学习，学习查询和LLM嵌入之间的关系，从而选择最合适的LLM。
  3. 实验表明，RouterDC在同分布和异分布任务上均显著优于现有方法和单个最佳LLM。

📝 摘要（中文）

本文提出了一种名为RouterDC的查询式路由方法，该方法通过双重对比学习来组装多个现成的大型语言模型（LLMs），以利用它们互补的能力。现有的路由模型在多个LLM对同一查询表现良好时效果不佳。为了解决这个问题，RouterDC模型由一个编码器和LLM嵌入组成，并提出了两个对比学习损失来训练RouterDC模型。实验结果表明，RouterDC在组装LLM方面是有效的，并且在同分布（+2.76%）和异分布（+1.90%）任务上都大大优于单个表现最佳的LLM以及现有的路由方法。

🔬 方法详解

问题定义：现有路由方法在多个LLM都能很好地处理同一查询时，无法有效区分并选择最佳的LLM。这导致模型无法充分利用不同LLM的互补优势，影响整体性能。现有方法的痛点在于缺乏对查询和LLM之间细粒度关系的建模能力。

核心思路：RouterDC的核心思路是通过对比学习，学习查询和LLM嵌入之间的关系。具体来说，模型学习将相似的查询与表现良好的LLM的嵌入拉近，而将不相似的查询与表现不佳的LLM的嵌入推远。通过这种方式，模型能够更准确地判断哪个LLM最适合处理给定的查询。

技术框架：RouterDC模型主要包含两个部分：一个编码器和一个LLM嵌入层。编码器负责将输入查询编码成一个向量表示。LLM嵌入层包含所有候选LLM的嵌入向量。训练过程中，模型通过双重对比学习来优化编码器和LLM嵌入，使得查询的编码向量能够更好地匹配到合适的LLM嵌入。在推理阶段，模型计算查询编码向量与所有LLM嵌入之间的相似度，选择相似度最高的LLM来处理该查询。

关键创新：RouterDC的关键创新在于提出了双重对比学习策略。一方面，模型通过查询对比学习，学习区分相似和不相似的查询。另一方面，模型通过LLM对比学习，学习区分适合和不适合处理特定查询的LLM。这两种对比学习策略相互促进，使得模型能够更准确地学习查询和LLM之间的关系。与现有方法相比，RouterDC能够更好地处理多个LLM都能很好地处理同一查询的情况。

关键设计：RouterDC使用了两个对比学习损失函数：查询对比损失和LLM对比损失。查询对比损失的目标是拉近相似查询的编码向量，推远不相似查询的编码向量。LLM对比损失的目标是拉近查询编码向量和适合处理该查询的LLM的嵌入向量，推远查询编码向量和不适合处理该查询的LLM的嵌入向量。具体来说，可以使用InfoNCE损失函数来实现这两个对比学习损失。此外，编码器可以使用预训练的Transformer模型，例如BERT或RoBERTa。LLM嵌入可以是随机初始化的，也可以使用预训练的LLM的输出表示。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，RouterDC在同分布任务上比最佳的单个LLM提高了2.76%，在异分布任务上提高了1.90%。此外，RouterDC也显著优于现有的路由方法，证明了其在组装LLM方面的有效性。这些结果表明，RouterDC能够有效地利用多个LLM的互补优势，从而提高整体性能。

🎯 应用场景

RouterDC可应用于各种需要组合多个LLM的场景，例如智能客服、内容生成、机器翻译等。通过动态选择最合适的LLM，可以提高系统的整体性能和效率。该研究有助于构建更强大、更灵活的LLM应用，并推动LLM在实际场景中的广泛应用。

📄 摘要（原文）

Recent works show that assembling multiple off-the-shelf large language models (LLMs) can harness their complementary abilities. To achieve this, routing is a promising method, which learns a router to select the most suitable LLM for each query. However, existing routing models are ineffective when multiple LLMs perform well for a query. To address this problem, in this paper, we propose a method called query-based Router by Dual Contrastive learning (RouterDC). The RouterDC model consists of an encoder and LLM embeddings, and we propose two contrastive learning losses to train the RouterDC model. Experimental results show that RouterDC is effective in assembling LLMs and largely outperforms individual top-performing LLMs as well as existing routing methods on both in-distribution (+2.76\%) and out-of-distribution (+1.90\%) tasks. Source code is available at https://github.com/shuhao02/RouterDC.