今天小編分享的科技經驗：蘋果正在與英偉達合作，想讓 AI 的響應速度更快，歡迎閱讀。
近日，蘋果與英偉達宣布合作，旨在加速和優化大語言模型（LLM）的推理性能。
為了改善傳統自回歸 LLM 推理效率低和内存帶寬小的問題，今年早些時候，蘋果機器學習的研究人員發布并開源了一種名為「ReDrafter」（Recurrent Drafter，循環草稿模型）的推測解碼技術。
▲圖源：GitHub
目前，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 驗證的次數，從而提高了整體的推理速度。
▲圖源：arXiv
動态樹注意力算法（Dynamic Tree Attention）則是一種優化束搜索結果的算法。
我們已經知道，在束搜索過程中會產生多個候選序列，而這些序列往往存在共享的前綴。動态樹注意力算法會識别出這些共享前綴，并将它們從需要驗證的 tokens 中去除，從而減少 LLM 需要處理的數據量。
某些情況下，該算法能将需要驗證的 tokens 數量減少 30% 到 60%。這意味着使用動态樹注意力算法後，ReDrafter 能夠更高效地利用計算資源，進一步提高推理速度。
▲圖源：NVIDIA
知識蒸餾是一種模型壓縮技術，它能夠将一個大型、復雜的模型（教師模型）的知識「蒸餾」到一個更小、更簡單的模型（學生模型）中。在 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 越來越多地用于驅動生產應用程式，提高推理效率既可以影響計算成本，也可以降低用戶端延遲」。
▲圖源：Apple
值得一提的是，在保持輸出質量的同時，ReDrafter 減少了對 GPU 資源的需求，這使得 LLM 在資源受限的環境中也能高效地運行，為 LLM 在各種硬體平台上的使用提供了新的可能性。
蘋果目前已經在 GitHub 上開源了這項技術，未來從中獲益的公司将很可能不止英偉達一家。
