當(dāng)下的手機(jī)及AIPC中都會(huì)安裝本地大模型，然而上下文長度增加，推理時(shí)的計(jì)算成本也會(huì)顯著增長。最明顯的一個(gè)后果就是，用戶輸入問題后需要等待很久才能看到結(jié)果。

為此，已有多種優(yōu)化方案提出，例如Flash Attention，而11月26日英偉達(dá)提出的Star Attention機(jī)制，可用于提升Transformer模型在處理長序列時(shí)的效率和準(zhǔn)確性。

值得一提的是，這篇文章受到了廣泛的關(guān)注，登頂Hugging Face每日論文榜首。

論文地址：https://arxiv.org/abs/2411.17116

Star Attention如何降低推理成本

在了解Star Attention如何改進(jìn)大模型推理前，讓我們先看看當(dāng)前大模型的推理過程涉及的兩個(gè)步驟：

1）prompt編碼，即模型處理輸入并在緩存中存儲(chǔ)KV（鍵值）向量；

2）token生成，即模型關(guān)注KV緩存并自回歸生成新令牌，同時(shí)用新 的KV向量更新緩存。

在許多長上下文任務(wù)中，輸入由一個(gè)長上下文后跟一個(gè)短查詢和一個(gè)短答案組成。當(dāng)大模型的上下文變得越來越長之后，回答查詢所需的信息通常局限在上下文的小部分內(nèi)，意味著上下文只需關(guān)注附近的token，而查詢token需要關(guān)注所有之前上下文涉及的內(nèi)容。

Star Attention下的兩階段推理

系統(tǒng)中所有設(shè)備被分組為多個(gè)主機(jī)（host），其中一個(gè)主機(jī)被標(biāo)記 為「查詢」主機(jī)。輸入序列分為兩個(gè)階段處理。

階段一：上下文編碼

輸入的上下文部分被分割成較小的塊，并分配到各個(gè)主機(jī)。除了第一個(gè)塊之外，所有塊的前面都加上一個(gè)初始?jí)K，稱為「錨點(diǎn)」塊（anchor block）。每個(gè)主機(jī)處理其分配的塊，并存儲(chǔ)非錨點(diǎn)部分的KV緩存。

階段二：查詢編碼和token生成

輸入查詢被廣播到所有主機(jī)，在每個(gè)主機(jī)中，它首先訪問在第一階段計(jì)算出的本地KV緩存。然后「查詢」主機(jī)通過聚合所有主機(jī)的softmax歸一化統(tǒng)計(jì)數(shù)據(jù)來計(jì)算全局注意力。這個(gè)過程對于每個(gè)生成的token都會(huì)重復(fù)。

用一個(gè)不那么嚴(yán)謹(jǐn)?shù)睦觼砀攀錾厦娴倪^程：想象一場烹飪比賽（上下文token），每個(gè)廚師（主機(jī)）負(fù)責(zé)準(zhǔn)備一道菜的一部分（塊）。

為了確保味道一致，每個(gè)廚師除了準(zhǔn)備自己的部分，還在前面加了一點(diǎn)「錨點(diǎn)」調(diào)料（錨點(diǎn)塊）。每個(gè)廚師準(zhǔn)備好自己的部分后，記住自己部分的口味（KV緩存）。

階段二的查詢編碼和token生成可視為：評(píng)委（查詢token）來品嘗菜肴，并決定下一道菜的口味（生成新token）。評(píng)委先品嘗每個(gè)廚師的部分，看看哪個(gè)部分最符合他們的口味。

最后，評(píng)委匯總所有廚師的意見，確定下一道菜的口味，并告訴廚師們。

Star Attention的性能提升

Star Attention帶來的性能提升，主要體現(xiàn)在以下兩個(gè)方面：

1）高達(dá)11倍的加速

在多個(gè)長上下文基準(zhǔn)測試上，Star Attention所加持的8B Llama3的推理速度顯著提升，隨著序列長度增加，加速比從1.1x提升到2.7x。

而在參數(shù)量更大的Llama3.1-70B上，推理的加速比提升更為顯著。

與此同時(shí)，對比采用全局注意力的基準(zhǔn)，Star Attention相對準(zhǔn)確率的降低只在0～3%范圍內(nèi)。

隨著上下文長度的增加，star attention推理的準(zhǔn)確性相比全局注意力幾乎相同，但推理計(jì)算成本顯著下降

在更長的上下文尺度（128K）中，上下文編碼過程中不同塊的大小，也會(huì)影響推理的準(zhǔn)確性和速度。塊尺寸越大，Star Attention 的準(zhǔn)確性越高。

在 RULER 基準(zhǔn)測試上，不同塊大小對Star Attention準(zhǔn)確性的影響，塊大小范圍從4K到32K，適用于序列長度為128K的Llama-3.1-8B instruct 模型

用于評(píng)估的RULER，包含了13個(gè)任務(wù)，分為4個(gè)領(lǐng)域：大海撈針 （檢索）、多跳追蹤、聚合和問答，

不同任務(wù)中，全局注意力和Star Attention的準(zhǔn)確性差異對比

而在上下文長度更大，達(dá)到1048K時(shí)，Star Attention的推理準(zhǔn)確性依舊保持在原基準(zhǔn)90%，推理加速比達(dá)到了10.8×～16.9×。

而在更大的Llama3.1-70B中，Star Attention能實(shí)現(xiàn)更大的加速比，同時(shí)保持相似水平的準(zhǔn)確率下降。

由于其運(yùn)行機(jī)制不涉及具體模型，Star Attention可以無縫集成到大多數(shù)通過全局注意力訓(xùn)練的基于Transformer的LLMs中，無需額外的模型微調(diào)。

由于減少了推理的計(jì)算成本，Star Attention顯著減少了內(nèi)存需求，使得在本地設(shè)備（如手機(jī)，筆記本中）用LLM處理更長的序列成為可能。

實(shí)驗(yàn)發(fā)現(xiàn)，將塊大小設(shè)置為總序列長度的約四分之一，可以在精度和速度之間取得最佳平衡。而用戶也可以根據(jù)需求調(diào)整塊大小，以在計(jì)算效率和精度之間進(jìn)行權(quán)衡。

結(jié)論

未來的研究，會(huì)嘗試將Star Attention擴(kuò)展到更長的序列（最長可達(dá)1M）和更大的模型，并希望能觀察到甚至更的加速，同時(shí)保持相似水平的準(zhǔn)確率。同時(shí)專注于優(yōu)化「錨塊」機(jī)制，并在更復(fù)雜的長上下文任務(wù)上提高性能，以增強(qiáng)Star Attention的可擴(kuò)展性和穩(wěn)健性。

總的來看，對于想要開發(fā)部署本地大模型的廠商，Star Attention是一項(xiàng)不容錯(cuò)過的技術(shù)。使用Star Attention后，本地LLM能夠更快地回復(fù)用戶，還可在有限的內(nèi)存中兼容更長的上下文序列，從而在RAG任務(wù)中閱讀更長的文本。

而對于云端大模型的提供商，Star Attention能夠在幾乎不影響用戶體現(xiàn)的前提下，顯著提升推理成本，實(shí)現(xiàn)「降本增效」，同時(shí)減少能源消費(fèi)（碳足跡）。

通過在多個(gè)主機(jī)間分配上下文處理，Star Attention使上下文長度能夠隨主機(jī)數(shù)量線性擴(kuò)展。