今天小編分享的科學經驗：标點符号成大模型訓練神器！KV緩存狂減一半，可處理400萬Tokens長序列，來自華為港大等，歡迎閱讀。
文字中貌似不起眼的标點符号，竟然可以顯著加速大模型的訓練和推理過程？
來自華為、港大、KAUST 和馬普所的研究者，就提出了一種新的自然語言建模視角——SepLLM。
起因是團隊發現某些看似無意義的分隔符，在注意力得分中占據了不成比例的重要地位。
于是，SepLLM 通過将一段文本中的信息壓縮進帳隔符（比如逗号，句号等）中，真的實現了加速效果，并且可以讓 KV 緩存減少一半。
自注意力機制的平方級復雜度，給計算存儲需求和訓練推理速度帶來了不小的挑戰。
為了降低推理的復雜度，大量節約 KV Cache 的稀疏化方法被提出。
然而這些方法大多是基于用戶的問題或者提示來篩選有用的 KV Cache。
這使得如果用戶再提出一個新的問題，模型回答的精度可能下降，因為包含答案信息的 KV 已經在上一次壓縮過程中被抛棄。
除此之外，免訓練方法通常無法相應地從頭訓練或者後訓練，導致了訓練和推理的流程差異性。
更重要的是現在主流的稀疏注意力改進方法，本質上更多是一種針對 KV Cache 存儲與計算的稀疏化管理，而不是對自然語言的自然且高效的建模。
用分隔符實現自然語言高效建模
SepLLM 通過将一段文本中的信息壓縮進帳隔符（比如逗号，句号等）中，顯著加速了大型語言模型的訓練和推理過程。
這一發現基于一個新穎且關鍵的模式：某些看似無意義的分隔符，在注意力得分中占據了不成比例的重要地位。
如下圖所示，注意力可視化顯示出一定的稀疏性，并且在分隔符處注意力明顯更大。
由此，可以将這些自然語言中分隔符所自然分割的語義段的信息有效地壓縮進帳隔符中，其他 tokens 直接丢棄，而不會造成信息損失。
除此之外，一般一個分割符所分割的語段的長度是有限且相對均衡的，因此用分割此語段的分隔符去濃縮語段信息，可以避免類似 RNN 當序列過長時而出現遺忘的問題。
因為這種基于分割符的語言建模視角反映了自然語言的自然而内在的稀疏性，而不是人為用類似 block/cluster 等概念預先定義的稀疏性，作者認為 SepLLM 可以作為大語言模型的原生稀疏注意力機制和原生基線模型。
具體來說，SepLLM 的基礎設計包含下列三種 tokens：
初始 tokens：使用稀疏注意力機制時，保留初始 tokens 可避免生成 tokens 的困惑度（ppl）顯著增加。
分隔符 tokens：看似 " 無意義 " 的分隔符 tokens 在給定輸入上下文中比有語義意義的 tokens 獲得更高的注意力分數。因此假設這些分隔符可壓縮其分割的文本片段信息，在免訓練（training-free）的場景中，基于此策略能在許多任務上取得與原始模型相似的結果；
相鄰 tokens：由于語言任務通常具有局部依賴性，相鄰 tokens 有助于形成局部平滑和連貫的上下文，所以在模型中考慮相鄰 tokens。
在預訓練或者後訓練的過程中，強迫模型當前的 token 只能看到前文每個片段中代表該片段的分隔符，使片段信息被強制濃縮到分隔符中。
實際上，每個分隔符（逗号、句号、分号、問号等）都是具備其特有的語義的，它們是對其分割段落的最原生和最細粒度的收尾與總結。
訓練階段，不需要将輸入上下文中所有 tokens 對應的 Query 向量與所有 Key 向量相乘，只需乘以掩碼矩陣中突出顯示元素對應的 Key 向量；
生成階段對 KV 緩存的管理較為直觀，只保留初始、分隔符和相鄰 tokens 的 KV Cache。
研究者還針對 Streaming 場景還提出了定制的設計，包括同時維護的四個專用緩存塊（初始緩存、分隔符緩存、過去視窗緩存和局部視窗緩存）及其功能，定義了四個緩存的運行時使用量和相鄰 tokens 數量的相關變量，并詳細說明了緩存系統的預設超參數。
在 Streaming 序列生成過程中，SepLLM 會按照一定規則填充和管理這些緩存，當緩存達到一定條件時會觸發壓縮操作。
算力緩存消耗均減少，推理速度也更快了
作者分析了 KV Cache 的平均使用情況，結果，SepLLM 在免訓練、預訓練和後訓練場景中都展現出了卓越的效率，首先進行一個簡單總結：
訓推效率提升：SepLLM 在免訓練、從頭預訓練和後訓練中都展現出了卓越的效率。特别是在使用 Llama-3-8B 模型時，SepLLM 在 GSM8K 和 MMLU 基準測試中減少了超過 50% 的 KV 緩存，同時保持了相當的性能表現。
無限長的流式處理能力：在無限長輸入的流式的場景中，SepLLM 能夠有效處理高達 400 萬甚至更多 tokens 的序列，同時保持一致的語言建模能力。
廣泛的實驗驗證與理論分析：通過在多種任務，各種基礎模型（Llama，Falcon， GPTNeoX 等）和多種數據集上的廣泛實驗，SepLLM 證明了其在不同設定下的有效性，包括免訓練、預訓練和後訓練。除此之外，作者還提供了對 SepLLM 架構通用近似（Universal Approximation）的詳細理論分析。
接下來看一下具體的實驗數據。
KV 緩存減少 50%
基于 Llama-3-8B 模型，SepLLM 實現了超過 50% 的 KV 緩存減少，推理開銷 / 顯存壓力大大降低，同時下遊任務的性能幾乎沒有損失。
SepLLM 的數學邏輯推理能力（GSM8K）/ 綜合知識面廣度（MMLU）在免訓練的場景下即可達到和 Llama-3-8B 幾乎一樣的性能。
基于 Pythia 模型的更多下遊任務上的結果，也驗證了 SepLLM 的優秀的計算和存儲效率與卓越的推理精度。
支持 400 萬 +Tokens 流式長序列生成
同時，SepLLM 可以輕松處理 400 萬 +Tokens 以上的超長流式（streaming）序列生成。
推理速度更快，困惑度更低
并且由于 SepLLM 優化了推理過程，生成速度更快，同時語言模型的困惑度也更低了，運行時的平均 KV Cache 同樣有所減小。
訓練 FLOPs 更低，速度 / 吞吐率更大
除了推理，訓練過程也用更低的 FLOPs 消耗，實現了更大的速度和吞吐率。
預訓練中，達到相同 Loss 的時間縮短 1.26 倍，并且達到 1.53 倍的訓練吞吐率和訓練加速比。
後訓練中，SepLLM 也可以在較短時間内通過後訓練恢復到原始 Full Attention 的訓練 loss，為基于大模型的高效後訓練提供了可能。
适配不同 backbone 模型架構
同時，SepLLM 可以适配各種 backbone 模型架構。
其中包括比如 Llama、Pythia、GPTNeoX、GPT2 以及 Falcon 等等。
對于這些架構，SepLLM 均能實現更低的平均運行時 KV Cache、更短的推理時間，以及更低的困惑度。
各種參數量模型均适配
SepLLM 還可以适配各種大小的模型。
從 Pythia-160M 到 Pythia-1.4B、6.9B，Llama3-8B，Falcon-40B 等等，SepLLM 均能實現更低的平均運行時 KV Cache、更短的推理時間和更低的困惑度。
最近，DeepSeek 的 NSA 與月之暗面的 MoBA 讓稀疏注意力機制受到了較大的關注，相較于上述工作采用固定 token 數來劃分壓縮區間，SepLLM 根據原生語義來劃分動态數量的 token 數。
研究者也針對靜态和動态 token 數壓縮做了讨論，在免訓練場景中，基于 SepLLM 的動态壓縮能在下遊任務中達到更好的準确率。
目前 SepLLM 的代碼庫已經公開，支持高效的多節點分布式訓練，并采用了加速注意力機制的模塊 Sep-Attention。
此外，它還支持多種現有的 Fusion Operators，如 fused rope 和 fused layer norm，以加速訓練過程。
項目地址：
https://sepllm.github.io/
論文地址：
https://arxiv.org/abs/2412.12094
代碼：
https://github.com/HKUDS/SepLLM
— 完 —
學術投稿請于工作日發郵件到：
标題注明【投稿】，告訴我們：
你是誰，從哪來，投稿内容
附上論文 / 項目主頁鏈接，以及聯系方式哦
我們會（盡量）及時回復你
一鍵關注 點亮星标
科技前沿進展每日見
一鍵三連「點贊」「轉發」「小心心」
歡迎在評論區留下你的想法！
