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