近日，蘋果與英偉達(dá)宣布合作，旨在加速和優(yōu)化大語(yǔ)言模型（LLM）的推理性能。
為了改善傳統(tǒng)自回歸 LLM 推理效率低和內(nèi)存帶寬小的問(wèn)題，今年早些時(shí)候，蘋果機(jī)器學(xué)習(xí)的研究人員發(fā)布并開(kāi)源了一種名為「ReDrafter」（Recurrent Drafter，循環(huán)草稿模型）的推測(cè)解碼技術(shù)。
▲圖源：GitHub
目前，ReDrafter 已經(jīng)整合到英偉達(dá)的可擴(kuò)展推理方案「TensorRT-LLM」當(dāng)中，后者是基于「TensorRT」深度學(xué)習(xí)編譯框架的專為優(yōu)化 LLM 推理而設(shè)計(jì)的開(kāi)源庫(kù)，支持包括「Medusa」等推測(cè)解碼方法。
不過(guò)，由于 ReDrafter 所包含的算法使用了之前從未用過(guò)的運(yùn)算符，因此英偉達(dá)方面添加了新的運(yùn)算符，或者公開(kāi)了現(xiàn)有的運(yùn)算符，大大提高了 TensorRT-LLM 適應(yīng)復(fù)雜模型和解碼方式的能力。
▲圖源：GitHub
據(jù)悉，ReDrafter 推測(cè)解碼通過(guò)三個(gè)關(guān)鍵技術(shù)來(lái)加速 LLM 的推理過(guò)程：
RNN 草稿模型
動(dòng)態(tài)樹(shù)注意力算法
知識(shí)蒸餾訓(xùn)練
RNN 草稿模型是 ReDrafter 的「核心」組件。它使用循環(huán)神經(jīng)網(wǎng)絡(luò)（Recurrent Neural Network），基于 LLM 的「隱藏狀態(tài)」來(lái)預(yù)測(cè)接下來(lái)可能出現(xiàn)的 tokens 序列，其能夠捕捉局部的時(shí)間依賴性，從而提高預(yù)測(cè)準(zhǔn)確性。
這個(gè)模型的工作原理是：LLM 在文本生成過(guò)程中首先生成一個(gè)初始 token，然后 RNN 草稿模型利用該 token 和 LLM 的最后一層隱藏狀態(tài)作為輸入進(jìn)行束搜索（Beam Search），進(jìn)而生成多個(gè)候選 tokens 序列。
與傳統(tǒng)自回歸 LLM 每次只生成一個(gè) token 不同，通過(guò) RNN 草稿模型的預(yù)測(cè)輸出，ReDrafter 能夠在每個(gè)解碼步驟生成多個(gè) tokens，大大減少了需要調(diào)用 LLM 驗(yàn)證的次數(shù)，從而提高了整體的推理速度。
▲圖源：arXiv
動(dòng)態(tài)樹(shù)注意力算法（Dynamic Tree Attention）則是一種優(yōu)化束搜索結(jié)果的算法。
我們已經(jīng)知道，在束搜索過(guò)程中會(huì)產(chǎn)生多個(gè)候選序列，而這些序列往往存在共享的前綴。動(dòng)態(tài)樹(shù)注意力算法會(huì)識(shí)別出這些共享前綴，并將它們從需要驗(yàn)證的 tokens 中去除，從而減少 LLM 需要處理的數(shù)據(jù)量。
某些情況下，該算法能將需要驗(yàn)證的 tokens 數(shù)量減少 30% 到 60%。這意味著使用動(dòng)態(tài)樹(shù)注意力算法后，ReDrafter 能夠更高效地利用計(jì)算資源，進(jìn)一步提高推理速度。
▲圖源：NVIDIA
知識(shí)蒸餾是一種模型壓縮技術(shù)，它能夠?qū)⒁粋�(gè)大型、復(fù)雜的模型（教師模型）的知識(shí)「蒸餾」到一個(gè)更小、更簡(jiǎn)單的模型（學(xué)生模型）中。在 ReDrafter 中，RNN 草稿模型作為學(xué)生模型通過(guò)知識(shí)蒸餾從 LLM（教師模型）中學(xué)習(xí)。
具體來(lái)講，蒸餾訓(xùn)練過(guò)程中，LLM 會(huì)給出一系列下一個(gè)可能詞的「概率分布」，開(kāi)發(fā)人員會(huì)基于這個(gè)概率分布數(shù)據(jù)訓(xùn)練 RNN 草稿模型，然后計(jì)算兩個(gè)模型概率分布之間的差異，并通過(guò)優(yōu)化算法使這個(gè)差異最小化。
在這個(gè)過(guò)程中，RNN 草稿模型不斷學(xué)習(xí) LLM 的概率預(yù)測(cè)模式，從而在實(shí)際應(yīng)用中能夠生成與 LLM 相似的文本。
通過(guò)知識(shí)蒸餾訓(xùn)練，RNN 草稿模型更好地捕捉到語(yǔ)言的規(guī)律和模式，從而更準(zhǔn)確地預(yù)測(cè) LLM 的輸出，并且因?yàn)槠漭^小的規(guī)模和較低的推理計(jì)算成本，顯著提高了 ReDrafter 在有限硬件條件下的整體性能。
▲圖源：阿里云開(kāi)發(fā)者社區(qū)
蘋果的基準(zhǔn)測(cè)試結(jié)果顯示，在 NVIDIA H100 GPU 上對(duì)數(shù)十億參數(shù)的生產(chǎn)模型使用集成了 ReDrafter 的 TensorRT-LLM 時(shí)，其貪心解碼（Greedy Decoding）每秒生成的 tokens 數(shù)量提高了 2.7 倍。
此外，在蘋果自家的 M2 Ultra Metal GPU 上，ReDrafter 也能實(shí)現(xiàn) 2.3 倍的推理速度提升。蘋果的研究人員表示「LLM 越來(lái)越多地用于驅(qū)動(dòng)生產(chǎn)應(yīng)用程序，提高推理效率既可以影響計(jì)算成本，也可以降低用戶端延遲」。
▲圖源：Apple
值得一提的是，在保持輸出質(zhì)量的同時(shí)，ReDrafter 減少了對(duì) GPU 資源的需求，這使得 LLM 在資源受限的環(huán)境中也能高效地運(yùn)行，為 LLM 在各種硬件平臺(tái)上的使用提供了新的可能性。
蘋果目前已經(jīng)在 GitHub 上開(kāi)源了這項(xiàng)技術(shù)，未來(lái)從中受益的公司將很可能不止英偉達(dá)一家。
