標點符號成大模型訓練神器！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）的詳細理論分析。