選擇性知識蒸餾精準過濾:推測解碼加速器AdaSPEC來了

選擇性知識蒸餾精準過濾:推測解碼加速器AdaSPEC來了

文章圖片

選擇性知識蒸餾精準過濾:推測解碼加速器AdaSPEC來了

文章圖片




本文共同第一作者為加州大學伯克利分校的博士生胡越舟與清華大學的本科生郭佳鑫 , 通訊作者為佐治亞理工學院的副教授趙拓 。

推測解碼(Speculative Decoding SD)通過使用一個較小的草稿模型(draft model)生成候選預測 , 再由更大的目標模型(target model)進行驗證 , 從而顯著加速大語言模型(LLM)的推理過程 。 SD 的加速效果在很大程度上取決于兩者之間的對齊程度 。

目前 , 最先進的對齊方法是使用知識蒸餾(Knowledge Distillation KD)在所有 token 上最小化 KL 散度 。 然而 , 最小化全局 KL 散度并不意味著 token 的接受率最大化 。 由于小模型容量受限 , 草稿模型往往難以完整吸收目標模型的知識 , 導致直接使用蒸餾方法的性能提升受限 。 在極限場景下 , 草稿模型和目標模型的巨大尺寸差異甚至可能導致訓練不收斂 。

為了解決這一問題 , 佐治亞理工、清華大學與加州大學伯克利分校的研究團隊提出 AdaSPEC , 一種引入選擇性 token 過濾機制的創新蒸餾方法 。 AdaSPEC 利用參考模型(reference model)識別并過濾出難以學習的 token , 使蒸餾過程更聚焦于「易學習」的部分 , 從而讓草稿模型在有限容量下更好地對齊目標模型 。

這種選擇性蒸餾策略在不降低生成質量的前提下 , 顯著提升了整體 token 接受率 。 我們在多個任務(算術推理、指令跟隨、代碼生成與文本摘要)和不同規模模型組合(31M/1.4B、350M/2.7B)上進行了系統評估 。 結果表明 , AdaSPEC 在所有任務上均超越當前最優的 DistillSpec 方法 , token 接受率最高提升達 15% 。


論文標題:AdaSPEC: Selective Knowledge Distillation for Efficient Speculative Decoders 論文鏈接:https://arxiv.org/abs/2510.19779 Github 鏈接:https://github.com/yuezhouhu/adaspec
研究背景

大型語言模型(LLM)在推理和生成任務中表現卓越 , 但其自回歸解碼機制導致推理延遲高、計算開銷大 , 成為實際部署的主要瓶頸 。 傳統加速方法如模型壓縮、量化或知識蒸餾雖能提升速度 , 但往往以犧牲生成質量為代價 。

近年來 , 推測解碼(Speculative Decoding SD)提供了一條新路徑:通過一個輕量級「草稿模型」并行生成多個候選 token , 再由原始「主模型」批量驗證 , 從而減少主模型的前向調用次數 。 理論上 , SD 可在不損失生成質量的前提下顯著加速推理 。 然而 , 其實際效果高度依賴草稿模型與主模型的預測一致性 —— 若草稿模型生成的候選 token 頻繁被主模型拒絕 , 加速收益將大打折扣 。
【選擇性知識蒸餾精準過濾:推測解碼加速器AdaSPEC來了】
當前一種做法是使用知識蒸餾(KD)讓草稿模型模仿主模型的輸出分布 。 但草稿模型通常比主模型小一個數量級 , 容量有限 , 強行擬合所有 token(尤其是罕見或上下文敏感的「難學 token」)不僅效率低下 , 還可能擠占其學習「易學 token」 的能力 , 反而降低整體接受率 。

針對這一問題 , 研究團隊提出 AdaSPEC—— 一種面向推測解碼的選擇性知識蒸餾方法 。 AdaSPEC 的核心思想是:讓草稿模型專注于學習那些它真正能掌握且對接受率貢獻大的「易學 token」 , 主動忽略難以擬合的 token 。 通過兩階段訓練(先識別難 token , 再在蒸餾中過濾) , AdaSPEC 更高效地利用草稿模型的有限容量 , 顯著提升其與主模型的一致性 。

實驗表明 , AdaSPEC 在多種模型和任務上 consistently 提高 token 接受率(最高提升 15%) , 在保持生成質量的同時 , 有效釋放了推測解碼的加速潛力 。

方法概述

AdaSPEC 的核心思想是:在蒸餾階段(如下圖 1)識別并過濾難以學習的 token , 讓知識遷移更聚焦、更有效 。


1、Selective KD 核心機制

通過引入參考模型(reference model) , 自動篩除訓練樣本中「難以對齊」的 token , 僅在「易學習」子集上進行蒸餾 , 從根本上緩解 draft–target 不匹配問題 。

2、雙階段訓練框架

AdaSPEC 首先在參考模型上執行初步蒸餾 , 得到參考模型 。 隨后使用參考模型過濾微調數據集 , 并在過濾后的子集上優化草稿模型 。 該方法顯著減少無效學習與梯度噪聲 , 既保持生成質量 , 又有效提升 token 接受率 。

3、通用適配性與輕量實現

AdaSPEC 具備極高的模塊化兼容性與結構清晰的設計 , 可無縫結合 EAGLE、vLLM 等高級推測解碼框架 。 核心實現不到百行 , 思路直觀、代碼簡潔 , 能直接適配任意 Transformer 架構的草稿–目標模型組合 , 便于研究者與開發者快速上手 。

實驗評估

研究團隊在多個模型家族(Pythia、CodeGen、Phi-2 等)和多種任務(GSM8K、Alpaca、MBPP、CNN/DailyMail、XSUM)上進行了系統驗證 , 在不同模型規模與任務類型下均展現出一致且穩健的提升效果 , 體現了方法的魯棒性與廣泛適用性 。 主要實驗結果表明:

Token 接受率全線超越基線方法 DistillSpec:在 GSM8K 上提升 5–6% , 在 MBPP 上最高提升 15% 。 實際加速顯著:經微調后使用 vLLM 框架測速 , end2end 推理速度提升可達 10–20% 。 進一步兼容提升:結合 EAGLE 框架微調 , 生成速度再提高 7.5% 。

總結與展望

AdaSPEC 為推測解碼提供了一種精準、高效、通用且具有廣泛適用性的加速新范式 。 它通過「選擇性蒸餾 + 自適應過濾」實現了 draft–target 的動態對齊 , 為未來 LLM 高效推理的研究與工業部署開辟了新方向 。

當前工作仍有兩個拓展方向值得探索:

進一步研究 token 難度的動態估計機制 , 實現更細粒度的選擇性蒸餾; 將 AdaSPEC 應用于多模態與推理型大模型中 , 以驗證其跨模態適配能力 。

    推薦閱讀