大模型“取长补短”新思路入选NeurIPS'24,显著优于现有路由方法,南科大港科大出品

科技   2024-11-15 00:02   吉林  


MLNLP社区是国内外知名的机器学习与自然语言处理社区,受众覆盖国内外NLP硕博生、高校老师以及企业研究人员。
社区的愿景是促进国内外自然语言处理,机器学习学术界、产业界和广大爱好者之间的交流和进步,特别是初学者同学们的进步。
转载自 | 量子位
作者 | RouterDC团队

高效组合多个大模型“取长补短”新思路,被顶会NeurIPS 2024接收。

名为RouterDC,是一种基于双重对比学习的路由架构,具有参数高效性(小于100M的参数)和计算高效性(不需要对于LLM进行梯度回传)的优势。

在具有挑战性语言理解、代码生成和数学推理等推理任务实验中,RouterDC在分布内(+2.76%)和分布外(+1.90%)设定下,都远超于现有的routing方法。

众所周知,LLM通常在不同数据集上预训练和微调,导致它们在不同任务上的性能强弱不同。

LLM路由则是一种组合多个LLM的新思路,它通过学习一个路由器(Router)来为每一个请求(query)选择最合适的LLM。在推理时,LLM路由只需要调用所选的LLM进行推理,使其在保持计算高效性的同时利用多个LLM的互补能力。

RouterDC这种新方法,包括一个较小的语言模型作为编码器和一系列与候选LLM对应的可学习的LLM embeddings。

对于训练数据中的每个query,首先将候选LLM的预测与真实标签进行比较获得表现最好和最差的LLM,然后构造两个对比损失:

  • sample-LLM对比损失:使得query embedding(由编码器提取)与表现最佳的LLM embeddings相似,同时与表现最差的 LLM embeddings不相似。

  • sample-sample对比损失:提高训练的稳定性,将所有训练query聚类成多个组,最大化同组query之间的相似性的同时最小化不同组query之间的相似性。

这项研究由来自南方科技大学,香港科技大学的研究团队提出,以下是更为详细的介绍。

双对比学习实现Router训练

Router架构

如图1所示,RouterDC包括一个较小的语言模型(mDeBERTaV3-base)作为编码器ε,和一系列的与候选LLM对应的可学习LLM嵌入kT。对于每个query xi,RouterDC生成对于T个LLMs的选择概率如下:

其中,sim(·,·)表示cosine相似度。

△图1:RouterDC方法示意图

sample-LLM对比损失

为了训练router,研究者将query的样本嵌入和在其上表现最好的K+个LLM对应嵌入拉进,和在其上表现最差的K-个LLM对应嵌入拉远。因此,样本-LLM对比损失可以表示为:

sample-sample对比损失

研究者通过实验发现,在routing问题中只使用样本-LLM对比损失并不稳定,使得相似的query可能具有不相似的嵌入。

为了提升训练的鲁棒性,训练样本被聚类成不同的组,从而在训练中拉近同一个组内的样本,拉远不同组的样本。和样本-LLM对比损失类似,样本-样本对比损失可以公式化为:

训练及推理

最终的优化目标为最小化样本-LLM对比损失和样本-样本对比损失的结合:

推理时,每个测试query只需要通过训练好的router选取概率最大的LLM,并使用选择的LLM对query进行回答。

RouterDC在训练时不需要任何经过LLM的梯度回传,并且在推理时只需要调用进行一次LLM,同时具有训练和推理的高效性。

实验效果如何?

主要结果

RouterDC在分布内数据集的测试准确率结果如表1所示。可以发现:

RouterDC显著好于最优的单个模型,平均具有3.98%性能提升。在单个任务的层面,RouterDC在三个任务上相比表现最优的单个模型取得了准确率的提升,其中GSM8K提升了0.51%,ARC-C提升了0.57%,HumanEval提升了1.63%。

和现有路由方法CosineClassifier以及ZOOTER对比,RouterDC在所有任务上都具有更好的表现。和LoraRetriever对比,RouterDC具有平均2.77%的准确率提升。

△表1:分布内任务的测试准确率(%)

为了评估RouterDC的泛化能力,表2展示了RouterDC在三个分布外数据集(PreAlgebra,MBPP,C-EVAL)的测试准确率。

可以看出,RouterDC再次达到最高的测试准确率,显著超过表现最佳的单个LLM(dolphin-2.9-llama3-8b)1.9%。

△表2:分布外任务的测试准确率(%)

sample-sample损失的作用

为了探究样本-样本损失的作用,图3展示了在是否有样本-样本损失的条件下训练和测试准确率曲线。可以看出,RouterDC(w/o Lsample-sample)有明显的震荡现象,而RouterDC则稳定得多。

△图2:RouterDC在GSM8K任务上的训练和测试准确率曲线

图3(a)可视化了使用RouterDC(w/o Lsample-sample)提取的训练样本的TSNE特征,可以看到,属于不同任务的训练样本粗略地混合在一起。而在结合Lsample-sample之后,训练样本有了清晰的聚类结构(如图3(b)所示)。

△图3:学习到的router所提取出训练样本embedding的t-SNE可视化

RouterDC具有成本高效性

由于价格(cost)同样是一个评估LLM的重要指标,研究者通过RouterBench上的两个任务的实验来格外考虑cost的影响。如图16所示,RouterDC相比于CosineClassifier和ZOOTER更加的成本高效。

△图4:在RouterBench上使用不同的Cost获取的测试准确率

论文地址:https://arxiv.org/abs/2409.19886
代码地址:https://github.com/shuhao02/RouterDC


技术交流群邀请函

△长按添加小助手

扫描二维码添加小助手微信

请备注:姓名-学校/公司-研究方向
(如:小张-哈工大-对话系统)
即可申请加入自然语言处理/Pytorch等技术交流群

关于我们

MLNLP 社区是由国内外机器学习与自然语言处理学者联合构建的民间学术社区,目前已经发展为国内外知名的机器学习与自然语言处理社区,旨在促进机器学习,自然语言处理学术界、产业界和广大爱好者之间的进步。
社区可以为相关从业者的深造、就业及研究等方面提供开放交流平台。欢迎大家关注和加入我们。

机器学习算法与自然语言处理
关注AI前沿技术,助力AI学者进步
 最新文章