由于LLM架構(gòu)固有的內(nèi)存限制，使得生成又慢又貴。

對此，很多大佬都在尋找一種挽救的方法。Karpathy曾多次提出，大模型「投機(jī)采樣」對于推理時(shí)間優(yōu)化是一個(gè)出色的方案。

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

今天，來自普林斯頓、UIUC等機(jī)構(gòu)的華人團(tuán)隊(duì)提出了全新的簡單框架：Medusa（美杜莎）。

沒有額外的草稿模型，研究人員只是引入了幾個(gè)額外的解碼頭，微調(diào)出「美杜莎頭」，能夠在單個(gè)A100-8G GPU，一天時(shí)間內(nèi)完成訓(xùn)練。

結(jié)果發(fā)現(xiàn)，Medusa直接讓模型推理加速約2倍。

Vicuna-7b與Medusa

為什么LLM生成效率低？

從系統(tǒng)角度來看，LLM生成遵循一種「以內(nèi)存為限制」的計(jì)算模式，主要的延遲瓶頸來自內(nèi)存讀取/寫入，而非算術(shù)計(jì)算。

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

即每次前向傳遞都需要將整個(gè)模型的參數(shù)，從高帶寬內(nèi)存（HBM）傳輸?shù)郊铀倨鞯挠?jì)算單元。

盡管每個(gè)樣本只產(chǎn)生一個(gè)token，但這個(gè)操作未能充分利用現(xiàn)代加速器的算術(shù)計(jì)算能力，由此導(dǎo)致模型效率低下。

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

但大模型的不斷涌現(xiàn)，讓情況變得更加復(fù)雜。

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

此外，這種低效率也體現(xiàn)在「成本結(jié)構(gòu)」上。

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

研究人員主要關(guān)注的重點(diǎn)是，改善LLM生成的延遲，同時(shí)Medusa也可以適用于需要平衡延遲和吞吐量的LLM服務(wù)。

每個(gè)token的價(jià)格

「投機(jī)解碼」是終極解決方案？

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

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

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

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

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

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

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

然而，這種方法并非沒有挑戰(zhàn)：

1. 尋找理想的「草稿模型」（Draft Model）：找到一個(gè)「小而強(qiáng)大」的草稿模型，與原始模型很好地協(xié)調(diào)，說起來容易，做起來難。

2. 系統(tǒng)復(fù)雜性：在一個(gè)系統(tǒng)中托管兩個(gè)不同的模型會引入多層的復(fù)雜性，不論是計(jì)算還是操作，尤其是在分布式環(huán)境中。

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

這些復(fù)雜性和權(quán)衡限制了投機(jī)解碼的廣泛采用。因此，雖然投機(jī)解碼前景廣闊，但并未被廣泛采用。

Medusa：將簡單與高效融合

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

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

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

回溯Transformer模型的發(fā)明，與其引入一個(gè)全新的「草稿」模型來預(yù)測后續(xù)token，為什么不簡單地?cái)U(kuò)展原始模型本身呢？

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

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

這些額外的解碼頭與原始模型無縫集成在一起，在每個(gè)生成的關(guān)鍵點(diǎn)產(chǎn)生token塊。

與草稿模型不同的是，Medusa head可以與原始模型一起訓(xùn)練，而原始模型在訓(xùn)練期間保持凍結(jié)狀態(tài)。

這種方法允許研究人員在單個(gè)GPU上微調(diào)大模型，充分利用強(qiáng)大的基礎(chǔ)模型學(xué)習(xí)到的表征。

此外，由于新的頭僅由一個(gè)與原始語言模型頭類似的層構(gòu)成，因此Medusa不會增加服務(wù)系統(tǒng)設(shè)計(jì)的復(fù)雜性，并且適用于分布式環(huán)境。

單靠Medusa head，并不能達(dá)到將處理速度提高一倍的目標(biāo)。

但這里有一個(gè)小技巧：將其與基于樹狀注意力機(jī)制配對使用時(shí)，就可以并行驗(yàn)證由Medusa head生成的多個(gè)候選項(xiàng)。

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

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

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

簡而言之，研究人員用一個(gè)簡單的系統(tǒng)解決了投機(jī)解碼的挑戰(zhàn)：

1. 沒有獨(dú)立的模型：不是引入一個(gè)新的草稿模型，而是在同一個(gè)模型上訓(xùn)練多個(gè)解碼頭。

2. 輕松集成到現(xiàn)有系統(tǒng)中：訓(xùn)練參數(shù)效率高，即使是GPU性能較差的情況下也可以進(jìn)行。而且由于沒有額外的模型，無需調(diào)整分布式計(jì)算設(shè)置。

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

Medusa概述

具體來說，Medusa在LLM的最后隱藏狀態(tài)之上引入了多個(gè)頭，使其能夠并行預(yù)測多個(gè)后續(xù)token。

在使用Medusa head擴(kuò)充模型時(shí)，原始模型在訓(xùn)練期間被凍結(jié)，只有Medusa head經(jīng)過微調(diào)。這種方法使得在單個(gè)GPU上對大型模型進(jìn)行微調(diào)成為可能。

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

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

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

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

Medusa head（美杜莎頭）

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

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

受到塊狀并行解碼方法的啟發(fā)，研究人員將每個(gè)Medusa head作為單層前饋網(wǎng)絡(luò)來實(shí)現(xiàn)，并增加了一個(gè)殘差連接。

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

在這個(gè)訓(xùn)練階段，原始模型保持不變；只有Medusa head經(jīng)過微調(diào)。

這種有針對性的訓(xùn)練會帶來一個(gè)參數(shù)效率極高的過程，并能迅速達(dá)到收斂。

尤其是，與在推測解碼方法中訓(xùn)練獨(dú)立的草稿模型的計(jì)算復(fù)雜性相比，優(yōu)勢更加突出。

在研究人員測試的Vicuna模型上，Medusa head在預(yù)測next-next token方面準(zhǔn)確率達(dá)到60％，位列第一。同時(shí)，它仍然有改進(jìn)的空間。

tree attention（樹狀注意力機(jī)制）

測試中，團(tuán)隊(duì)發(fā)現(xiàn)一些引人注目的指標(biāo)：盡管對于預(yù)測next-next token的第1名準(zhǔn)確率大約在60%左右，但第5名的準(zhǔn)確率卻超過了80%。

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

為了實(shí)現(xiàn)這個(gè)目標(biāo)，研究人員首先通過從每個(gè)Medusa head的頂級預(yù)測中，獲取笛卡爾積來創(chuàng)建一組候選項(xiàng)。

然后，根據(jù)圖神經(jīng)網(wǎng)絡(luò)的思路，將依賴關(guān)系圖編碼到注意力機(jī)制中，以便可以并行處理多個(gè)候選項(xiàng)。

比如，使用第一個(gè)Medusa head的前2個(gè)預(yù)測，以及第二個(gè)Medusa head的前3個(gè)預(yù)測，如下所示。

樹狀注意力

上圖中的可視化效果，展示了使用樹狀注意力同時(shí)處理多個(gè)候選項(xiàng)的過程。

在這種情況下，第一個(gè)頭的任何預(yù)測都可以與第二個(gè)頭的任何預(yù)測配對，形成一個(gè)多級樹結(jié)構(gòu)。

這個(gè)樹的每一層，都對應(yīng)著一個(gè)Medusa head的預(yù)測。由此，能夠產(chǎn)生2*3=6個(gè)候選項(xiàng)。

這些候選項(xiàng)中的每一個(gè)對應(yīng)于樹結(jié)構(gòu)內(nèi)的一個(gè)不同的分支。

為了確保每個(gè)token只訪問其前置token，研究人員設(shè)計(jì)了一種注意力掩碼，專門允許注意流從當(dāng)前token返回到其前置token。

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

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

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

典型接受

在早期關(guān)于投機(jī)解碼的研究中，「重要性采樣」被用來產(chǎn)生與原始模型預(yù)測密切相關(guān)的各種輸出結(jié)果。

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

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

然而，重要性采樣可能會在中間環(huán)節(jié)，拒絕這個(gè)解決方案。

在現(xiàn)實(shí)世界中，我們經(jīng)常調(diào)整采樣溫度，只是為了控制模型的創(chuàng)造力，不一定是為了匹配原始模型的分布。

那么為什么不把重點(diǎn)放在，可接受可信的候選項(xiàng)上呢？

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

從現(xiàn)有的截?cái)嗖蓸樱╰runcation sampling）中汲取靈感，普林斯頓研究人員目標(biāo)是根據(jù)原始模型選擇足夠可能的候選項(xiàng)。

研究人員根據(jù)根據(jù)原始模型的預(yù)測概率設(shè)置一個(gè)閾值，如果候選項(xiàng)超過這個(gè)閾值，就會被接受。

用專業(yè)術(shù)語來說，研究人員采用硬閾值，以及熵相關(guān)閾值中的最小值，來決定是否像截?cái)嗖蓸幽菢咏邮芎蜻x項(xiàng)。

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

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

這一方法的高明之處在于，適應(yīng)性強(qiáng)。

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

并且，研究人員還已經(jīng)通過嚴(yán)格的測試證實(shí)了這一說法。

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

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

研究人員使用了專門針對聊天應(yīng)用進(jìn)行微調(diào)的Vicuna模型對Medusa進(jìn)行了測試。

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

研究的目標(biāo)是，測量Medusa在實(shí)際聊天機(jī)器人環(huán)境中如何加速這些模型。

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

而且，關(guān)鍵是，根據(jù)模型參數(shù)的大小，整個(gè)訓(xùn)練過程可以在幾小時(shí)到一天之內(nèi)完成，僅需要一塊A100-80G GPU。

值得注意的是，Medusa可以輕松與量化的基礎(chǔ)模型結(jié)合使用，以減少內(nèi)存需求。

研究人員利用這一優(yōu)勢，在訓(xùn)練33B模型時(shí)使用了8位量化。為了模擬實(shí)際環(huán)境，他們使用了MT bench進(jìn)行評估。

憑借其簡單的設(shè)計(jì)，Medusa在各種用例中，始終能夠?qū)崿F(xiàn)約2倍的絕對時(shí)間加速。

值得注意的是，通過Medusa的優(yōu)化，33B參數(shù)的Vicuna模型可以像13B模型一樣迅速運(yùn)行。

消融實(shí)驗(yàn)

Medusa head的配置

在利用美杜莎頭的預(yù)測能力時(shí)，研究人員靈活地選擇了每個(gè)頭應(yīng)考慮多少個(gè)頂級候選項(xiàng)。

例如可以選擇第一個(gè)頭的前3個(gè)預(yù)測和第二個(gè)頭的前2個(gè)預(yù)測。當(dāng)對這些頂級候選項(xiàng)進(jìn)行笛卡爾積運(yùn)算時(shí)，就生成了6個(gè)延續(xù)選項(xiàng)供模型評估。

這種可配置性，是需要進(jìn)行權(quán)衡的。

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

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

典型接受的閾值

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

閾值越高，接受標(biāo)準(zhǔn)越嚴(yán)格，從而影響通過這種方法獲得的整體加速。

研究人員通過在MT bench的2個(gè)以創(chuàng)造力為導(dǎo)向的任務(wù)上進(jìn)行實(shí)驗(yàn)，來探討質(zhì)量和加速之間的這種權(quán)衡。

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

作者介紹

Tianle Cai（蔡天樂）

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

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

他曾在北大拿到了本科學(xué)位，主修應(yīng)用數(shù)學(xué)，同時(shí)也主修計(jì)算機(jī)科學(xué)雙學(xué)位，在Liwei Wang教授的指導(dǎo)下，開始了機(jī)器學(xué)習(xí)領(lǐng)域的研究。

蔡天樂的學(xué)術(shù)興趣跨越了機(jī)器學(xué)習(xí)的廣泛領(lǐng)域，包括優(yōu)化、表征學(xué)習(xí)、架構(gòu)設(shè)計(jì)（重點(diǎn)是Transfomer、圖神經(jīng)網(wǎng)絡(luò)等），以及最近的系統(tǒng)架構(gòu)協(xié)同設(shè)計(jì)。

Yuhong Li

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

此前，他在北京郵電大學(xué)獲得了學(xué)士學(xué)位，興趣是高效機(jī)器學(xué)習(xí)。

Zhengyang Geng（耿正陽）

Zhengyang Geng是卡內(nèi)基梅隆大學(xué)（CMU）的計(jì)算機(jī)科學(xué)博士生，由J. Zico Kolter指導(dǎo)。

之前，他在北京大學(xué)做研究助理，由Zhouchen Lin指導(dǎo)。致力于識別和開發(fā)能夠自組織復(fù)雜系統(tǒng)的結(jié)構(gòu)。

Hongwu Peng

Hongwu Peng是康涅狄格大學(xué)計(jì)算機(jī)科學(xué)與工程系的博士生。

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

Tri Dao

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

此前，他在斯坦福大學(xué)獲得了計(jì)算機(jī)博士學(xué)位，導(dǎo)師是Christopher Ré和Stefano Ermon。

Tri Dao的研究興趣在于機(jī)器學(xué)習(xí)和系統(tǒng)，重點(diǎn)關(guān)注：高效Transformer訓(xùn)練和推理；遠(yuǎn)程記憶的序列模型；緊湊型深度學(xué)習(xí)模型的結(jié)構(gòu)化稀疏性。

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