英偉達提出全新Star Attention，10倍加速LLM推理！登頂Hugging Face論文榜
當下的手機及AIPC中都會安裝本地大模型，然而上下文長度增加，推理時的計算成本也會顯著增長。最明顯的一個后果就是，用戶輸入問題后需要等待很久才能看到結果。
為此，已有多種優化方案提出，例如Flash Attention，而11月26日英偉達提出的Star Attention機制，可用于提升Transformer模型在處理長序列時的效率和準確性。
值得一提的是，這篇文章受到了廣泛的關注，登頂Hugging Face每日論文榜首。
論文地址：https://arxiv.org/abs/2411.17116
Star Attention如何降低推理成本
在了解Star Attention如何改進大模型推理前，讓我們先看看當前大模型的推理過程涉及的兩個步驟：
1）prompt編碼，即模型處理輸入并在緩存中存儲KV（鍵值）向量；
2）token生成，即模型關注KV緩存并自回歸生成新令牌，同時用新 的KV向量更新緩存。
在許多長上下文任務中，輸入由一個長上下文后跟一個短查詢和一個短答案組成。當大模型的上下文變得越來越長之后，回答查詢所需的信息通常局限在上下文的小部分內，意味著上下文只需關注附近的token，而查詢token需要關注所有之前上下文涉及的內容。
Star Attention下的兩階段推理
系統中所有設備被分組為多個主機（host），其中一個主機被標記 為「查詢」主機。輸入序列分為兩個階段處理。
階段一：上下文編碼
輸入的上下文部分被分割成較小的塊，并分配到各個主機。除了第一個塊之外，所有塊的前面都加上一個初始塊，稱為「錨點」塊（anchor block）。每個主機處理其分配的塊，并存儲非錨點部分的KV緩存。
階段二：查詢編碼和token生成
輸入查詢被廣播到所有主機，在每個主機中，它首先訪問在第一階段計算出的本地KV緩存。然后「查詢」主機通過聚合所有主機的softmax歸一化統計數據來計算全局注意力。這個過程對于每個生成的token都會重復。
用一個不那么嚴謹的例子來概述上面的過程：想象一場烹飪比賽（上下文token），每個廚師（主機）負責準備一道菜的一部分（塊）。
為了確保味道一致，每個廚師除了準備自己的部分，還在前面加了一點「錨點」調料（錨點塊）。每個廚師準備好自己的部分后，記住自己部分的口味（KV緩存）。
階段二的查詢編碼和token生成可視為：評委（查詢token）來品嘗菜肴，并決定下一道菜的口味（生成新token）。評委先品嘗每個廚師的部分，看看哪個部分最符合他們的口味。
最后，評委匯總所有廚師的意見，確定下一道菜的口味，并告訴廚師們。
Star Attention的性能提升
Star Attention帶來的性能提升，主要體現在以下兩個方面：
1）高達11倍的加速
在多個長上下文基準測試上，Star Attention所加持的8B Llama3的推理速度顯著提升，隨著序列長度增加，加速比從1.1x提升到2.7x。
而在參數量更大的Llama3.1-70B上，推理的加速比提升更為顯著。
與此同時，對比采用全局注意力的基準，Star Attention相對準確率的降低只在0～3%范圍內。
隨著上下文長度的增加，star attention推理的準確性相比全局注意力幾乎相同，但推理計算成本顯著下降
在更長的上下文尺度（128K）中，上下文編碼過程中不同塊的大小，也會影響推理的準確性和速度。塊尺寸越大，Star Attention 的準確性越高。
在 RULER 基準測試上，不同塊大小對Star Attention準確性的影響，塊大小范圍從4K到32K，適用于序列長度為128K的Llama-3.1-8B instruct 模型
用于評估的RULER，包含了13個任務，分為4個領域：大海撈針 （檢索）、多跳追蹤、聚合和問答，
不同任務中，全局注意力和Star Attention的準確性差異對比
而在上下文長度更大，達到1048K時，Star Attention的推理準確性依舊保持在原基準90%，推理加速比達到了10.8×～16.9×。
而在更大的Llama3.1-70B中，Star Attention能實現更大的加速比，同時保持相似水平的準確率下降。
由于其運行機制不涉及具體模型，Star Attention可以無縫集成到大多數通過全局注意力訓練的基于Transformer的LLMs中，無需額外的模型微調。
由于減少了推理的計算成本，Star Attention顯著減少了內存需求，使得在本地設備（如手機，筆記本中）用LLM處理更長的序列成為可能。
實驗發現，將塊大小設置為總序列長度的約四分之一，可以在精度和速度之間取得最佳平衡。而用戶也可以根據需求調整塊大小，以在計算效率和精度之間進行權衡。
結論
未來的研究，會嘗試將Star Attention擴展到更長的序列（最長可達1M）和更大的模型，并希望能觀察到甚至更的加速，同時保持相似水平的準確率。同時專注于優化「錨塊」機制，并在更復雜的長上下文任務上提高性能，以增強Star Attention的可擴展性和穩健性。
總的來看，對于想要開發部署本地大模型的廠商，Star Attention是一項不容錯過的技術。使用Star Attention后，本地LLM能夠更快地回復用戶，還可在有限的內存中兼容更長的上下文序列，從而在RAG任務中閱讀更長的文本。
而對于云端大模型的提供商，Star Attention能夠在幾乎不影響用戶體現的前提下，顯著提升推理成本，實現「降本增效」，同時減少能源消費（碳足跡）。
通過在多個主機間分配上下文處理，Star Attention使上下文長度能夠隨主機數量線性擴展。
