一、背景

本文中我們繼續介紹一種 LLM 推理優化相關的工作，通過路由的方式組合多個模型；其與投機采樣類似，通過多個不同規模和性能的模型組合來降本增效，然而又有本質的區別。投機采樣在一個 Query 內會反復調用大小模型，而路由方式在調用之前已經確定好需要調用哪個模型，直到調用結束。

目前常見有兩種路由的范式：

• 按意圖路由：與傳統意圖識別思路類似。其思路是雖然小模型可能整體實力不如大模型，但是在某些垂類可能與大模型相當，比如代碼、數學等，此時如果判斷是代碼相關 Query 就可以直接路由到專業的代碼小模型。
• 按難易路由：其核心思路是說小模型雖然處理復雜問題能力不行，但是處理簡單問題時與大模型相當，那么簡單問題用小模型足以。比如 LeetCode 的 Easy 題目讓小模型做即可，Hard 題目還是交給大模型比較靠譜。

當然，我們并不認為下面要介紹的部分方法已經充分達到可以在實際業務場景中落地應用的程度，這里只是提供一種新的優化思路。比如說，當前的很多方案還都是在有損甚至比較大損失的前提下，也許通過模型壓縮量化實現的收益和路由的方式類似，反而實現起來更加簡單。此外，也有很多方面是值得繼續探索的，比如：

• 針對更多模型的 Router 方案。
• 提升 Router 魯棒性，進一步擴展對不同數據分布下的適應性。
• 成本與質量更好的動態平衡，比如說流量比較高、服務壓力比較大時可以降低質量要求，更多的流量到達小模型，壓力比較小時可以更多流量到達大模型。

二、引言

2.1 投機采樣

投機采樣核心思路如下圖所示，首先以低成本的方式快速生成多個候選 Token（小模型，多頭，檢索，Early Exit 等方式），然后通過一次并行驗證階段快速驗證多個 Token，進而減少大模型的 Decoding Step，實現加速的目的：

投機采樣可以有效減少 Decoding Step 數量，這也是其存在的意義，然而驗證的接受率會很大程度上影響最終的加速比，接受率越高，減少的 Decoding Step 數量就越多，因未接收而浪費的計算就越少（實際上只要不是接受率 100%，就一定存在計算的浪費）。除此之外，當序列比較長時，由于減少 Decoding Step 而減少的對全局 KV Cache 的訪問更加可觀，相當于在 Memory Bound 的時候用 Compute 換 IO。

2.2 BART Score

BART Score（[2106.11520] BARTScore: Evaluating Generated Text as Text Generation） 是一種用于評價自然語言生成任務中模型輸出質量的指標，它是基于 BART 模型的（BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension）。BART 模型在預訓練時，會使用多種噪聲對原始文本進行破壞，然后通過雙向 Transformer 模型重建原始文本。

BART Score 的優勢在于它利用了預訓練的 BART 模型所捕捉到的豐富語言信息和上下文關系，能夠更好地反映生成文本的質量，特別是在考慮句子間的語義相似性和一致性方面。相比傳統的 BLEU、ROUGE 等指標，BART Score 更加靈活和精確，能夠更好地捕捉到文本生成任務中的細微差別。

如下圖所示為 BART Score 的計算公式，其中 w 為權重，值大于等于 0；p 為概率，介于 [0, 1] 之間，因此 BART Score 始終小于等于 0：

在使用中，BART Score 有幾種變體，用于不同的生成方向，具體包括：

• Faithfulness(s -> h)：從源文本到假設（生成文本）的概率 p(h|s, θ)。衡量從源文本生成假設文本的可能性。
• Precision(r -> h)：從參考文本生成假設的概率 p(r|h, θ)。衡量基于 Gold 參考文本生成假設文本的可能性。
• Recall(h -> r): 從生成的文本到參考文本的概率 p(r|h, θ)。衡量從 Gold 參考文本能多容易地生成假設文本。
• F score(r <-> h): 同時考慮 Precision 和 Recall 兩個方向，并使用它們的算術平均值。這個版本可以廣泛用于參考文本和生成文本之間的語義重疊。

二、SambaNova CoE

如下圖 Fig 2 所示（來自 [2405.07518] SambaNova SN40L: Scaling the AI Memory Wall with Dataflow and Composition of Experts），SambaNova 很早之前就提出了 CoE（Composition of Experts） 的方案，其 Samba-CoE 可能包含數十個專家模型，這些專家模型都來自開源社區，作者在編碼、數學和翻譯等特定領域進行微調。在這些模型之前會有一個 Router 模型，它能動態地將每個輸入 Prompt 分配給最相關的專家，比如數學模型將被路由到數學專家。

Samba-CoE 的靈感來自混合專家（Mixtral of Expert, MoE），但有一些關鍵區別。盡管 MoE 和 CoE 都比傳統的密集模型更稀疏，但 MoE 的靈活性不如 CoE。MoE 需要作為單一模型進行訓練/微調，類似于整體模型，而 CoE 由獨立和異構的專家模型組成，這些模型彼此獨立地進行訓練/微調。CoE 的能力也更強：先前的研究表明，CoE 的性能既優于 MoE，也優于 GPT-3.5 和 GPT-4 等大型單體模型。此外，CoE 和 MoE 也可以相互結合：CoE 可以在內部使用 MoE 實現的專家模型。

如下圖所示，Samba 最近升級了 Samba-CoE 模型，在 Router 中添加了不確定性量化來提升 Router 質量，如下圖所示為 Samba-CoE-v0.3（Samba-CoE v0.3: The Power of Routing ML Models at Scale） 的性能（看著 v0.3 主要提升來自 MMLU），可以看出其超越了很多常見的大模型：

三、Hybrid LLM

3.1 摘要

在 [2404.14618] Hybrid LLM: Cost-Efficient and Quality-Aware Query Routing 中，作者提出了一種利用低成本小模型和高質量大模型進行混合推理的方案。該方案使用一個 Router 模型，根據 Query 的難度和所需的質量水平將其分配給小模型或大模型。并且可以在使用時動態的調整質量水平，以便根據場景需求無縫的用質量換成本。評估結果表明，提出的方法可以在降低對大模型 40% 調用量的情況下不降低響應質量。

PS：當然，Hybrid LLM 中依然有一些局限性，比如沒有討論如果是多個模型時如何路由，如果數據分布改變或者大小模型改變后可能需要重新訓練 Router。

對應的代碼庫：GitHub - lm-sys/RouteLLM: A framework for serving and evaluating LLM routers - save LLM costs without compromising quality!

3.2 方法

3.2.1 概覽

如下圖 Figure 2 所示，其核心思路是通過 Router 來判斷 Query 難易程度，Easy Query 之間在 Edge 側使用小模型計算，Hard Query 在 Cloud 側使用大模型計算。因為小模型和大模型都是事先訓練好的，因此關鍵問題變成了如何訓練一個 Router，盡可能的把 Query 路由到小模型，同時又保證生成質量盡可能接近甚至超過大模型。

3.2.2 Router 訓練

作者采用 DeBERTa 模型作為 Router，使用 H(x) := q(S(x)) - q(L(x)) 作為兩個模型的質量差距，其中 S(x) 表示小模型結果，L(x) 表示大模型結果，q() 表示模型質量。

• 使用 Pr[H(x) >= 0] = Pr[q(S(x)) >= q(L(x))] 作為Router Score，如果對于 x 來說 Pr[H(x) >= 0] 的概率很高，則可以路由 x 到小模型。
• 通常來說小模型的質量可能遠小于大模型，此時可以設置一個閾值 t，Pr[H(x) >= -t] = Pr[q(S(x)) >= q(L(x)) - t] 時，表示如果小模型和大模型的差距不是特別大，也可以路由到小模型。t 值的大小可以作為生成質量以及成本之間的 tradeoff。

為了訓練 Router，作者設計了 3 種損失函數：

• 確定性 Router的損失函數：基于確定性假設，認為模型是確定性函數，將輸入特征映射到輸出空間的某個點。它通過為每個訓練 Query 從每個模型中采樣單個響應，并使用 BART Score 作為質量函數 q() 來分配二進制標簽。這種 Router 稱為rdet。
• 概率 Router的損失函數：考慮 NLP 任務的復雜性，LLM 生成結果通常有一定的隨機性（PS：GPT-4 模型即使設置 temperature 為 0 也無法保證結果完全相同）。該損失是將 hard 標簽轉換為 soft 標簽實現的，每個模型每個 Query 采樣 10 個響應并計算相應指示函數值的平均來估計。對應的 Router 稱為rprob。
• 帶有數據轉換的概率 Router損失函數：當小模型和大模型的差距比較大時，直接估計 Pr[H(x) >= 0] 可能得到的信號非常弱，導致訓練效果不佳。為了解決這個問題，作者定義了放松的條件 Pr[H(x) >= -t] 來修正損失函數。對應的 Router 稱為 rtrans。如下圖 Figure 4 為一個示例：?

3.3 實驗和評估

3.3.1 實驗配置

使用 MixInstruct 數據集評估 Router 的有效性，使用 DeBERTa-v3-large 作為 Router 的基座模型，使用 BART Score 作為質量指標，使用 BART Score Drop 表示質量差異，使用路由到小模型的 Query 比例作為效率指標（Cost Advantage）。選擇了 3 對大小模型，分別為 LLaMA2-7B 和 LLaMA2-13B，LLaMA2-13B 和 GPT-3.5 以及 FLAN-t5（800M） 和 LLaMA2-13B。

3.3.2 路由性能

如下圖 Table 1 所示，作者評估了在不同的 Cost Advantage 下幾種 Router 的效果，可以看出：

• LLaMA2-7B 和 LLaMA2-13B：模型差距很小，甚至在 40% 流量到達 LLaMA2-7B（cost advantage），依然可以保持基本無損：
• LLaMA2-13B 和 GPT-3.5：模型差距中等，20% 流量到達 LLaMA2-13B 就會有一些損失。
• FLAN-t5 和 LLaMA2-13B：模型差距比較大，10% 流量到達 FLAN-t5 即可能存在一些損失。當然，此時也可以看出 rtrans 的優勢。?

如下圖 Figure 6 所示，作者進一步進行了詳細的對比，并與隨機路由（Query 按照比例隨機分發到小模型和大模型）進行了對比。Avg Quality Gap Diff 表示的是，對于被路由到小模型的 Query 與被路由到大模型的 Query，它們之間平均質量差距的差異。正值意味著小模型的平均質量差距比大模型的小，這通常表示小模型的響應質量接近大型模型，從而實現了成本優勢。Random 時，路由到小模型和大模型的 Query 分布一致，平均質量差距類似，因此相應的 Avg Quality Gap Diff 接近于 0。

四、Router LLM

4.1 摘要

在 [2406.18665] RouteLLM: Learning to Route LLMs with Preference Data 中作者提出了 RouterLLM，其同樣是為了結合小模型低成本和大模型高質量來獲得生成質量和成本的平衡。為了解決這個問題，作者提出了幾種高效的 Router 模型，并開發了一個訓練框架，利用人類偏好數據和數據增強技術來提高 Router 性能。

在幾種常見的基準上評估表明，提出的方法可以在不影響響應質量的情況下顯著降低成本，某些情況可以降低 2x 成本；與此同時，提出的模型還顯示了顯著的遷移能力，即使在測試時更改大小模型也能保持性能。這凸顯了這些 Router 為部署 LLM 提供經濟高效且高性能的解決方案的潛力。

4.2 評估指標

RouterLLM 中作者主要提出了 3 種評估指標：PGR、APGR 和 CPT。

PGR（Performance Gain Ration）主要評估質量增益比例，該值越接近 1，表明獲得的質量越接近于大模型的性能。

CPT（Call Performance Threshold）是調用和質量閾值，用于衡量在給定的質量目標下，需要調用大模型的的最小百分比。比如說，GPT-4-1106-Preview 的 MT-Bench 指標為 9.3，Mixtral 8x7B 的 MT-Bench 指標為 8.3，則 CPT(50%)=37% 則表示要想 MT-Bench 維持在 (9.3-8.3)*50+8.3=8.8 的水平，至少需要有 37% 的請求調用 GPT-4-1106-Preview。

APGR（Average Performance Gain Recovered）：平均質量增益恢復，是一個綜合指標，用于衡量在不同成本約束下恢復質量的能力。

4.3 方法

4.3.1 Router 方案

作者提出了 4 種 Router 方案：

相似性加權排序（Similarity-weighted ranking, SW）：采用 Bradley-Terry（BT）模型，通過計算用戶 Query q 與訓練集中每個 Query qi 的相似性權重 wi 來預測大模型的獲勝概率。相似度 S 是通過 Query 對應 Embedding ε 的點積或歐式距離計算。其中的 BT 系數 ? 通過二元交叉熵損失函數來學習。

矩陣分解（Matrix factorization）：這種方法受到推薦系統中矩陣分解模型的啟發，目的是捕獲 user-item 的低秩結構。它通過一個隱藏的評分函數 s：M x Q -> R 來表示模型 Mw 對 Query q 的回答質量。評分函數 s 被建模為模型和 Query 的雙線性函數，并通過訓練優化 Bradley-Terry 關系來學習。

BERT 分類器（BERT classifier）：使用 BERT 模型訓練一個分類器，作者在 BERT 模型上進行了全參微調，在 2*L4 GPU 上訓練了 2000 個 step，batch size 為 16，最大序列長度為 512。

因果 LLM 分類器（Causal LLM classifier）：使用 LLaMA 3 8B 模型來作為 Router，采用指令跟隨范式，將用戶 Query 作為輸入指令，以下一個 Token 預測的方式輸出獲勝概率。與使用單獨的分類 head 不同，這里將比較標簽作為額外的 Token 添加到詞匯表中，并在標簽類別 L 上計算獲勝概率的 softmax。在 8xA100 80GB GPU 上訓練 2000 個 step，batch size 為 8，最大訓練長度為 2048。

4.3.2 數據

在 RouteLLM 中，作者使用了兩種類型的數據增強方面來增強 Router 模型的性能：

• Golden-Labeled Dataset：使用自動生成的準確標簽來增強訓練數據。一個具體的例子是 MMLU 基準測試，作者使用其驗證集的 1500 個問題及對應的標簽來生成 Dgold。
• LLM-Judge-labeled Dataset：使用一個 LLM 作為裁判來生成偏好標簽。具體來說，分別從大模型 GPT-4 和小模型 Mixtral-8x7B 來生成答案，然后讓 GPT-4 作為裁判來生成對應的比較標簽。作者通過這種方式收集了大約 120K 個樣本的偏好數據集 Djudge。

4.4 實驗和評估

4.4.1 結果

如下圖 Table 1 所示，作者在 MT-Bench 上評估了不同 Router 方法、訓練數據情況下的性能，其中的大小模型為 GPT-4 和 Mixtral-8x7B。可以看出，Matrix factorization 和 Similarity-weighted ranking 獲得更好的結果，其中 25.32%，表示保證 CPT(50%) 也就是 MT-bench 為 8.8 的情況下最少只需 25.32% 的流量路由到大模型 GPT-4，越低越好。APGR 表示質量恢復到 GPT-4 的水平，越高越好。

如下圖 Table 3 所示，作者進一步在 8-shot GSM8K 上進行評估，結論稍有不同，不使用 Djudge 數據集結果和隨機路由差不多，使用后有所改善：

4.4.2 遷移能力

如下圖 Table 5 所示，直接將訓練好的 Router 模型應用到 Claude 3 Opus 和 LLaMA3-8B 上，同樣可以獲得一定的收益，表明其具有不錯的遷移能力：

4.4.3 成本分析

如下圖 Table 6 所示，作者進一步評估了在不同質量要求下成本節約情況，在 MT-Bench 上，CPT(80%) 依然可以節約 2.49x 成本：

五、參考鏈接

1. ??https://arxiv.org/abs/2106.11520??
2. ??https://arxiv.org/abs/1910.13461??
3. ??https://arxiv.org/abs/2405.07518??
4. ??https://sambanova.ai/blog/samba-coe-the-power-of-routing-ml-models-at-scale??
5. ??https://arxiv.org/abs/2404.14618??
6. ??https://arxiv.org/abs/2406.18665??

本文轉載自 ??AI閑談??，作者： AI閑談