大模型“取長補短”新思路入選NeurIPS’24，顯著優于現有路由方法，南科大港科大出品

西風 2024-11-14 17:24:43 來源：量子位

RouterDC團隊 投稿

量子位 | 公眾號 QbitAI

高效組合多個大模型“取長補短”新思路，被頂會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所示?？梢园l現：

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