在端側設備上處理長文本常常面臨計算和內存瓶頸。

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