蘋果正在與英偉達合作，想讓AI的響應速度更快

近日，蘋果與英偉達宣布合作，旨在加速和優化大語言模型（LLM）的推理性能。

為了改善傳統自回歸 LLM 推理效率低和內存帶寬小的問題，今年早些時候，蘋果機器學習的研究人員發布並開源了一種名為「ReDrafter」（Recurrent Drafter，循環草稿模型）的推測解碼技術。

目前，ReDrafter 已經整合到英偉達的可擴展推理方案「TensorRT-LLM」當中，後者是基於「TensorRT」深度學習編譯框架的專為優化 LLM 推理而設計的開源庫，支持包括「Medusa」等推測解碼方法。

不過，由於 ReDrafter 所包含的算法使用了之前從未用過的運算符，因此英偉達方面添加了新的運算符，或者公開了現有的運算符，大大提高了 TensorRT-LLM 適應複雜模型和解碼方式的能力。

據悉，ReDrafter 推測解碼通過三個關鍵技術來加速 LLM 的推理過程：

• RNN 草稿模型

• 動態樹注意力算法

• 知識蒸餾訓練

RNN 草稿模型是 ReDrafter 的「核心」組件。它使用循環神經網路（Recurrent Neural Network），基於 LLM 的「隱藏狀態」來預測接下來可能出現的 tokens 序列，其能夠捕捉局部的時間依賴性，從而提高預測準確性。

這個模型的工作原理是：LLM 在文本生成過程中首先生成一個初始 token，然後 RNN 草稿模型利用該 token 和 LLM 的最後一層隱藏狀態作為輸入進行束搜索（Beam Search），進而生成多個候選 tokens 序列。

與傳統自回歸 LLM 每次只生成一個 token 不同，通過 RNN 草稿模型的預測輸出，ReDrafter 能夠在每個解碼步驟生成多個 tokens，大大減少了需要調用 LLM 驗證的次數，從而提高了整體的推理速度。

動態樹注意力算法（Dynamic Tree Attention）則是一種優化束搜索結果的算法。

我們已經知道，在束搜索過程中會產生多個候選序列，而這些序列往往存在共享的前綴。動態樹注意力算法會識別出這些共享前綴，並將它們從需要驗證的 tokens 中去除，從而減少 LLM 需要處理的數據量。

某些情況下，該算法能將需要驗證的 tokens 數量減少 30% 到 60%。這意味著使用動態樹注意力算法後，ReDrafter 能夠更高效地利用計算資源，進一步提高推理速度。

知識蒸餾是一種模型壓縮技術，它能夠將一個大型、複雜的模型（教師模型）的知識「蒸餾」到一個更小、更簡單的模型（學生模型）中。在 ReDrafter 中，RNN 草稿模型作為學生模型通過知識蒸餾從 LLM（教師模型）中學習。

具體來講，蒸餾訓練過程中，LLM 會給出一系列下一個可能詞的「概率分布」，開發人員會基於這個概率分布數據訓練 RNN 草稿模型，然後計算兩個模型概率分布之間的差異，並通過優化算法使這個差異最小化。

在這個過程中，RNN 草稿模型不斷學習 LLM 的概率預測模式，從而在實際應用中能夠生成與 LLM 相似的文本。

通過知識蒸餾訓練，RNN 草稿模型更好地捕捉到語言的規律和模式，從而更準確地預測 LLM 的輸出，並且因為其較小的規模和較低的推理計算成本，顯著提高了 ReDrafter 在有限硬體條件下的整體性能。

蘋果的基準測試結果顯示，在 NVIDIA H100 GPU 上對數十億參數的生產模型使用集成了 ReDrafter 的 TensorRT-LLM 時，其貪心解碼（Greedy Decoding）每秒生成的 tokens 數量提高了 2.7 倍。

此外，在蘋果自家的 M2 Ultra Metal GPU 上，ReDrafter 也能實現 2.3 倍的推理速度提升。蘋果的研究人員表示「LLM 越來越多地用於驅動生產應用程序，提高推理效率既可以影響計算成本，也可以降低用戶端延遲」。

值得一提的是，在保持輸出質量的同時，ReDrafter 減少了對 GPU 資源的需求，這使得 LLM 在資源受限的環境中也能高效地運行，為 LLM 在各種硬體平台上的使用提供了新的可能性。

蘋果目前已經在 GitHub 上開源了這項技術，未來從中受益的公司將很可能不止英偉達一家。