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

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

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

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

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

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

然而這些方法大多是基于用戶的問題或者提示來篩選有用的KV Cache。

這使得如果用戶再提出一個新的問題，模型回答的精度可能下降，因為包含答案信息的KV已經(jīng)在上一次壓縮過程中被拋棄。

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

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

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

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

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

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

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

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

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

具體來說，SepLLM的基礎(chǔ)設(shè)計包含下列三種tokens：

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

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

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

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

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

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

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

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

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

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

接下來看一下具體的實驗數(shù)據(jù)。

KV緩存減少50%

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

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

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

支持400萬+Tokens流式長序列生成

同時，SepLLM可以輕松處理400萬+Tokens以上的超長流式（streaming）序列生成。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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