大模型如今已具有越來越長(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)注，登頂 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 向量更新緩存。
在許多長(zhǎng)上下文任務(wù)中，輸入由一個(gè)長(zhǎng)上下文后跟一個(gè)短查詢和一個(gè)短答案組成。當(dāng)大模型的上下文變得越來越長(zhǎ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è)過程對(duì)于每個(gè)生成的 token 都會(huì)重復(fù)。
用一個(gè)不那么嚴(yán)謹(jǐn)?shù)睦觼砀攀錾厦娴倪^程：想象一場(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）來品嘗菜肴，并決定下一道菜的口味（生成新 token）。評(píng)委先品嘗每個(gè)廚師的部分，看看哪個(gè)部分最符合他們的口味。
最后，評(píng)委匯總所有廚師的意見，確定下一道菜的口味，并告訴廚師們。
Star Attention 的性能提升
Star Attention 帶來的性能提升，主要體現(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 可以無縫集成到大多數(shù)通過全局注意力訓(xùn)練的基于 Transformer 的 LLMs 中，無需額外的模型微調(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é)論
未來的研究，會(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)健性。
總的來看，對(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
本文來自微信公眾號(hào)：新智元（ID：AI_era）
廣告聲明：文內(nèi)含有的對(duì)外跳轉(zhuǎn)鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節(jié)省甄選時(shí)間，結(jié)果僅供參考，IT之家所有文章均包含本聲明。