文字中貌似不起眼的標(biāo)點(diǎn)符號(hào)，竟然可以顯著加速大模型的訓(xùn)練和推理過(guò)程？

來(lái)自華為、港大、KAUST和馬普所的研究者，就提出了一種新的自然語(yǔ)言建模視角——SepLLM。

起因是團(tuán)隊(duì)發(fā)現(xiàn)某些看似無(wú)意義的分隔符，在注意力得分中占據(jù)了不成比例的重要地位。

于是，SepLLM通過(guò)將一段文本中的信息壓縮進(jìn)分隔符（比如逗號(hào)，句號(hào)等）中，真的實(shí)現(xiàn)了加速效果，并且可以讓KV緩存減少一半。

自注意力機(jī)制的平方級(jí)復(fù)雜度，給計(jì)算存儲(chǔ)需求和訓(xùn)練推理速度帶來(lái)了不小的挑戰(zhàn)。

為了降低推理的復(fù)雜度，大量節(jié)約KV Cache的稀疏化方法被提出。

然而這些方法大多是基于用戶(hù)的問(wèn)題或者提示來(lái)篩選有用的KV Cache。

這使得如果用戶(hù)再提出一個(gè)新的問(wèn)題，模型回答的精度可能下降，因?yàn)榘鸢感畔⒌腒V已經(jīng)在上一次壓縮過(guò)程中被拋棄。

除此之外，免訓(xùn)練方法通常無(wú)法相應(yīng)地從頭訓(xùn)練或者后訓(xùn)練，導(dǎo)致了訓(xùn)練和推理的流程差異性。

更重要的是現(xiàn)在主流的稀疏注意力改進(jìn)方法，本質(zhì)上更多是一種針對(duì)KV Cache存儲(chǔ)與計(jì)算的稀疏化管理，而不是對(duì)自然語(yǔ)言的自然且高效的建模。

用分隔符實(shí)現(xiàn)自然語(yǔ)言高效建模

SepLLM通過(guò)將一段文本中的信息壓縮進(jìn)分隔符（比如逗號(hào)，句號(hào)等）中，顯著加速了大型語(yǔ)言模型的訓(xùn)練和推理過(guò)程。

這一發(fā)現(xiàn)基于一個(gè)新穎且關(guān)鍵的模式：某些看似無(wú)意義的分隔符，在注意力得分中占據(jù)了不成比例的重要地位。

如下圖所示，注意力可視化顯示出一定的稀疏性，并且在分隔符處注意力明顯更大。

由此，可以將這些自然語(yǔ)言中分隔符所自然分割的語(yǔ)義段的信息有效地壓縮進(jìn)分隔符中，其他tokens直接丟棄，而不會(huì)造成信息損失。

除此之外，一般一個(gè)分割符所分割的語(yǔ)段的長(zhǎng)度是有限且相對(duì)均衡的，因此用分割此語(yǔ)段的分隔符去濃縮語(yǔ)段信息，可以避免類(lèi)似RNN當(dāng)序列過(guò)長(zhǎng)時(shí)而出現(xiàn)遺忘的問(wèn)題。

因?yàn)檫@種基于分割符的語(yǔ)言建模視角反映了自然語(yǔ)言的自然而內(nèi)在的稀疏性，而不是人為用類(lèi)似block/cluster等概念預(yù)先定義的稀疏性，作者認(rèn)為SepLLM可以作為大語(yǔ)言模型的原生稀疏注意力機(jī)制和原生基線(xiàn)模型。

具體來(lái)說(shuō)，SepLLM的基礎(chǔ)設(shè)計(jì)包含下列三種tokens：

• 初始tokens：使用稀疏注意力機(jī)制時(shí)，保留初始tokens可避免生成tokens的困惑度（ppl）顯著增加。
• 分隔符tokens：看似“無(wú)意義”的分隔符tokens在給定輸入上下文中比有語(yǔ)義意義的tokens獲得更高的注意力分?jǐn)?shù)。因此假設(shè)這些分隔符可壓縮其分割的文本片段信息，在免訓(xùn)練（training-free）的場(chǎng)景中，基于此策略能在許多任務(wù)上取得與原始模型相似的結(jié)果；
• 相鄰tokens：由于語(yǔ)言任務(wù)通常具有局部依賴(lài)性，相鄰tokens有助于形成局部平滑和連貫的上下文，所以在模型中考慮相鄰tokens。

在預(yù)訓(xùn)練或者后訓(xùn)練的過(guò)程中，強(qiáng)迫模型當(dāng)前的token只能看到前文每個(gè)片段中代表該片段的分隔符，使片段信息被強(qiáng)制濃縮到分隔符中。

實(shí)際上，每個(gè)分隔符（逗號(hào)、句號(hào)、分號(hào)、問(wèn)號(hào)等）都是具備其特有的語(yǔ)義的，它們是對(duì)其分割段落的最原生和最細(xì)粒度的收尾與總結(jié)。

訓(xùn)練階段，不需要將輸入上下文中所有tokens對(duì)應(yīng)的Query向量與所有Key向量相乘，只需乘以掩碼矩陣中突出顯示元素對(duì)應(yīng)的Key向量；

生成階段對(duì)KV緩存的管理較為直觀，只保留初始、分隔符和相鄰tokens的KV Cache。

研究者還針對(duì)Streaming場(chǎng)景還提出了定制的設(shè)計(jì)，包括同時(shí)維護(hù)的四個(gè)專(zhuān)用緩存塊（初始緩存、分隔符緩存、過(guò)去窗口緩存和局部窗口緩存）及其功能，定義了四個(gè)緩存的運(yùn)行時(shí)使用量和相鄰tokens數(shù)量的相關(guān)變量，并詳細(xì)說(shuō)明了緩存系統(tǒng)的預(yù)設(shè)超參數(shù)。

在Streaming序列生成過(guò)程中，SepLLM會(huì)按照一定規(guī)則填充和管理這些緩存，當(dāng)緩存達(dá)到一定條件時(shí)會(huì)觸發(fā)壓縮操作。

算力緩存消耗均減少，推理速度也更快了

作者分析了KV Cache的平均使用情況，結(jié)果，SepLLM在免訓(xùn)練、預(yù)訓(xùn)練和后訓(xùn)練場(chǎng)景中都展現(xiàn)出了卓越的效率，首先進(jìn)行一個(gè)簡(jiǎn)單總結(jié)：

• 訓(xùn)推效率提升：SepLLM在免訓(xùn)練、從頭預(yù)訓(xùn)練和后訓(xùn)練中都展現(xiàn)出了卓越的效率。特別是在使用Llama-3-8B模型時(shí)，SepLLM在GSM8K和MMLU基準(zhǔn)測(cè)試中減少了超過(guò)50%的KV緩存，同時(shí)保持了相當(dāng)?shù)男阅鼙憩F(xiàn)。
• 無(wú)限長(zhǎng)的流式處理能力：在無(wú)限長(zhǎng)輸入的流式的場(chǎng)景中，SepLLM能夠有效處理高達(dá)400萬(wàn)甚至更多tokens的序列，同時(shí)保持一致的語(yǔ)言建模能力。
• 廣泛的實(shí)驗(yàn)驗(yàn)證與理論分析：通過(guò)在多種任務(wù)，各種基礎(chǔ)模型（Llama，F(xiàn)alcon， GPTNeoX等）和多種數(shù)據(jù)集上的廣泛實(shí)驗(yàn)，SepLLM證明了其在不同設(shè)置下的有效性，包括免訓(xùn)練、預(yù)訓(xùn)練和后訓(xùn)練。除此之外，作者還提供了對(duì)SepLLM架構(gòu)通用近似（Universal Approximation）的詳細(xì)理論分析。

接下來(lái)看一下具體的實(shí)驗(yàn)數(shù)據(jù)。

KV緩存減少50%

基于Llama-3-8B模型，SepLLM實(shí)現(xiàn)了超過(guò)50%的KV緩存減少，推理開(kāi)銷(xiāo)/顯存壓力大大降低，同時(shí)下游任務(wù)的性能幾乎沒(méi)有損失。

SepLLM的數(shù)學(xué)邏輯推理能力（GSM8K）/綜合知識(shí)面廣度（MMLU）在免訓(xùn)練的場(chǎng)景下即可達(dá)到和Llama-3-8B幾乎一樣的性能。

基于Pythia模型的更多下游任務(wù)上的結(jié)果，也驗(yàn)證了SepLLM的優(yōu)秀的計(jì)算和存儲(chǔ)效率與卓越的推理精度。

支持400萬(wàn)+Tokens流式長(zhǎng)序列生成

同時(shí)，SepLLM可以輕松處理400萬(wàn)+Tokens以上的超長(zhǎng)流式（streaming）序列生成。

推理速度更快，困惑度更低

并且由于SepLLM優(yōu)化了推理過(guò)程，生成速度更快，同時(shí)語(yǔ)言模型的困惑度也更低了，運(yùn)行時(shí)的平均KV Cache同樣有所減小。

訓(xùn)練FLOPs更低，速度/吞吐率更大

除了推理，訓(xùn)練過(guò)程也用更低的FLOPs消耗，實(shí)現(xiàn)了更大的速度和吞吐率。

預(yù)訓(xùn)練中，達(dá)到相同Loss的時(shí)間縮短1.26倍，并且達(dá)到1.53倍的訓(xùn)練吞吐率和訓(xùn)練加速比。

后訓(xùn)練中，SepLLM也可以在較短時(shí)間內(nèi)通過(guò)后訓(xùn)練恢復(fù)到原始Full Attention的訓(xùn)練loss，為基于大模型的高效后訓(xùn)練提供了可能。

適配不同backbone模型架構(gòu)

同時(shí)，SepLLM可以適配各種backbone模型架構(gòu)。

其中包括比如Llama、Pythia、GPTNeoX、GPT2以及Falcon等等。

對(duì)于這些架構(gòu)，SepLLM均能實(shí)現(xiàn)更低的平均運(yùn)行時(shí)KV Cache、更短的推理時(shí)間，以及更低的困惑度。

各種參數(shù)量模型均適配

SepLLM還可以適配各種大小的模型。

從Pythia-160M到Pythia-1.4B、6.9B，Llama3-8B，F(xiàn)alcon-40B等等，SepLLM均能實(shí)現(xiàn)更低的平均運(yùn)行時(shí)KV Cache、更短的推理時(shí)間和更低的困惑度。

最近，DeepSeek的NSA與月之暗面的MoBA讓稀疏注意力機(jī)制受到了較大的關(guān)注，相較于上述工作采用固定token數(shù)來(lái)劃分壓縮區(qū)間，SepLLM根據(jù)原生語(yǔ)義來(lái)劃分動(dòng)態(tài)數(shù)量的token數(shù)。

研究者也針對(duì)靜態(tài)和動(dòng)態(tài)token數(shù)壓縮做了討論，在免訓(xùn)練場(chǎng)景中，基于SepLLM的動(dòng)態(tài)壓縮能在下游任務(wù)中達(dá)到更好的準(zhǔn)確率。

目前SepLLM的代碼庫(kù)已經(jīng)公開(kāi)，支持高效的多節(jié)點(diǎn)分布式訓(xùn)練，并采用了加速注意力機(jī)制的模塊Sep-Attention。

此外，它還支持多種現(xiàn)有的Fusion Operators，如fused rope和fused layer norm，以加速訓(xùn)練過(guò)程。

項(xiàng)目地址：https://sepllm.github.io/
論文地址：https://arxiv.org/abs/2412.12094
代碼：https://github.com/HKUDS/SepLLM