一、背景

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

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

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

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

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

二、引言

2.1 投機采樣

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

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

2.2 BART Score

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

BART Score 的優(yōu)勢在于它利用了預訓練的 BART 模型所捕捉到的豐富語言信息和上下文關系，能夠更好地反映生成文本的質(zhì)量，特別是在考慮句子間的語義相似性和一致性方面。相比傳統(tǒng)的 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 可能包含數(shù)十個專家模型，這些專家模型都來自開源社區(qū)，作者在編碼、數(shù)學和翻譯等特定領域進行微調(diào)。在這些模型之前會有一個 Router 模型，它能動態(tài)地將每個輸入 Prompt 分配給最相關的專家，比如數(shù)學模型將被路由到數(shù)學專家。

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

如下圖所示，Samba 最近升級了 Samba-CoE 模型，在 Router 中添加了不確定性量化來提升 Router 質(zhì)量，如下圖所示為 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 中，作者提出了一種利用低成本小模型和高質(zhì)量大模型進行混合推理的方案。該方案使用一個 Router 模型，根據(jù) Query 的難度和所需的質(zhì)量水平將其分配給小模型或大模型。并且可以在使用時動態(tài)的調(diào)整質(zhì)量水平，以便根據(jù)場景需求無縫的用質(zhì)量換成本。評估結果表明，提出的方法可以在降低對大模型 40% 調(diào)用量的情況下不降低響應質(zhì)量。

PS：當然，Hybrid LLM 中依然有一些局限性，比如沒有討論如果是多個模型時如何路由，如果數(shù)據(jù)分布改變或者大小模型改變后可能需要重新訓練 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 路由到小模型，同時又保證生成質(zhì)量盡可能接近甚至超過大模型。

3.2.2 Router 訓練

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

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

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

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

3.3 實驗和評估

3.3.1 實驗配置

使用 MixInstruct 數(shù)據(jù)集評估 Router 的有效性，使用 DeBERTa-v3-large 作為 Router 的基座模型，使用 BART Score 作為質(zhì)量指標，使用 BART Score Drop 表示質(zhì)量差異，使用路由到小模型的 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 的優(yōu)勢。?

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

四、Router LLM

4.1 摘要

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

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

4.2 評估指標

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

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

CPT（Call Performance Threshold）是調(diào)用和質(zhì)量閾值，用于衡量在給定的質(zhì)量目標下，需要調(diào)用大模型的的最小百分比。比如說，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% 的請求調(diào)用 GPT-4-1106-Preview。

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

4.3 方法

4.3.1 Router 方案

作者提出了 4 種 Router 方案：

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

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

BERT 分類器（BERT classifier）：使用 BERT 模型訓練一個分類器，作者在 BERT 模型上進行了全參微調(diào)，在 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 數(shù)據(jù)

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

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

4.4 實驗和評估

4.4.1 結果

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

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

4.4.2 遷移能力

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

4.4.3 成本分析

如下圖 Table 6 所示，作者進一步評估了在不同質(zhì)量要求下成本節(jié)約情況，在 MT-Bench 上，CPT(80%) 依然可以節(jié)約 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閑談