如你我所知，在大型語言模型（LLM）的運行邏輯中，隨著規模大小的增加，語言生成的質量會隨著提高。不過，這也導致了推理延遲的增加，從而對實際應用構成了重大挑戰。

從系統角度來看，LLM 推理主要受內存限制，主要延遲瓶頸源于加速器的內存帶寬而非算術計算。這一瓶頸是自回歸解碼的順序性所固有的，其中每次前向傳遞都需要將完整的模型參數從高帶寬內存傳輸到加速器緩存。該過程僅生成了單個的 token，沒有充分利用現代加速器的算術計算潛力，導致了效率低下。

為了解決這一問題，加速 LLM 推理的方法被提出，既可以增加解碼過程的算術強度（FLOPs 與總數據移動的比率），也能減少解碼步驟數量。這類方法以推測解碼（speculative decoding）為代表，使用較小的草稿（draft） 模型在每一步生成 token 序列，然后通過較大的原始模型進行細化以獲得可接受的延續。不過獲得合適的草稿模型仍然具有挑戰性，并且將草稿模型集成到分布式系統中更加困難。

在本文中，來自普林斯頓大學、Together.AI、伊利諾伊大學厄巴納 - 香檳分校等機構的研究者沒有使用單獨的草稿模型來順序生成候選輸出，而是重新審視并完善了在主干模型之上使用多個解碼頭加速推理的概念。他們發現，如果該技術得到有效應用，可以克服推測解碼的挑戰，從而無縫地集成到現有 LLM 系統中。

具體來講， 研究者提出了 MEDUSA，一種通過集成額外解碼頭（能夠同時預測多個 tokens）來增強 LLM 推理的方法。這些頭以參數高效的方式進行微調，并可以添加到任何現有模型中。至此，不需要任何新模型，MEDUSA 就可以輕松地集成地當前的 LLM 系統中（包括分布式環境），以確保友好用戶體驗。

值得關注的是，該論文作者之一 Tri Dao 是近來非常火爆的 Transformer 替代架構 Mamba 的兩位作者之一。他是 Together.AI 首席科學家，并即將成為普林斯頓大學計算機科學助理教授。

• 論文地址：https://arxiv.org/pdf/2401.10774.pdf
• GitHub 地址：https://arxiv.org/pdf/2401.10774.pdf

在具體實現中，研究者通過兩個關鍵見解進一步增強了 MEDUSA。首先，當前在每個解碼步驟生成單個候選延續的方法導致了可接受長度受限和計算資源的低效使用。為了解決這個問題，他們建議使用 MEDUSA 頭來生成多個候選延續，并通過對注意力掩碼的簡單調整來進行驗證。其次可以使用類似于推測解碼中的拒絕采樣方案來生成與原始模型具有相同分布的響應，但對于很多 LLM 應用來說通常不必要。

因此，研究者考慮或許可以引入一種典型的可接受方案，即從 MEDUSA 輸出中選擇合理的候選者。他們使用溫度作為閾值來管理原始模型預測的偏差，為拒絕采樣提供了一種有效的替代方案。這種方法有效地解決了拒絕采樣的局限性，比如在較高溫度下速度降低。

此外，為了給 LLM 配備預測性的 MEDUSA 頭，研究者提出了兩種針對不同場景量身定制的微調程序。對于計算資源有限或者目標是將 MEDUSA 納入現有模型而不影響其性能的情況，他們建議使用 MEDUSA-1。該方法需要的內存最少，并且可以使用類似于 QLoRA 中的量化技術來進一步優化，而不會因固定主干模型影響生成質量。

不過，對于 MEDUSA-1，主干模型的全部潛力無法得到充分利用。因此可以進一步進行微調，以提高 MEDUSA 頭的預測精度，并直接帶來更大加速。因此研究者提出了 MEDUSA - 2，它適用于計算資源充足或從基礎模型進行直接監督微調的場景。MEDUSA-2 的關鍵是一個訓練協議，它能夠對 MEDUSA 頭和主干模型進行聯合訓練，而不會影響模型下一個 token 的預測能力和輸出質量。

在實驗部分，研究者主要關注批大小為 1 的場景，這代表了 LLM 本地托管以供個人使用的用例。他們在不同大小和訓練設置下測試了 MEDUSA，包括 Vicuna-7B 和 13B（使用公共數據集訓練）、Vicuna -33B（使用私有數據集訓練）、Zephyr-7B（使用監督微調和對齊訓練）。 

結果表明，MEDUSA 在不影響生成質量的情況下，可以在不同的 promt 類型中實現 2.3 至 3.6 的推理加速。如下動圖為 Vicuna-7b 上有無 Medusa-1 時推理速度比較。

論文共同一作 Tianle Cai 表示，自 Medusa 項目推出以來，它在 TensorRT、TGI 以及眾多開源項目和公司中得到采用。在新的技術論文中，我們推出了用于全模型調優的 Medusa-2 方案、用于將 Medusa 集成到任何微調 LLM 的自蒸餾以及其他更多加速技術。

對于這項研究，Lepton AI 創始人賈揚清表示，Medusa 可能是他們見過的最優雅的加速推理解決方案之一，能夠與 int8/fp8、編譯等互補，在實踐中實現 2 倍性能增益。

并且，他們已將 Medusa 與很多現有優化方法、混合加速方案進行集成，結果在合理的并發下，加速保持正值，并在 A100 和 H100 等卡中尤其有效。此外，他們還已經為 Llama 模型訓練了通用 Medusa 頭。

方法概覽

MEDUSA 遵循推測解碼框架，其中每個解碼步驟主要由三個子步驟組成：(1) 生成候選者，(2) 處理候選者， (3) 接受候選者。對于 MEDUSA，(1) 是通過 MEDUSA 頭（head）實現的，(2) 是通過樹注意力（tree attention）實現的，并且由于 MEDUSA 頭位于原始主干模型之上，因此 (2) 中計算的 logits 可以用于子步驟 (1) 的下一個解碼步驟。最后一步 (3) 可以通過拒絕采樣（rejection sampling）或典型接受（typical acceptance）來實現。MEDUSA 的整體流程如下圖 1 所示。

關鍵組件

MEDUSA 的關鍵組件主要包括 MEDUSA 頭和樹注意力。

首先，MEDUSA 頭與原始主干模型一起進行訓練。其中，原始主干模型可以在訓練期間保持凍結狀態 (MEDUSA-1) 或一起訓練 (MEDUSA-2)。這種方法甚至可以在單個 GPU 上微調大模型，利用強大的基礎模型學得的表征。

此外，MEDUSA 頭的分布確保與原始模型的分布一致，從而緩解了分布偏移問題，并且 MEDUSA 不會增加服務系統設計的復雜性，對分布式設置很友好。

由于候選者增加會提高計算需求，該研究采用樹狀結構的注意力機制來同時處理多個候選者。這種注意力機制不同于傳統的因果注意力范式。在其框架內，只有來自同一 continuation 的 token 才被視為歷史數據。受圖神經網絡領域提出的將圖結構嵌入注意力的啟發，研究團隊還將樹結構合并到注意力掩碼中，如下圖 2 所示。

訓練策略

凍結主干模型來訓練 MEDUSA 頭的方法很簡單，并且需要的計算資源很少，但是將主干網絡與 MEDUSA 頭結合訓練可以顯著提高 MEDUSA 頭的準確性。因此，根據計算資源和用例的具體要求，研究團隊為 MEDUSA 頭提出了兩個級別的訓練策略，即 MEDUSA-1：凍結主干網絡，MEDUSA-2：聯合訓練。

最后，該研究提出了 MEDUSA 的兩個擴展，包括自蒸餾（self-distillation）和典型接受（typical acceptance），分別用于處理 MEDUSA 沒有可用訓練數據的情況和提高解碼過程的效率。

實驗

為了證明 MEDUSA 在不同設置下的有效性，該研究進行了兩組實驗：首先，在 Vicuna-7B/13B 模型上評估 MEDUSA，以展示 MEDUSA-1 和 MEDUSA-2 的性能；其次，在 Vicuna-33B 和 Zephyr-7B 模型上評估 MEDUSA，以研究自蒸餾的有效性，因為 Vicuna-33B 模型的訓練數據集不公開，而 Zephyr-7B 模型使用 RLHF 進行訓練。

用例研究 1：在 Vicuna-7B/13B 模型上評估 MEDUSA

在 Vicuna-7B/13B 模型上評估 MEDUSA-1、MEDUSA-2 的結果如下圖 4 所示。

用例研究 2：在 Vicuna-33B 和 Zephyr-7B 使用自蒸餾訓練

研究者關注了需要自蒸餾的情況，使用 Vicuna-33B 和 Zephyr-7B 作為示例。他們首先使用一些種子 prompt 來生成數據集，然后將 ShareGPT 和 UltraChat 作為種子數據集，并為以上兩個示例收集了包含大約 100k 樣本的數據集。

下表 1 展示了不同 MEDUSA-2 模型在 MT-Bench 基準下的加速比、開銷和質量。

下圖 5 為使用 MEDUSA-2 時不同模型的加速情況。

消融實驗

下圖 6a 比較了隨機采樣密集樹設置（藍點）和優化稀疏樹設置（紅星）的加速率。6b 比較了密集和稀疏樹設置的速度。

下圖 7 展示了不同采樣設置下，模型性能的比較分析。

兩階段微調的有效性。研究者針對 Vicuna-7B 模型，評估了兩種微調策略下的性能差異。