在人工智能快速發(fā)展的今天，大型語言模型（LLMs）展現(xiàn)出令人印象深刻的推理能力，但這些能力往往伴隨著巨大的推理開銷，給部署帶來重大挑戰(zhàn)。近日，來自清華大學、無限極AI和上海交通大學的研究團隊發(fā)表了一篇題為《R2R: Efficiently Navigating Divergent Reasoning Paths with Small-Large Model Token Routing》的研究論文，提出了一種名為"通向羅馬之路"（Roads to Rome，簡稱R2R）的創(chuàng)新方法，通過在小型和大型語言模型之間進行智能令牌路由，顯著提升了推理效率。該研究已提交至審稿階段，代碼已在GitHub上開源（https://github.com/thu-nics/R2R）。

雖然經(jīng)過蒸餾的小型語言模型（SLMs）可以大幅提高生成效率，但它們在嘗試跟隨大型語言模型的推理路徑時常常會出現(xiàn)性能下降。研究團隊發(fā)現(xiàn)了一個有趣的現(xiàn)象：只有一小部分令牌真正導(dǎo)致LLMs和SLMs之間的推理路徑發(fā)生分歧。大多數(shù)生成的令牌要么完全相同，要么只有中性差異，如縮寫或表達方式的細微變化。

想象一下，當你要解一道復(fù)雜的數(shù)學題時，你可以請一位經(jīng)驗豐富的數(shù)學教授（大模型）或一位熱心的高中生（小模型）來幫忙。大部分時候，高中生的解題思路與教授相似，只是偶爾在某些關(guān)鍵步驟上會走偏。如果我們能在高中生即將犯錯的關(guān)鍵時刻讓教授臨時指導(dǎo)一下，然后再讓高中生繼續(xù)，這不就既保證了解題質(zhì)量，又節(jié)省了教授的時間嗎？

這正是R2R方法的核心思想。研究團隊基于這一發(fā)現(xiàn)，設(shè)計了一種神經(jīng)令牌路由方法，只在那些關(guān)鍵的、會導(dǎo)致路徑分歧的令牌上選擇性地使用大型模型，而將大部分令牌生成任務(wù)交給小型模型。此外，團隊還開發(fā)了一個自動數(shù)據(jù)生成管道，用于識別分歧令牌并生成令牌級路由標簽，以訓(xùn)練輕量級路由器。

讓我們深入了解這項研究的細節(jié)和成果。

一、研究背景與挑戰(zhàn)

大型語言模型展示出強大的能力，但它們的推理過程需要大量計算資源。通常，這些模型會生成詳細的思維鏈（Chain-of-Thought，CoT）推理，然后才給出最終答案。然而，這種方法需要擁有數(shù)百億參數(shù)的大型模型為每個查詢生成數(shù)千個令牌，導(dǎo)致推理開銷巨大。

經(jīng)過蒸餾的小型語言模型（SLMs）只有幾十億參數(shù)，生成效率高得多。通過對大模型響應(yīng)進行監(jiān)督微調(diào)，這些小模型可以模仿大模型的推理行為。但問題是，SLMs在推理過程中仍可能產(chǎn)生與大模型不同的推理路徑，導(dǎo)致性能嚴重下降。

研究團隊舉例說明：與32B參數(shù)的R1-32B大模型相比，1.5B參數(shù)的R1-1.5B小模型在AIME基準測試中45%的問題上給出了不同的最終答案，準確率下降了4.8倍。

然而，研究團隊發(fā)現(xiàn)，SLMs和LLMs在給定相同上下文的情況下，往往在下一個令牌預(yù)測上達成一致。它們之間的性能差距主要來自于累積錯誤：在部分響應(yīng)中出現(xiàn)了一些關(guān)鍵差異后，它們的推理路徑逐漸偏離。

通過詳細分析，團隊發(fā)現(xiàn)在32B模型生成的總計760萬個令牌中，1.5B模型僅在11%的令牌上產(chǎn)生了不同的預(yù)測結(jié)果。更重要的是，這些差異中的大部分只是中性變化，例如縮寫或表達方式的不同（如"let's"與"let us"），并不影響推理結(jié)果。真正導(dǎo)致推理路徑分歧的是一小部分特殊令牌，團隊將其稱為"分歧令牌"（divergent tokens）。這些令牌真正改變了當前句子的含義、邏輯或結(jié)論，從而導(dǎo)致后續(xù)推理路徑的偏離。

這一發(fā)現(xiàn)引發(fā)了一個關(guān)鍵問題：SLMs是否可以通過僅替換分歧令牌來跟隨LLM的推理路徑？如果能解決這個問題，我們就能在保持LLM高質(zhì)量推理的同時，充分利用SLM的效率優(yōu)勢。

二、R2R方法：通向羅馬之路

解決SLM-LLM混合推理的主要挑戰(zhàn)包括兩個相互關(guān)聯(lián)的部分：在特定目標下標記首選模型，以及設(shè)計路由方案在推理過程中執(zhí)行這種選擇。

以往的方法通常在查詢級別進行路由，為整個響應(yīng)選擇SLM或LLM，以在成本預(yù)算內(nèi)最大化人類偏好的勝率。然而，這些方法依賴于人工標注和復(fù)雜的路由器設(shè)計，其數(shù)據(jù)標注和路由方案對于細粒度的令牌級路由來說都太昂貴了。

另一方面，推測性解碼方法（speculative decoding）旨在令牌級別上實現(xiàn)SLM和LLM之間的完全相同輸出。它們使用SLM草擬輸出，并定期用LLM驗證。雖然這種方法精確，但嚴格的驗證導(dǎo)致接受率低。此外，中間的差異會使所有后續(xù)令牌無效，嚴重限制了接受長度。

為了解決這些挑戰(zhàn)，研究團隊提出了"通向羅馬之路"（R2R）方法，這是一種令牌級路由方法，僅在SLM生成過程中針對路徑分歧令牌選擇性地使用LLM。

### 1. 令牌級路由公式化

在自回歸語言模型中，推理可以表示為一系列下一個令牌預(yù)測的序列。研究團隊專注于貪婪采樣，通過以下方式定義：

yi = argmax_y P_mi(y|x0, ..., xn-1, y0, ..., yi-1) = argmax_y P_mi(y|S<i)

其中，xi和yi分別表示輸入和輸出令牌，S<i表示在步驟i的令牌序列。下一個令牌概率P_mi由步驟i處的模型mi預(yù)測，mi可以是SLM（θs）或LLM（θl）。

路由策略的本質(zhì)是定義一個路由函數(shù)R，為每個解碼步驟選擇模型：

mi = R(S<i, θs, θl)

團隊的目標是在確保輸出序列與僅使用LLM的輸出質(zhì)量相匹配的同時，最小化總生成成本。

### 2. 路徑跟隨路由策略

最優(yōu)地解決令牌級路由問題在計算上是難以實現(xiàn)的，特別是對于大規(guī)模數(shù)據(jù)生成。雖然可能存在更好的路由序列（可能與LLM的推理路徑不同），但找到它們需要徹底搜索O(2^n)的空間，并為每次搜索生成數(shù)千個輸出令牌。

為了克服這一實際限制，研究團隊提出了一種貪婪的、句子級路徑跟隨路由策略，將搜索復(fù)雜度降低到O(n)。該策略不是探索所有可能的模型選擇，而是逐步將混合模型生成與LLM建立的推理路徑對齊。

具體來說，在每個生成步驟，策略都會優(yōu)先選擇高效的SLM，除非這會導(dǎo)致與LLM預(yù)期推理路徑的有意義偏離。為了確定差異是中性的還是分歧的，團隊使用了一種延續(xù)驗證機制：

1. 首先比較SLM和LLM的下一個令牌預(yù)測 2. 如果預(yù)測相同，則選擇SLM 3. 如果預(yù)測不同，則構(gòu)建兩個候選序列，分別從SLM和LLM的預(yù)測開始 4. 使用LLM繼續(xù)生成這兩個序列，直到滿足停止標準（如生成EOS令牌） 5. 如果SLM的延續(xù)序列在驗證函數(shù)V下與LLM的延續(xù)序列質(zhì)量相當，則差異被視為中性的，否則被視為分歧的 6. 根據(jù)這一判斷選擇對應(yīng)的模型

雖然這種策略可以保證混合生成的令牌序列與其僅LLM的對應(yīng)物達到相同的質(zhì)量，但全面延續(xù)在計算上對大規(guī)模數(shù)據(jù)生成來說太昂貴了。此外，當前差異對數(shù)千個令牌之后的最終輸出質(zhì)量的影響對于要訓(xùn)練的神經(jīng)路由器來說太難學習了。

因此，在實踐中，團隊使用句子級路徑跟隨路由，其中延續(xù)在當前句子結(jié)束時結(jié)束。他們在延續(xù)過程中監(jiān)控句子結(jié)束符號（如句號），并使用現(xiàn)有的語義句子分隔器在句子真正結(jié)束時結(jié)束生成。為了驗證這種局部延續(xù)，一個強大的LLM作為句子級驗證器，用于比較延續(xù)并確定初始令牌差異是否引入了與LLM預(yù)期推理路徑的有意義偏離。

團隊驗證了句子級路徑跟隨路由的有效性，使用Qwen2.5-72B作為驗證模型。在R1-32B在8K令牌限制內(nèi)正確解決的17個AIME-24問題中，他們的路徑跟隨策略實現(xiàn)了可比的準確率（16個問題正確回答），同時僅在3%的生成令牌上依賴較小的R1-1.5B模型。

然而，這種策略直接用于實時推理是不切實際的，因為它依賴于昂貴的LLM延續(xù)和驗證。相反，團隊設(shè)計并訓(xùn)練了一個輕量級神經(jīng)路由器，它可以有效地近似這種策略，僅依靠SLM輸出來確定在推理過程中何時使用LLM。

三、令牌級神經(jīng)路由器

### 1. 訓(xùn)練數(shù)據(jù)生成

研究團隊使用句子級路徑跟隨路由為神經(jīng)路由器生成訓(xùn)練標簽，并采用了幾種優(yōu)化來控制數(shù)據(jù)標注開銷。

數(shù)據(jù)生成管道包括以下步驟：

1. 從現(xiàn)有數(shù)據(jù)集獲取查詢，獲取完整的LLM響應(yīng)（直接從數(shù)據(jù)集獲取或通過批處理LLM推理） 2. 使用高度并行的SLM預(yù)填充，高效地識別SLM預(yù)測與LLM相同的令牌，允許排除約90%無需進一步處理的令牌 3. 對剩余10%的不同令牌，從每個SLM預(yù)測執(zhí)行批處理LLM延續(xù) 4. 為提高效率，應(yīng)用前綴緩存重用多個延續(xù)之間共享上下文前綴的KV-Cache計算 5. 驗證模型比較延續(xù)并標記路由偏好

使用這個管道，團隊在約2.3天內(nèi)使用8個A800 GPU高效生成了760萬個路由標簽，涵蓋了Bespoke-Stratos數(shù)據(jù)集中的數(shù)學、編碼和問答主題。

### 2. 分歧的預(yù)測指標

團隊探索了可以幫助識別分歧令牌的預(yù)測指標，專注于僅在SLM下一個令牌預(yù)測期間可以獲取的指標。

SLM的輸出logits：數(shù)據(jù)顯示，分歧令牌在SLM的輸出logits中表現(xiàn)出明顯更高的熵，平均值是非分歧令牌的3.8倍。團隊觀察到使用其他不確定性度量（如隨機不確定性）也有類似趨勢。這些實證結(jié)果表明，SLM預(yù)測中的不確定性增加與令牌分歧強烈相關(guān)。

令牌頻率：數(shù)據(jù)集中的低頻令牌更可能是分歧的。這可能是由于訓(xùn)練數(shù)據(jù)中的長尾令牌分布，使得罕見令牌由于SLM的有限容量而更難有效建模。

基于這些見解，路由器將SLM的top-100 logit值作為輸入特征之一，并通過使用令牌嵌入作為路由器輸入來明確合并令牌頻率偏差。

### 3. 路由器設(shè)計和路由方案

模型架構(gòu)：團隊將神經(jīng)路由器設(shè)計為一個輕量級的五層前饋網(wǎng)絡(luò)（FFN），參數(shù)為56M。它接收SLM的輸出logits和標記化嵌入，以及最后一層隱藏狀態(tài)以獲取額外的語義上下文。所有輸入都經(jīng)過線性投影，連接起來，并送入FFN主干。路由器輸出一個二元分類概率，指示當前令牌是否偏離LLM的推理路徑。

訓(xùn)練方案：使用交叉熵損失和第4.1節(jié)描述的標記數(shù)據(jù)訓(xùn)練路由器。為了解決由低分歧率引起的類別不平衡，團隊對損失進行了反比于類別頻率的重新加權(quán)。訓(xùn)練后，使用驗證集選擇滿足用戶定義的LLM使用率的路由閾值。

路由方案：與定期驗證SLM輸出的推測性解碼方法不同，研究團隊的路由方案旨在立即決定是否接受每個SLM令牌，從而消除回滾的需要。如圖2所示，這種方法減少了不必要的草稿和驗證計算，這在計算密集型批處理服務(wù)場景中特別有益。具體來說，神經(jīng)路由器在每個生成步驟使用SLM輸出估計分歧概率。當這個概率超過預(yù)定義的閾值pth時，就會調(diào)用LLM來糾正當前的輸出令牌。

四、實驗結(jié)果與分析

### 1. 實驗設(shè)置

基線：研究團隊使用DeepSeek-R1-Distill-Qwen模型作為基線，用R1-MB表示，其中M表示模型大小（單位：十億）。他們將R1-1.5B和R1-32B分別指定為SLM和LLM，中間大?。?B，14B）捕獲蒸餾縮放行為。團隊比較了RouteLLM框架中的各種查詢級路由（QR）方法，包括相似度加權(quán)排名（QR-SW）、矩陣分解（QR-MF）、基于BERT的分類器（QR-BERT）和基于Llama3-8B的分類器（QR-LLM）。對于推測性解碼，他們采用了EAGLE2和HASS，LLM為R1-32B。

R2R設(shè)置：R2R在R1-1.5B和R1-32B之間路由，使用輕量級56M參數(shù)FFN路由器，該路由器在第4.1節(jié)描述的760萬個令牌級路由標簽上訓(xùn)練。路由閾值pth在驗證集上選擇，平均參數(shù)使用為6B。性能效率權(quán)衡僅通過調(diào)整pth控制，無需重新訓(xùn)練路由器。

基準測試：評估方法涵蓋具有挑戰(zhàn)性的推理基準，包括數(shù)學（AIME 2024-2025），研究生級問答（GPQA-Diamond）和編碼任務(wù)（LiveCodeBench 2024-08-2025-01）。所有實驗使用32K令牌的最大輸出長度和零生成溫度以確?？芍貜?fù)性。

效率指標：研究團隊遵循之前的工作，使用每令牌平均激活參數(shù)作為與硬件無關(guān)的效率指標，簡稱為平均參數(shù)（M）。對于查詢級路由，M根據(jù)SLM和LLM在所有輸出中的激活比率計算其參數(shù)的加權(quán)平均值。對于R2R，M包括SLM和路由器參數(shù)，以及根據(jù)LLM激活比率加權(quán)的LLM參數(shù)。團隊還報告了總成本（C），定義為平均激活參數(shù)乘以每個查詢的平均輸出令牌數(shù)。平均參數(shù)和總成本分別反映了平均解碼速度和總延遲。此外，團隊還使用SGLang框架在NVIDIA A800-80GB GPU上報告了特定硬件的解碼速度。

### 2. 性能結(jié)果

縮放行為：圖5顯示了平均激活參數(shù)與準確率的縮放關(guān)系。查詢級路由（QR）方法從1.5B到32B參數(shù)顯示出近線性的準確率縮放。經(jīng)過廣泛訓(xùn)練的蒸餾模型（R1-7B，R1-14B）實現(xiàn)了超線性增益，在14B時達到R1-32B準確率的88%，參數(shù)大小僅為50%。通過僅將分歧令牌路由到LLM，R2R實現(xiàn)了92%的平均準確率，平均參數(shù)僅為17%，達到了更好的縮放效果，形成了新的Pareto前沿。此外，由于輸出長度減少，R2R在準確率與總測試時間成本C的權(quán)衡方面提供了更好的效果。R2R中的路由閾值還能夠在訓(xùn)練后靈活控制這種權(quán)衡。

數(shù)值比較：表2顯示了7B和14B左右平均參數(shù)大小的模型性能的數(shù)值細節(jié)。平均參數(shù)大小為5.6B的R2R分別比最佳查詢級路由方法（在7B和14B中）提高了1.4-2.4倍和1.2-1.4倍。與蒸餾模型相比，R2R比R1-7B提高了1.4-1.7倍的準確率，甚至在平均準確率上超過了R1-14B 1.1倍。相對于極端情況，R2R實現(xiàn)了比R1-1.5B高4.6倍的準確率，同時保持了R1-32B 92%的準確率，而LLM的使用僅為11-15%。

### 3. 實際速度

表3報告了AIME基準上所有方法的實際延遲和速度。所有基線都使用官方的、高效的SGLang框架，并在兩個NVIDIA A800-80GB GPU上使用張量并行性進行評估。R2R使用與表2相同的閾值；查詢級路由方法使用14B版本以獲得可比的性能。R2R分別比R1-14B和R1-32B實現(xiàn)了1.62倍和2.76倍的生成速度。與查詢級路由相比，R2R提供了1.48-1.52倍的加速。它還優(yōu)于具有樹狀草稿的高度優(yōu)化推測性解碼方法，這些方法主要在當前的單批處理設(shè)置中加速。進一步的系統(tǒng)級優(yōu)化可以為R2R帶來更大的收益。

### 4. 消融研究

從表4第一行的R2R開始，團隊通過重新訓(xùn)練具有替代目標或減少輸入的路由器來評估設(shè)計的有效性，保持所有其他設(shè)置不變。所有實驗都在AIME基準上進行。

路由目標：如第3節(jié)所述，團隊將不同的下一個令牌預(yù)測分類為中性或分歧。R2R通過容忍中性差異并僅將真正分歧的令牌路由到LLM來提高效率。當路由器被訓(xùn)練為將所有不同的令牌都使用LLM時，它無法在相同數(shù)量的LLM使用下達到原始準確率，面臨1.4倍的準確率下降，如表4第二行所示。這證實了將LLM使用限制在僅分歧令牌對于在保持高準確率的同時降低成本至關(guān)重要。

路由器輸入：如第4節(jié)所述，SLM logits和令牌嵌入都是分歧的強指標，用作路由器輸入。當這些特征逐漸被移除時，路由準確率下降了高達1.3倍，凸顯了它們的重要性。值得注意的是，雖然SLM logits理論上可以從路由器內(nèi)的最后一層隱藏狀態(tài)計算，但這需要234M參數(shù)嵌入層的容量，超過了56M參數(shù)神經(jīng)路由器的容量。

### 5. 路由結(jié)果觀察

研究團隊分析了R2R在AIME基準上的路由行為，考慮了在32K令牌限制內(nèi)完成的響應(yīng)。圖6(a)顯示了跨響應(yīng)位置的LLM使用率。每個響應(yīng)被分為思考過程和后續(xù)回復(fù)，位置歸一化為[0, 1]。子圖寬度反映了它們各自的平均長度。結(jié)果顯示，R2R在回復(fù)階段路由到LLM的令牌明顯減少。這反映了一種直覺，即在內(nèi)部思考之后，回復(fù)本身是直接的，要求較低。

遵循先前的工作，團隊進一步將思考過程基于諸如"Wait"和"Alternatively"等令牌分割為順序思考。圖6(b)檢查了每個思考內(nèi)的LLM使用率。它顯示R2R在每個思考的開始和結(jié)束時更多地依賴LLM。這與預(yù)期一致，即初始令牌為思考設(shè)定方向，而結(jié)束令牌決定是否結(jié)束思考、分支到替代方案或繼續(xù)更深入的推理。值得注意的是，這些路由模式不是手工制作的，而是自然地從路由器訓(xùn)練中出現(xiàn)的。它幫助R2R有效地分配LLM以獲得更好的測試時間縮放。

五、結(jié)論與展望

通過這項研究，研究團隊提出了R2R，一種令牌級路由方法，通過糾正僅分歧令牌使SLM能夠跟隨LLM推理路徑。他們設(shè)計了一種路徑跟隨路由策略進行有效的數(shù)據(jù)標注，并研究了預(yù)測指標，使神經(jīng)路由器能夠準確選擇令牌。在具有挑戰(zhàn)性的基準測試中的實驗表明，R2R以不到7B的平均參數(shù)超越了R1-14B模型，將SLM性能提高了4.6倍，LLM使用率不到15%，并在可比準確率下實現(xiàn)了比LLM快2.8倍的實際速度。

然而，當前的路由策略假設(shè)貪婪采樣。探索替代采樣策略可能進一步增強多功能性。此外，需要更專門的系統(tǒng)級優(yōu)化來充分實現(xiàn)R2R的理論成本優(yōu)勢。

總的來說，R2R展示了一種有前景的方法，通過智能令牌路由有效地導(dǎo)航分歧推理路徑，為大型語言模型的高效部署鋪平了道路。