由于LLM架構固有的內存限制，使得生成又慢又貴。

對此，很多大佬都在尋找一種挽救的方法。Karpathy曾多次提出，大模型「投機采樣」對于推理時間優化是一個出色的方案。

但是，盡管投機解碼能夠加速生成速度，但因其太過復雜，并沒有被許多開發者采用。

今天，來自普林斯頓、UIUC等機構的華人團隊提出了全新的簡單框架：Medusa（美杜莎）。

沒有額外的草稿模型，研究人員只是引入了幾個額外的解碼頭，微調出「美杜莎頭」，能夠在單個A100-8G GPU，一天時間內完成訓練。

結果發現，Medusa直接讓模型推理加速約2倍。

Vicuna-7b與Medusa

為什么LLM生成效率低？

從系統角度來看，LLM生成遵循一種「以內存為限制」的計算模式，主要的延遲瓶頸來自內存讀取/寫入，而非算術計算。

這一問題的根源在于，自回歸解碼過程中固有的順序性。

即每次前向傳遞都需要將整個模型的參數，從高帶寬內存（HBM）傳輸到加速器的計算單元。

盡管每個樣本只產生一個token，但這個操作未能充分利用現代加速器的算術計算能力，由此導致模型效率低下。

在LLM崛起之前，應對這種低效率的常見方法是，簡單地「增加批大小」，從而實現更多token的并行生成。

但大模型的不斷涌現，讓情況變得更加復雜。

在這種情況下，增加批大小不僅會帶來更高的延遲，還會大大增加Transformer模型的鍵-值緩存的內存需求。

此外，這種低效率也體現在「成本結構」上。

截止到2023年9月，與僅處理提示相比，GPT-4的生成成本約高出2倍，Claude 2的生成成本大約高出3倍。

研究人員主要關注的重點是，改善LLM生成的延遲，同時Medusa也可以適用于需要平衡延遲和吞吐量的LLM服務。

每個token的價格

「投機解碼」是終極解決方案？

基于上述的挑戰，加速文本生成的一種吸引人策略是：更高效地利用計算資源。

具體來說，通過并行處理更多的token。

這種方法，采用了一個簡化的「草稿」模型，每一步都能快速生成一批token的候選項。

然后，這些候選token將通過原始的、全尺寸的語言模型進行驗證，以確定最合理的文本延續。

這一基本邏輯基于一個有趣的假設：「草稿」模型雖然小，但應該足夠嫻熟，能夠生成原始模型可以接受的序列。

如果這個假設成立，「草稿」模型可以迅速生成token序列，同時原始模型可以高效地并行驗證多個token，從而最大化計算吞吐量。

最近的研究表明，通過精心調整的草稿模型，投機解碼可以將延遲降低高達2.5倍，令人印象深刻。

然而，這種方法并非沒有挑戰：

1. 尋找理想的「草稿模型」（Draft Model）：找到一個「小而強大」的草稿模型，與原始模型很好地協調，說起來容易，做起來難。

2. 系統復雜性：在一個系統中托管兩個不同的模型會引入多層的復雜性，不論是計算還是操作，尤其是在分布式環境中。

3. 采樣效率低：使用投機解碼進行采樣時，需要使用一種重要性采樣方案。這會帶來額外的生成開銷，尤其是在較高的采樣溫度下。

這些復雜性和權衡限制了投機解碼的廣泛采用。因此，雖然投機解碼前景廣闊，但并未被廣泛采用。

Medusa：將簡單與高效融合

為了滿足對更加用戶友好，且功能強大的解決方案的需求，普林斯頓研究團隊推出了創新框架Medusa（美杜莎）。

它不僅加速了模型的生成，甚至讓LLM能夠讓更多人去訪問和使用。

據介紹，最新方法重新回顧了「Blockwise Parallel Decoding for Deep Autoregressive Models」論文中，一個被低估的寶藏：

回溯Transformer模型的發明，與其引入一個全新的「草稿」模型來預測后續token，為什么不簡單地擴展原始模型本身呢？

論文地址：https://arxiv.org/abs/1811.03115

這就是「Medusa head」（美杜莎頭）發揮作用的地方。

這些額外的解碼頭與原始模型無縫集成在一起，在每個生成的關鍵點產生token塊。

與草稿模型不同的是，Medusa head可以與原始模型一起訓練，而原始模型在訓練期間保持凍結狀態。

這種方法允許研究人員在單個GPU上微調大模型，充分利用強大的基礎模型學習到的表征。

此外，由于新的頭僅由一個與原始語言模型頭類似的層構成，因此Medusa不會增加服務系統設計的復雜性，并且適用于分布式環境。

單靠Medusa head，并不能達到將處理速度提高一倍的目標。

但這里有一個小技巧：將其與基于樹狀注意力機制配對使用時，就可以并行驗證由Medusa head生成的多個候選項。

這樣一來，Medusa head的預測能力速度提升2倍-3倍。

另外，研究人員摒棄了傳統的重要性采樣方案，專門為Medusa head生成創建了一種高效且高質量的替代方案。

這種新方法完全繞過了采樣開銷，甚至進一步提升了Medusa的加速步驟。

簡而言之，研究人員用一個簡單的系統解決了投機解碼的挑戰：

1. 沒有獨立的模型：不是引入一個新的草稿模型，而是在同一個模型上訓練多個解碼頭。

2. 輕松集成到現有系統中：訓練參數效率高，即使是GPU性能較差的情況下也可以進行。而且由于沒有額外的模型，無需調整分布式計算設置。

3. 將采樣視為一種放松：放寬與原始模型分布相匹配的要求，使得「非貪心生成」甚至比「貪心解碼」更快。

Medusa概述

具體來說，Medusa在LLM的最后隱藏狀態之上引入了多個頭，使其能夠并行預測多個后續token。

在使用Medusa head擴充模型時，原始模型在訓練期間被凍結，只有Medusa head經過微調。這種方法使得在單個GPU上對大型模型進行微調成為可能。

在推理過程中，每個頭為其指定的位置生成多個頂級預測。這些預測被組合成候選項，并使用基于樹狀注意力機制并行處理。

最后一步是，使用典型接受方案選擇合理的延續，被接受的最長候選項prefix將用于下一階段的解碼。

這樣，Medusa通過同時接受更多token，從而減少所需的解碼步驟，提高了解碼過程的效率。

接下來，讓我們深入了解Medusa的三個組成部分：Medusa head（美杜莎頭）、tree attention（樹狀注意力機制）和典型接受方案。

Medusa head（美杜莎頭）

那么，Medusa head究竟是什么呢？

它們類似于原始架構中的語言模型頭（因果Transformer模型的最后一層），但有一個變化：它們預測多個即將出現的token，而不僅僅是下一個token。

受到塊狀并行解碼方法的啟發，研究人員將每個Medusa head作為單層前饋網絡來實現，并增加了一個殘差連接。

訓練這些頭非常簡單。你可以使用訓練原始模型的相同語料庫，或者使用模型本身生成新的語料庫。

在這個訓練階段，原始模型保持不變；只有Medusa head經過微調。

這種有針對性的訓練會帶來一個參數效率極高的過程，并能迅速達到收斂。

尤其是，與在推測解碼方法中訓練獨立的草稿模型的計算復雜性相比，優勢更加突出。

在研究人員測試的Vicuna模型上，Medusa head在預測next-next token方面準確率達到60％，位列第一。同時，它仍然有改進的空間。

tree attention（樹狀注意力機制）

測試中，團隊發現一些引人注目的指標：盡管對于預測next-next token的第1名準確率大約在60%左右，但第5名的準確率卻超過了80%。

這個顯著的提高表明，如果可以有效利用Medusa head產生多個排名靠前的預測，就可以增加每個解碼步驟生成的token數量。

為了實現這個目標，研究人員首先通過從每個Medusa head的頂級預測中，獲取笛卡爾積來創建一組候選項。

然后，根據圖神經網絡的思路，將依賴關系圖編碼到注意力機制中，以便可以并行處理多個候選項。

比如，使用第一個Medusa head的前2個預測，以及第二個Medusa head的前3個預測，如下所示。

樹狀注意力

上圖中的可視化效果，展示了使用樹狀注意力同時處理多個候選項的過程。

在這種情況下，第一個頭的任何預測都可以與第二個頭的任何預測配對，形成一個多級樹結構。

這個樹的每一層，都對應著一個Medusa head的預測。由此，能夠產生2*3=6個候選項。

這些候選項中的每一個對應于樹結構內的一個不同的分支。

為了確保每個token只訪問其前置token，研究人員設計了一種注意力掩碼，專門允許注意流從當前token返回到其前置token。

通過這樣做，并相應地位位置編碼設置位置索引，可以在不需要增加批大小的情況下，同時處理各種候選項。

研究人員還指出，一些研究也采用了非常相似的樹狀注意力思想。

與它們相比，最新的方法更傾向于一種更簡單的樹狀注意力形式，在推理期間，樹狀模式是規則的且固定的，使得樹狀注意力掩碼進行預處理，進一步提高了效率。

典型接受

在早期關于投機解碼的研究中，「重要性采樣」被用來產生與原始模型預測密切相關的各種輸出結果。

然而，后來的研究表明，隨著采樣溫度升高，這種方法的效率往往會降低。

簡單來說，如果你的草稿模型和你原來的模型一樣好，理想情況下，你應該接受它的所有輸出，從而使這個過程變得超級高效。

然而，重要性采樣可能會在中間環節，拒絕這個解決方案。

在現實世界中，我們經常調整采樣溫度，只是為了控制模型的創造力，不一定是為了匹配原始模型的分布。

那么為什么不把重點放在，可接受可信的候選項上呢？

對此，研究人員引入了「典型接受」方案。

從現有的截斷采樣（truncation sampling）中汲取靈感，普林斯頓研究人員目標是根據原始模型選擇足夠可能的候選項。

研究人員根據根據原始模型的預測概率設置一個閾值，如果候選項超過這個閾值，就會被接受。

用專業術語來說，研究人員采用硬閾值，以及熵相關閾值中的最小值，來決定是否像截斷采樣那樣接受候選項。

這確保了在解碼過程中，選擇有意義的token和合理延續。

研究人員總是使用貪心解碼接受第一個token，確保在每個步驟中至少生成一個token。最終輸出是通過研究接受測試的最長序列。

這一方法的高明之處在于，適應性強。

如果你將采樣溫度設置為0，它就會恢復到最有效的形式——貪心解碼。當溫度升高時，最新研究的方法變得更加高效，可以接受更長的序列。

并且，研究人員還已經通過嚴格的測試證實了這一說法。

因此，從本質上講，最新的典型接受方案提供了一種更有效的方式，讓LLM進行更創造性地輸出。

Llama「吐口水」可以有多快？

研究人員使用了專門針對聊天應用進行微調的Vicuna模型對Medusa進行了測試。

這些模型的大小不同，參數數量分別為7B、13B和33B。

研究的目標是，測量Medusa在實際聊天機器人環境中如何加速這些模型。

在訓練Medusa head時，研究人員采用了一種簡單的方法，使用了公開可用的ShareGPT數據集。這是最初用于Vicuna模型的訓練數據的子集，僅進行了一個epoch的訓練。

而且，關鍵是，根據模型參數的大小，整個訓練過程可以在幾小時到一天之內完成，僅需要一塊A100-80G GPU。

值得注意的是，Medusa可以輕松與量化的基礎模型結合使用，以減少內存需求。

研究人員利用這一優勢，在訓練33B模型時使用了8位量化。為了模擬實際環境，他們使用了MT bench進行評估。

憑借其簡單的設計，Medusa在各種用例中，始終能夠實現約2倍的絕對時間加速。

值得注意的是，通過Medusa的優化，33B參數的Vicuna模型可以像13B模型一樣迅速運行。

消融實驗

Medusa head的配置

在利用美杜莎頭的預測能力時，研究人員靈活地選擇了每個頭應考慮多少個頂級候選項。

例如可以選擇第一個頭的前3個預測和第二個頭的前2個預測。當對這些頂級候選項進行笛卡爾積運算時，就生成了6個延續選項供模型評估。

這種可配置性，是需要進行權衡的。

一方面，選擇更多的頂級預測會增加模型接受生成token的可能性。另一方面，它還會增加每個解碼步驟的計算開銷。

為了找到最佳平衡，研究人員嘗試了各種配置，并確定了最有效的設置，如附圖所示。

典型接受的閾值

在典型接受方案中，一個關鍵的超參數，稱為「閾值」，幫助研究人員根據模型自身的預測，來確定生成的token是否合理。

閾值越高，接受標準越嚴格，從而影響通過這種方法獲得的整體加速。

研究人員通過在MT bench的2個以創造力為導向的任務上進行實驗，來探討質量和加速之間的這種權衡。

圖中顯示的結果表明，與貪心解碼方法相比，典型接受能夠加速10%。這種加速明顯優于使用隨機采樣的投機解碼方法，后者與貪心解碼相比實際減慢了過程。

作者介紹

Tianle Cai（蔡天樂）

共同一作蔡天樂是普林斯頓大學的博士生，導師為Kai Li、Jason D. Lee。

就讀博士期間，他在Xuezhi Wang和Denny Zhou指導下在Google DeepMind實習，還在Sébastien Bubeck和Debadeepta Dey的指導下在Microsoft Research完成了實習。

他曾在北大拿到了本科學位，主修應用數學，同時也主修計算機科學雙學位，在Liwei Wang教授的指導下，開始了機器學習領域的研究。

蔡天樂的學術興趣跨越了機器學習的廣泛領域，包括優化、表征學習、架構設計（重點是Transfomer、圖神經網絡等），以及最近的系統架構協同設計。

Yuhong Li

共同一作Yuhong (Jesse) Li是伊利諾伊大學厄巴納-香檳分校（UIUC）的ECE博士生，導師是Deming Chen教授。

此前，他在北京郵電大學獲得了學士學位，興趣是高效機器學習。

Zhengyang Geng（耿正陽）

Zhengyang Geng是卡內基梅隆大學（CMU）的計算機科學博士生，由J. Zico Kolter指導。

之前，他在北京大學做研究助理，由Zhouchen Lin指導。致力于識別和開發能夠自組織復雜系統的結構。

Hongwu Peng

Hongwu Peng是康涅狄格大學計算機科學與工程系的博士生。

此前，他于2018年獲得華中科技大學電氣工程學士學位，于2020年獲得阿肯色大學電氣工程碩士學位。

Tri Dao

Tri Dao是生成式AI初創公司Together AI的首席科學家。2024年9月起，將出任普林斯頓大學計算機科學助理教授。

此前，他在斯坦福大學獲得了計算機博士學位，導師是Christopher Ré和Stefano Ermon。

Tri Dao的研究興趣在于機器學習和系統，重點關注：高效Transformer訓練和推理；遠程記憶的序列模型；緊湊型深度學習模型的結構化稀疏性。

項目鳴謝：Zhuohan Li，Shaojie Bai，Denny Zhou，Yanping Huang，stability.ai，together.ai，ChatGPT。