新智元報(bào)道
編輯：peter東 喬楊
【新智元導(dǎo)讀】大模型如今已具有越來(lái)越長(zhǎng)的上下文，而與之相伴的是推理成本的上升。英偉達(dá)最新提出的Star Attention，能夠在不損失精度的同時(shí)，顯著減少推理計(jì)算量，從而助力邊緣計(jì)算。
當(dāng)下的手機(jī)及AIPC中都會(huì)安裝本地大模型，然而上下文長(zhǎng)度增加，推理時(shí)的計(jì)算成本也會(huì)顯著增長(zhǎng)。最明顯的一個(gè)后果就是，用戶輸入問題后需要等待很久才能看到結(jié)果。
為此，已有多種優(yōu)化方案提出，例如Flash Attention，而11月26日英偉達(dá)提出的Star Attention機(jī)制，可用于提升Transformer模型在處理長(zhǎng)序列時(shí)的效率和準(zhǔn)確性。
值得一提的是，這篇文章受到了廣泛的關(guān)注，登頂H ug ging F ace每日論文榜首。
論文地址：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向量更新緩存。
在許多長(zhǎng)上下文任務(wù)中，輸入由一個(gè)長(zhǎng)上下文后跟一個(gè)短查詢和一個(gè)短答案組成。當(dāng)大模型的上下文變得越來(lái)越長(zhǎng)之后，回答查詢所需的信息通常局限在上下文的小部分內(nèi)，意味著上下文只需關(guān)注附近的token，而查詢token需要關(guān)注所有之前上下文涉及的內(nèi)容。
系統(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ù)來(lái)計(jì)算全局注意力。這個(gè)過程對(duì)于每個(gè)生成的token都會(huì)重復(fù)。
用一個(gè)不那么嚴(yán)謹(jǐn)?shù)睦觼?lái)概述上面的過程：想象一場(chǎng)烹飪比賽（上下文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）來(lái)品嘗菜肴，并決定下一道菜的口味（生成新token）。評(píng)委先品嘗每個(gè)廚師的部分，看看哪個(gè)部分最符合他們的口味。
最后，評(píng)委匯總所有廚師的意見，確定下一道菜的口味，并告訴廚師們。
Star Attention的性能提升
Star Attention帶來(lái)的性能提升，主要體現(xiàn)在以下兩個(gè)方面：
1）高達(dá)11倍的加速
在多個(gè)長(zhǎng)上下文基準(zhǔn)測(cè)試上，Star Attention所加持的8B Llama3的推理速度顯著提升，隨著序列長(zhǎng)度增加，加速比從1.1x提升到2.7x。
而在參數(shù)量更大的Llama3.1-70B上，推理的加速比提升更為顯著。
與此同時(shí)，對(duì)比采用全局注意力的基準(zhǔn)，Star Attention相對(duì)準(zhǔn)確率的降低只在0～3%范圍內(nèi)。
隨著上下文長(zhǎng)度的增加，star attention推理的準(zhǔn)確性相比全局注意力幾乎相同，但推理計(jì)算成本顯著下降
在更長(zhǎng)的上下文尺度（128K）中，上下文編碼過程中不同塊的大小，也會(huì)影響推理的準(zhǔn)確性和速度。塊尺寸越大，Star Attention 的準(zhǔn)確性越高。
在 RULER 基準(zhǔn)測(cè)試上，不同塊大小對(duì)Star Attention準(zhǔn)確性的影響，塊大小范圍從4K到32K，適用于序列長(zhǎng)度為128K的Llama-3.1-8B instruct 模型
用于評(píng)估的RULER，包含了13個(gè)任務(wù)，分為4個(gè)領(lǐng)域：大海撈針 （檢索）、多跳追蹤、聚合和問答，
不同任務(wù)中，全局注意力和Star Attention的準(zhǔn)確性差異對(duì)比
而在上下文長(zhǎng)度更大，達(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可以無(wú)縫集成到大多數(shù)通過全局注意力訓(xùn)練的基于Transformer的LLMs中，無(wú)需額外的模型微調(diào)。
由于減少了推理的計(jì)算成本，Star Attention顯著減少了內(nèi)存需求，使得在本地設(shè)備（如手機(jī)，筆記本中）用LLM處理更長(zhǎng)的序列成為可能。
實(shí)驗(yàn)發(fā)現(xiàn)，將塊大小設(shè)置為總序列長(zhǎng)度的約四分之一，可以在精度和速度之間取得最佳平衡。而用戶也可以根據(jù)需求調(diào)整塊大小，以在計(jì)算效率和精度之間進(jìn)行權(quán)衡。
結(jié)論
未來(lái)的研究，會(huì)嘗試將Star Attention擴(kuò)展到更長(zhǎng)的序列（最長(zhǎng)可達(dá)1M）和更大的模型，并希望能觀察到甚至更的加速，同時(shí)保持相似水平的準(zhǔn)確率。同時(shí)專注于優(yōu)化「錨塊」機(jī)制，并在更復(fù)雜的長(zhǎng)上下文任務(wù)上提高性能，以增強(qiáng)Star Attention的可擴(kuò)展性和穩(wěn)健性。
總的來(lái)看，對(duì)于想要開發(fā)部署本地大模型的廠商，Star Attention是一項(xiàng)不容錯(cuò)過的技術(shù)。使用Star Attention后，本地LLM能夠更快地回復(fù)用戶，還可在有限的內(nèi)存中兼容更長(zhǎng)的上下文序列，從而在RAG任務(wù)中閱讀更長(zhǎng)的文本。
而對(duì)于云端大模型的提供商，Star Attention能夠在幾乎不影響用戶體現(xiàn)的前提下，顯著提升推理成本，實(shí)現(xiàn)「降本增效」，同時(shí)減少能源消費(fèi)（碳足跡）。
通過在多個(gè)主機(jī)間分配上下文處理，Star Attention使上下文長(zhǎng)度能夠隨主機(jī)數(shù)量線性擴(kuò)展。
參考資料：
https://arxiv.org/abs/2411.17116
熱門跟貼