作者：Winnie

今天為大家介紹一個新技術—Medusa，它旨在加速大型語言模型（LLM）的生成。盡管其設計簡單，但 Medusa能夠將LLM的生成效率提高約2倍。讓我們看看它是怎么做到的吧！

為什么LLM生成低效？

LLM在生成時的效率問題主要是由內存讀/寫操作帶來的延遲，而這個問題源自自動回歸解碼過程的順序性特點。每次的前向傳播都需要頻繁地移動模型參數，盡管這只產生一個結果，但卻沒有完全利用現代硬件的計算潛能。傳統的解決方式（如增大批次大小）在LLM的場景下卻不再適用，因為這不僅會增加延遲，還會引發內存問題。

不僅如此，這種低效還帶來了額外的生成成本。例如，GPT-4的生成成本比僅僅處理prompt高了兩倍，Claude2則大約高出3倍。因此，加速LLM的低效生成是一個亟待解決的問題。

Medusa來了！

面對推測性解碼的復雜性，研究人員推出了Medusa技術，這個框架回歸了Transformer模型的本質，減少了復雜度，增強了效率，讓每個生成階段都能快速產出結果。當將Medusa與基于樹的注意機制結合時，生成速度提高了2到3倍。

接下來，讓我們看一看Msdusa都做了哪些改進吧！

Medusa總體框架

Medusa的核心在于它在LLM的最后隱藏層上增加的多個Heads，使它們并行工作，預測接下來的內容。

當將Medusa Heads加入模型時，你會發現，原始模型保持不變，而只有這些Medusa Heads進行微調。在真正使用時，每個Medusa Head都會為其位置產生預測，這些預測會被組合、處理，最終給出最佳結果。

通過同時接受更多的tokens來增強解碼過程的效率，從而減少了所需的解碼步驟數量。

Medusa Heads

Medusa Heads與原有的語言模型頭相似，但卻擁有一個獨特的優勢：它們可以預測多個即將出現的token，而不僅僅是下一個。這種方法從Blockwise Parallel Decoding方法中汲取靈感，將每個Medusa頭設計為一個單層的前饋網絡，且增強了殘差連接。

訓練這些Medusa Heads非常方便！你可以使用用于訓練原始模型的同一語料庫，或者使用模型本身生成一個新的語料庫來訓練它們。在訓練階段，原始的模型保持靜態，僅Medusa Heads進行微調。這種有針對性的訓練產生了一個參數效率極高的過程，可以迅速實現收斂—尤其是與speculative decoding方法中訓練單獨的draft model的計算密集度相比。

Medusa Heads的表現相當出色，它在預測“下一個”token時的top-1準確率約為60%。但這僅僅是個開始，它還有很大的提升空間。

Tree attention

通過Medusa Heads的測試，研究人員發現：雖然預測“下下一個”token的top-1準確率僅約為60%，但top-5準確率卻飆升至超過80%。這一顯著的提高表明，如果我們可以巧妙地利用Medusa Heads做出的多個top排名預測，就可以顯著增加每個解碼步驟生成的tokens數量。

實現這一目標的方式是首先構造一個候選集，這個集合由每個Medusa Head的預測結果的笛卡爾積形成。然后，依賴圖被編碼到注意力機制中，允許多個候選項目并行處理，這是受到圖神經網絡思想的啟發。例如，在一個實際應用中，可以從第一個Medusa頭部獲取前兩個預測，從第二個頭部獲取前三個預測，并將它們組合成一個多層樹結構。在這種結構中，一個注意力掩碼被實施，僅限制注意力于一個token的前一個token，從而保持歷史上下文。通過這種方式，可以同時處理多個候選項，而無需增加批次大小。

下圖是Tree attention機制用于并行處理多個候選項目的一個可視化示例。在一個示例中，來自第一個Medusa頭部的前兩個預測和來自第二個頭部的前三個預測產生了2*3=6個候選項。這些候選項中的每一個都對應于樹結構中的一個不同分支。為了保證每個token只能訪問其前面的token，注意力掩碼，該掩碼僅允許注意力從當前token流向其前面的token。位置編碼的位置指數將根據這種結構進行調整。通過這種方式，可以確保歷史上下文的完整性和連貫性，同時提高解碼步驟的效率和準確性。

值得注意的是，與一些獨立的研究相比，該方法傾向于使用簡化形式的樹狀注意力，其中樹的模式在推斷期間是規則和固定的，這允許預處理樹狀注意力掩碼，進而提高效率。通過創新這種解碼方法，它不僅提供了一個新的解碼路徑，而且為更精確和高效的未來預測打開了新的可能性。

Typical acceptance

在早期關于投機解碼的研究中，重要性采樣技術用于產生與原始模型預測緊密相符的多樣化輸出。但隨后的研究表明，隨著“creativity dial”或采樣溫度的增加，這種方法的效率會降低。簡而言之，如果一個draft model與原始模型一樣優秀，理論上應接受其所有輸出，使過程極為高效。但是重要性采樣可能會在某個階段拒絕這種方案。

實際上，人們常常僅調整采樣溫度來控制模型的創造力，而不是嚴格匹配原始模型的分布。那么為什么不只是接受看似合理的候選項呢？Typical acceptance策略受到截斷采樣的啟發，目的是選取根據原始模型被視為足夠可能的候選項。通過設置基于原始模型預測概率的閾值，如果候選項超過這個閾值，則將其接受。

在技術語言中，我們采用硬閾值和依賴于熵的閾值中的最小值來決定是否接受一個候選項，如截斷采樣中所做。這確保在解碼期間選擇了有意義的標記和合理的延續。第一個標記總是通過貪婪解碼被接受，確保每一步至少生成一個標記。最終輸出是通過接受測試的最長序列。這種方法的優點在于其適應性。如果將采樣溫度設置為零，它將簡單地退化為最有效的形式——貪婪解碼。提高溫度會使方法變得更加高效，允許更長的接受序列，這一點已通過嚴格測試得到驗證。

性能測試

在Vicuna模型上測試了Medusa，這些模型是特別為聊天應用優化和調整的羊駝模型，其大小不同，參數數量分別為7B、13B和33B。目標是衡量Medusa在現實世界的聊天機器人環境中能夠多大程度上加速這些模型的運行。

訓練Medusa頭部選擇了簡單的方式，使用了公開的ShareGPT數據集，這是最初用于訓練Vicuna模型的數據的一個子集，只進行了一個時代的訓練。

這里的重點是——整個訓練過程可以在幾小時到一天之內完成，具體取決于模型的大小，全部在單個A100-80G GPU上完成。顯著的是，Medusa可以與量化基模型輕松結合，從而減少內存需求。為了利用這一優勢，在訓練33B模型時使用了8位量化。

為模擬現實環境，采用了MT測試臺進行評估。結果是令人鼓舞的：Medusa借助其簡單的設計，在各種用例中穩定實現了約2倍的實際運行時間加速。顯著的是，有了Medusa的優化，33B參數的Vicuna模型可以與13B模型一樣快速運行。

結語

Medusa技術致力于通過多層頭部預測方法來加速LLM的語言生成速度。該研究中引入了多個Medusa頭和Tree attention機制，通過預測多個即將出現的標記而非一個來優化生成速度，同時還保持了高準確率。此外，研究還提出了Typical acceptance方案，它基于原始模型的預測概率來選擇候選項，而不是依賴重要性抽樣，使得創意輸出更為高效和自適應。

在實際測試中，Medusa成功地將Vicuna模型的運行速度提高了大約兩倍，證明了其在現實世界的聊天機器人環境中的實用性和效果。整體來看，Medusa為開發更快、更有效的聊天機器人開辟了新的可能，顯示出在語言模型生成領域的巨大潛力。

編輯：黃飛