作者：Winnie

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

為什么LLM生成低效？

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

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

Medusa來了！

面對推測性解碼的復(fù)雜性，研究人員推出了Medusa技術(shù)，這個框架回歸了Transformer模型的本質(zhì)，減少了復(fù)雜度，增強(qiáng)了效率，讓每個生成階段都能快速產(chǎn)出結(jié)果。當(dāng)將Medusa與基于樹的注意機(jī)制結(jié)合時，生成速度提高了2到3倍。

接下來，讓我們看一看Msdusa都做了哪些改進(jìn)吧！

Medusa總體框架

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

當(dāng)將Medusa Heads加入模型時，你會發(fā)現(xiàn)，原始模型保持不變，而只有這些Medusa Heads進(jìn)行微調(diào)。在真正使用時，每個Medusa Head都會為其位置產(chǎn)生預(yù)測，這些預(yù)測會被組合、處理，最終給出最佳結(jié)果。

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

Medusa Heads

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

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

Medusa Heads的表現(xiàn)相當(dāng)出色，它在預(yù)測“下一個”token時的top-1準(zhǔn)確率約為60%。但這僅僅是個開始，它還有很大的提升空間。

Tree attention

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

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

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

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

Typical acceptance

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

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

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

性能測試

在Vicuna模型上測試了Medusa，這些模型是特別為聊天應(yīng)用優(yōu)化和調(diào)整的羊駝模型，其大小不同，參數(shù)數(shù)量分別為7B、13B和33B。目標(biāo)是衡量Medusa在現(xiàn)實世界的聊天機(jī)器人環(huán)境中能夠多大程度上加速這些模型的運行。

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

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

為模擬現(xiàn)實環(huán)境，采用了MT測試臺進(jìn)行評估。結(jié)果是令人鼓舞的：Medusa借助其簡單的設(shè)計，在各種用例中穩(wěn)定實現(xiàn)了約2倍的實際運行時間加速。顯著的是，有了Medusa的優(yōu)化，33B參數(shù)的Vicuna模型可以與13B模型一樣快速運行。

結(jié)語

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

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

編輯：黃飛