vivo端側大模型團隊 投稿量子位 | 公眾號 QbitAI
在端側裝置上處理長文字常常面臨計算和記憶體瓶頸。
vivo AI研究院推出的EdgeInfinite演算法專為端側裝置設計，讓裝置處理超長文字時更加高效流暢，該方法能夠在不到10GB GPU記憶體的裝置上處理長達128K tokens的輸入。
該研究成果已中稿ACL 2025。
以下是更多詳細內容介紹。
EdgeInfinite:解決端側裝置長文字處理的高效演算法
端側LLM在實際應用中會遇到很多長文字輸入的場景（例如通話摘要和個人文件總結），但由於端側裝置的資源限制，現有的LLM在部署到端側後都無法處理很長的上下文。
這是由於現在LLM都是基於Transformer架構，其計算耗時和記憶體佔用會隨著輸入長度增加而顯著增長，尤其當需要將Transformer類模型部署到端側裝置上時，面臨的挑戰會愈發突出。
為了解決這類問題，vivo AI研究院提出了一種用於端側裝置的長文字演算法——EdgeInfinite，該演算法透過一個可訓練的門控記憶模組將記憶壓縮演算法整合到了Transformer架構中。
本方法與原生的Transformer架構完全相容，在訓練時只需要微調一小部分引數，就可以在長文字任務上取得不錯的效果，同時保持了相對高效的推理效能，非常適合在端側裝置上高效地處理長文字任務。
EdgeInfinite架構解析
EdgeInfinite的架構如上圖所示，主要包括三個核心部分：結合ROPE的分塊注意力模組、記憶壓縮與解壓縮模組和自適應的門控記憶模組。
1. 結合ROPE的分塊注意力
EdgeInfinite會把輸入的文字序列，按照一定長度切成小的片段。對每個片段分別計算 Q、K、V 值。同時，對每個片段會加入位置編碼ROPE，讓模型更好地理解小段內每個token之間的相對位置關係，這樣在計算注意力的時候就更準確。
2. 記憶壓縮與解壓縮
EdgeInfinite引入了記憶的壓縮與解壓縮模組，將過去KV狀態分塊並存儲為固定長度的記憶塊，並在計算attention時候解壓縮出來進行計算。由於記憶中編碼了之前片段的KV對的關聯，解壓縮使我們能夠計算當前Q狀態和過去的KV狀態之間的注意力。這個過程使得塊狀計算能夠近似原始長序列的注意力計算。
3. 自適應的門控記憶模組
EdgeInfinite透過自適應的門控記憶模組將基於記憶的注意力與基於區域性片段的注意力相結合，從而增強模型處理長距離依賴關係的能力。在訓練時，EdgeInfinite只需要對記憶門控模組進行微調。
推理策略
EdgeInfinite在推理時採用了兩個策略：
1.保留特定token的kv cache：在推理過程中的固定保留了兩種特殊token的kv cache，即sink token（序列最開始的一些token）和window token（序列末尾的一些token），這些token對於保留語義和位置資訊非常重要，需要將它們保留為未壓縮狀態來確保高質量的推理輸出。
2.長短文字任務路由機制：EdgeInfinite可以靈活的和已有的基礎模型相結合，來提升基礎模型的長文字能力，同時也不影響基礎模型的短文字能力，這是透過推理時動態的長短文字任務路由來實現的。
實驗結果
研究人員使用vivo自研的 BlueLM-3B 作為基礎模型，在 LongBench 這個包含多種長文字任務的資料集上測試 EdgeInfinite 的效能，同時將EdgeInfinite與三種KV Cache最佳化方法（SnapKV、PyramidKV和StreamingLLM）以及保留完整KV Cache的原始模型（FullKV）進行比較。
實驗結果（如上圖）顯示，EdgeInfinite在多文件問答和少樣本學習這些任務上，相比其他的方法有明顯優勢；在其中部分任務上還會優於原始的 FullKV 模型，整體模型效果還是很有競爭力的。
研究人員還做了消融實驗來驗證推理時保留特定token是很重要的。實驗結果（如上圖）表明，如果去掉sink token或者window token，推理效果會受到很大影響。
和原始的BlueLM-3B模型相比，EdgeInfinite在處理長文字輸入時，首詞出詞時間更短，佔用的記憶體也更少；即使輸入文字長度增加，EdgeInfinite 的記憶體佔用也保持在相對穩定的水平。
未來，EdgeInfinite有望在更多資源受限的裝置上廣泛應用，提升各類長文字處理任務的效率。比如在智慧語音助手、移動辦公文件處理等場景中，讓使用者獲得更流暢的體驗。
論文連結：https://arxiv.org/pdf/2503.22196
一鍵三連「點贊」「轉發」「小心心」
歡迎在評論區留下你的想法！
— 完 —
🌟 點亮星標 🌟