大模型“取長(zhǎng)補(bǔ)短”新思路入選NeurIPS’24，顯著優(yōu)于現(xiàn)有路由方法，南科大港科大出品

西風(fēng) 2024-11-14 17:24:43 來(lái)源：量子位

RouterDC團(tuán)隊(duì) 投稿

量子位 | 公眾號(hào) QbitAI

高效組合多個(gè)大模型“取長(zhǎng)補(bǔ)短”新思路，被頂會(huì)NeurIPS 2024接收。

名為RouterDC，是一種基于雙重對(duì)比學(xué)習(xí)的路由架構(gòu)，具有參數(shù)高效性（小于100M的參數(shù)）和計(jì)算高效性（不需要對(duì)于LLM進(jìn)行梯度回傳）的優(yōu)勢(shì)。

在具有挑戰(zhàn)性語(yǔ)言理解、代碼生成和數(shù)學(xué)推理等推理任務(wù)實(shí)驗(yàn)中，RouterDC在分布內(nèi)（+2.76%）和分布外（+1.90%）設(shè)定下，都遠(yuǎn)超于現(xiàn)有的routing方法。

眾所周知，LLM通常在不同數(shù)據(jù)集上預(yù)訓(xùn)練和微調(diào)，導(dǎo)致它們?cè)诓煌蝿?wù)上的性能強(qiáng)弱不同。

LLM路由則是一種組合多個(gè)LLM的新思路，它通過(guò)學(xué)習(xí)一個(gè)路由器（Router）來(lái)為每一個(gè)請(qǐng)求（query）選擇最合適的LLM。在推理時(shí)，LLM路由只需要調(diào)用所選的LLM進(jìn)行推理，使其在保持計(jì)算高效性的同時(shí)利用多個(gè)LLM的互補(bǔ)能力。

RouterDC這種新方法，包括一個(gè)較小的語(yǔ)言模型作為編碼器和一系列與候選LLM對(duì)應(yīng)的可學(xué)習(xí)的LLM embeddings。

對(duì)于訓(xùn)練數(shù)據(jù)中的每個(gè)query，首先將候選LLM的預(yù)測(cè)與真實(shí)標(biāo)簽進(jìn)行比較獲得表現(xiàn)最好和最差的LLM，然后構(gòu)造兩個(gè)對(duì)比損失：

• sample-LLM對(duì)比損失：使得query embedding（由編碼器提?。┡c表現(xiàn)最佳的LLM embeddings相似，同時(shí)與表現(xiàn)最差的 LLM embeddings不相似。
• sample-sample對(duì)比損失：提高訓(xùn)練的穩(wěn)定性，將所有訓(xùn)練query聚類成多個(gè)組，最大化同組query之間的相似性的同時(shí)最小化不同組query之間的相似性。

這項(xiàng)研究由來(lái)自南方科技大學(xué)，香港科技大學(xué)的研究團(tuán)隊(duì)提出，以下是更為詳細(xì)的介紹。

雙對(duì)比學(xué)習(xí)實(shí)現(xiàn)Router訓(xùn)練

Router架構(gòu)

如圖1所示，RouterDC包括一個(gè)較小的語(yǔ)言模型（mDeBERTaV3-base）作為編碼器ε，和一系列的與候選LLM對(duì)應(yīng)的可學(xué)習(xí)LLM嵌入kT。對(duì)于每個(gè)query xi，RouterDC生成對(duì)于T個(gè)LLMs的選擇概率如下：

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

△圖1：RouterDC方法示意圖

sample-LLM對(duì)比損失

為了訓(xùn)練router，研究者將query的樣本嵌入和在其上表現(xiàn)最好的K+個(gè)LLM對(duì)應(yīng)嵌入拉進(jìn)，和在其上表現(xiàn)最差的K-個(gè)LLM對(duì)應(yīng)嵌入拉遠(yuǎn)。因此，樣本-LLM對(duì)比損失可以表示為：

sample-sample對(duì)比損失

研究者通過(guò)實(shí)驗(yàn)發(fā)現(xiàn)，在routing問(wèn)題中只使用樣本-LLM對(duì)比損失并不穩(wěn)定，使得相似的query可能具有不相似的嵌入。

為了提升訓(xùn)練的魯棒性，訓(xùn)練樣本被聚類成不同的組，從而在訓(xùn)練中拉近同一個(gè)組內(nèi)的樣本，拉遠(yuǎn)不同組的樣本。和樣本-LLM對(duì)比損失類似，樣本-樣本對(duì)比損失可以公式化為：

訓(xùn)練及推理

最終的優(yōu)化目標(biāo)為最小化樣本-LLM對(duì)比損失和樣本-樣本對(duì)比損失的結(jié)合：

推理時(shí)，每個(gè)測(cè)試query只需要通過(guò)訓(xùn)練好的router選取概率最大的LLM，并使用選擇的LLM對(duì)query進(jìn)行回答。

RouterDC在訓(xùn)練時(shí)不需要任何經(jīng)過(guò)LLM的梯度回傳，并且在推理時(shí)只需要調(diào)用進(jìn)行一次LLM，同時(shí)具有訓(xùn)練和推理的高效性。

實(shí)驗(yàn)效果如何？

主要結(jié)果

RouterDC在分布內(nèi)數(shù)據(jù)集的測(cè)試準(zhǔn)確率結(jié)果如表1所示。可以發(fā)現(xiàn)：

RouterDC顯著好于最優(yōu)的單個(gè)模型，平均具有3.98%性能提升。在單個(gè)任務(wù)的層面，RouterDC在三個(gè)任務(wù)上相比表現(xiàn)最優(yōu)的單個(gè)模型取得了準(zhǔn)確率的提升，其中GSM8K提升了0.51%，ARC-C提升了0.57%，HumanEval提升了1.63%。

和現(xiàn)有路由方法CosineClassifier以及ZOOTER對(duì)比，RouterDC在所有任務(wù)上都具有更好的表現(xiàn)。和LoraRetriever對(duì)比，RouterDC具有平均2.77%的準(zhǔn)確率提升。

△表1：分布內(nèi)任務(wù)的測(cè)試準(zhǔn)確率（%）

為了評(píng)估RouterDC的泛化能力，表2展示了RouterDC在三個(gè)分布外數(shù)據(jù)集（PreAlgebra，MBPP，C-EVAL）的測(cè)試準(zhǔn)確率。

可以看出，RouterDC再次達(dá)到最高的測(cè)試準(zhǔn)確率，顯著超過(guò)表現(xiàn)最佳的單個(gè)LLM（dolphin-2.9-llama3-8b）1.9%。

△表2：分布外任務(wù)的測(cè)試準(zhǔn)確率（%）

sample-sample損失的作用

為了探究樣本-樣本損失的作用，圖3展示了在是否有樣本-樣本損失的條件下訓(xùn)練和測(cè)試準(zhǔn)確率曲線?？梢钥闯?，RouterDC（w/o Lsample-sample）有明顯的震蕩現(xiàn)象，而RouterDC則穩(wěn)定得多。

△圖2：RouterDC在GSM8K任務(wù)上的訓(xùn)練和測(cè)試準(zhǔn)確率曲線

圖3（a）可視化了使用RouterDC（w/o Lsample-sample）提取的訓(xùn)練樣本的TSNE特征，可以看到，屬于不同任務(wù)的訓(xùn)練樣本粗略地混合在一起。而在結(jié)合Lsample-sample之后，訓(xùn)練樣本有了清晰的聚類結(jié)構(gòu)（如圖3（b）所示）。

△圖3：學(xué)習(xí)到的router所提取出訓(xùn)練樣本embedding的t-SNE可視化

RouterDC具有成本高效性

由于價(jià)格（cost）同樣是一個(gè)評(píng)估LLM的重要指標(biāo)，研究者通過(guò)RouterBench上的兩個(gè)任務(wù)的實(shí)驗(yàn)來(lái)格外考慮cost的影響。如圖16所示，RouterDC相比于CosineClassifier和ZOOTER更加的成本高效。

△圖4：在RouterBench上使用不同的Cost獲取的測(cè)試準(zhǔn)確率

論文地址：https://arxiv.org/abs/2409.19886
代碼地址：https://github.com/shuhao02/RouterDC