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

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

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

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

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

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

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

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

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

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

Router架構(gòu)

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

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

△圖1：RouterDC方法示意圖

sample-LLM對比損失

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

sample-sample對比損失

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

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

訓(xùn)練及推理

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

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

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

實驗效果如何？

主要結(jié)果

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

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

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

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

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

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

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

sample-sample損失的作用

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

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

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

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

RouterDC具有成本高效性

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

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

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