一、背景

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

目前常見(jiàn)有兩種路由的范式：

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

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

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

二、引言

2.1 投機(jī)采樣

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

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

2.2 BART Score

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

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

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

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

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

二、SambaNova CoE

如下圖 Fig 2 所示（來(lái)自 [2405.07518] SambaNova SN40L: Scaling the AI Memory Wall with Dataflow and Composition of Experts），SambaNova 很早之前就提出了 CoE（Composition of Experts） 的方案，其 Samba-CoE 可能包含數(shù)十個(gè)專家模型，這些專家模型都來(lái)自開(kāi)源社區(qū)，作者在編碼、數(shù)學(xué)和翻譯等特定領(lǐng)域進(jìn)行微調(diào)。在這些模型之前會(huì)有一個(gè) Router 模型，它能動(dòng)態(tài)地將每個(gè)輸入 Prompt 分配給最相關(guān)的專家，比如數(shù)學(xué)模型將被路由到數(shù)學(xué)專家。

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

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

三、Hybrid LLM

3.1 摘要

在 [2404.14618] Hybrid LLM: Cost-Efficient and Quality-Aware Query Routing 中，作者提出了一種利用低成本小模型和高質(zhì)量大模型進(jìn)行混合推理的方案。該方案使用一個(gè) Router 模型，根據(jù) Query 的難度和所需的質(zhì)量水平將其分配給小模型或大模型。并且可以在使用時(shí)動(dòng)態(tài)的調(diào)整質(zhì)量水平，以便根據(jù)場(chǎng)景需求無(wú)縫的用質(zhì)量換成本。評(píng)估結(jié)果表明，提出的方法可以在降低對(duì)大模型 40% 調(diào)用量的情況下不降低響應(yīng)質(zhì)量。

PS：當(dāng)然，Hybrid LLM 中依然有一些局限性，比如沒(méi)有討論如果是多個(gè)模型時(shí)如何路由，如果數(shù)據(jù)分布改變或者大小模型改變后可能需要重新訓(xùn)練 Router。

對(duì)應(yīng)的代碼庫(kù)：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 所示，其核心思路是通過(guò) Router 來(lái)判斷 Query 難易程度，Easy Query 之間在 Edge 側(cè)使用小模型計(jì)算，Hard Query 在 Cloud 側(cè)使用大模型計(jì)算。因?yàn)樾∧Ｐ秃痛竽Ｐ投际鞘孪扔?xùn)練好的，因此關(guān)鍵問(wèn)題變成了如何訓(xùn)練一個(gè) Router，盡可能的把 Query 路由到小模型，同時(shí)又保證生成質(zhì)量盡可能接近甚至超過(guò)大模型。

3.2.2 Router 訓(xùn)練

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

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

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

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

3.3 實(shí)驗(yàn)和評(píng)估

3.3.1 實(shí)驗(yàn)配置

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

3.3.2 路由性能

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

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

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

四、Router LLM

4.1 摘要

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

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

4.2 評(píng)估指標(biāo)

RouterLLM 中作者主要提出了 3 種評(píng)估指標(biāo)：PGR、APGR 和 CPT。

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

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

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

4.3 方法

4.3.1 Router 方案

作者提出了 4 種 Router 方案：

相似性加權(quán)排序（Similarity-weighted ranking, SW）：采用 Bradley-Terry（BT）模型，通過(guò)計(jì)算用戶 Query q 與訓(xùn)練集中每個(gè) Query qi 的相似性權(quán)重 wi 來(lái)預(yù)測(cè)大模型的獲勝概率。相似度 S 是通過(guò) Query 對(duì)應(yīng) Embedding ε 的點(diǎn)積或歐式距離計(jì)算。其中的 BT 系數(shù) ? 通過(guò)二元交叉熵?fù)p失函數(shù)來(lái)學(xué)習(xí)。

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

BERT 分類器（BERT classifier）：使用 BERT 模型訓(xùn)練一個(gè)分類器，作者在 BERT 模型上進(jìn)行了全參微調(diào)，在 2*L4 GPU 上訓(xùn)練了 2000 個(gè) step，batch size 為 16，最大序列長(zhǎng)度為 512。

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

4.3.2 數(shù)據(jù)

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

• Golden-Labeled Dataset：使用自動(dòng)生成的準(zhǔn)確標(biāo)簽來(lái)增強(qiáng)訓(xùn)練數(shù)據(jù)。一個(gè)具體的例子是 MMLU 基準(zhǔn)測(cè)試，作者使用其驗(yàn)證集的 1500 個(gè)問(wèn)題及對(duì)應(yīng)的標(biāo)簽來(lái)生成 Dgold。
• LLM-Judge-labeled Dataset：使用一個(gè) LLM 作為裁判來(lái)生成偏好標(biāo)簽。具體來(lái)說(shuō)，分別從大模型 GPT-4 和小模型 Mixtral-8x7B 來(lái)生成答案，然后讓 GPT-4 作為裁判來(lái)生成對(duì)應(yīng)的比較標(biāo)簽。作者通過(guò)這種方式收集了大約 120K 個(gè)樣本的偏好數(shù)據(jù)集 Djudge。

4.4 實(shí)驗(yàn)和評(píng)估

4.4.1 結(jié)果

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

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

4.4.2 遷移能力

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

4.4.3 成本分析

如下圖 Table 6 所示，作者進(jìn)一步評(píng)估了在不同質(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??

本文轉(zhuǎn)載自 ??AI閑談??，作者： AI閑談