在端側(cè)設(shè)備上處理長(zhǎng)文本常常面臨計(jì)算和內(nèi)存瓶頸。

vivo AI 研究院推出的 EdgeInfinite 算法專為端側(cè)設(shè)備設(shè)計(jì)，讓設(shè)備處理超長(zhǎng)文本時(shí)更加高效流暢，該方法能夠在不到 10GB GPU 內(nèi)存的設(shè)備上處理長(zhǎng)達(dá) 128K tokens 的輸入。

該研究成果已中稿 ACL 2025。

以下是更多詳細(xì)內(nèi)容介紹。

EdgeInfinite: 解決端側(cè)設(shè)備長(zhǎng)文本處理的高效算法

端側(cè) LLM 在實(shí)際應(yīng)用中會(huì)遇到很多長(zhǎng)文本輸入的場(chǎng)景（例如通話摘要和個(gè)人文檔總結(jié)），但由于端側(cè)設(shè)備的資源限制，現(xiàn)有的 LLM 在部署到端側(cè)后都無(wú)法處理很長(zhǎng)的上下文。

這是由于現(xiàn)在 LLM 都是基于 Transformer 架構(gòu)，其計(jì)算耗時(shí)和內(nèi)存占用會(huì)隨著輸入長(zhǎng)度增加而顯著增長(zhǎng)，尤其當(dāng)需要將 Transformer 類(lèi)模型部署到端側(cè)設(shè)備上時(shí)，面臨的挑戰(zhàn)會(huì)愈發(fā)突出。

為了解決這類(lèi)問(wèn)題，vivo AI 研究院提出了一種用于端側(cè)設(shè)備的長(zhǎng)文本算法——EdgeInfinite，該算法通過(guò)一個(gè)可訓(xùn)練的門(mén)控記憶模塊將記憶壓縮算法集成到了Transformer 架構(gòu)中。

本方法與原生的 Transformer 架構(gòu)完全兼容，在訓(xùn)練時(shí)只需要微調(diào)一小部分參數(shù)，就可以在長(zhǎng)文本任務(wù)上取得不錯(cuò)的效果，同時(shí)保持了相對(duì)高效的推理性能，非常適合在端側(cè)設(shè)備上高效地處理長(zhǎng)文本任務(wù)。

EdgeInfinite 架構(gòu)解析

EdgeInfinite 的架構(gòu)如上圖所示，主要包括三個(gè)核心部分：結(jié)合 ROPE 的分塊注意力模塊、記憶壓縮與解壓縮模塊和自適應(yīng)的門(mén)控記憶模塊。

1. 結(jié)合 ROPE 的分塊注意力

EdgeInfinite 會(huì)把輸入的文本序列，按照一定長(zhǎng)度切成小的片段。對(duì)每個(gè)片段分別計(jì)算 Q、K、V 值。同時(shí)，對(duì)每個(gè)片段會(huì)加入位置編碼 ROPE，讓模型更好地理解小段內(nèi)每個(gè) token 之間的相對(duì)位置關(guān)系，這樣在計(jì)算注意力的時(shí)候就更準(zhǔn)確。

2. 記憶壓縮與解壓縮

EdgeInfinite 引入了記憶的壓縮與解壓縮模塊，將過(guò)去 KV 狀態(tài)分塊并存儲(chǔ)為固定長(zhǎng)度的記憶塊，并在計(jì)算 attention 時(shí)候解壓縮出來(lái)進(jìn)行計(jì)算。由于記憶中編碼了之前片段的 KV 對(duì)的關(guān)聯(lián)，解壓縮使我們能夠計(jì)算當(dāng)前 Q 狀態(tài)和過(guò)去的 KV 狀態(tài)之間的注意力。這個(gè)過(guò)程使得塊狀計(jì)算能夠近似原始長(zhǎng)序列的注意力計(jì)算。

3. 自適應(yīng)的門(mén)控記憶模塊

EdgeInfinite 通過(guò)自適應(yīng)的門(mén)控記憶模塊將基于記憶的注意力與基于局部片段的注意力相結(jié)合，從而增強(qiáng)模型處理長(zhǎng)距離依賴關(guān)系的能力。在訓(xùn)練時(shí)，EdgeInfinite 只需要對(duì)記憶門(mén)控模塊進(jìn)行微調(diào)。

推理策略

EdgeInfinite 在推理時(shí)采用了兩個(gè)策略：

1.保留特定 token 的 kv cache：在推理過(guò)程中的固定保留了兩種特殊 token 的 kv cache，即sink token（序列最開(kāi)始的一些 token）和window token（序列末尾的一些 token），這些 token 對(duì)于保留語(yǔ)義和位置信息非常重要，需要將它們保留為未壓縮狀態(tài)來(lái)確保高質(zhì)量的推理輸出。

2.長(zhǎng)短文本任務(wù)路由機(jī)制：EdgeInfinite 可以靈活的和已有的基礎(chǔ)模型相結(jié)合，來(lái)提升基礎(chǔ)模型的長(zhǎng)文本能力，同時(shí)也不影響基礎(chǔ)模型的短文本能力，這是通過(guò)推理時(shí)動(dòng)態(tài)的長(zhǎng)短文本任務(wù)路由來(lái)實(shí)現(xiàn)的。

實(shí)驗(yàn)結(jié)果

研究人員使用 vivo 自研的 BlueLM-3B 作為基礎(chǔ)模型，在 LongBench 這個(gè)包含多種長(zhǎng)文本任務(wù)的數(shù)據(jù)集上測(cè)試 EdgeInfinite 的性能，同時(shí)將 EdgeInfinite 與三種 KV Cache 優(yōu)化方法（SnapKV、PyramidKV 和 StreamingLLM）以及保留完整 KV Cache 的原始模型（FullKV）進(jìn)行比較。

實(shí)驗(yàn)結(jié)果（如上圖）顯示，EdgeInfinite 在多文檔問(wèn)答和少樣本學(xué)習(xí)這些任務(wù)上，相比其他的方法有明顯優(yōu)勢(shì)；在其中部分任務(wù)上還會(huì)優(yōu)于原始的 FullKV 模型，整體模型效果還是很有競(jìng)爭(zhēng)力的。

研究人員還做了消融實(shí)驗(yàn)來(lái)驗(yàn)證推理時(shí)保留特定 token 是很重要的。實(shí)驗(yàn)結(jié)果（如上圖）表明，如果去掉 sink token 或者 window token，推理效果會(huì)受到很大影響。

和原始的 BlueLM-3B 模型相比，EdgeInfinite 在處理長(zhǎng)文本輸入時(shí)，首詞出詞時(shí)間更短，占用的內(nèi)存也更少；即使輸入文本長(zhǎng)度增加，EdgeInfinite 的內(nèi)存占用也保持在相對(duì)穩(wěn)定的水平。

未來(lái)，EdgeInfinite 有望在更多資源受限的設(shè)備上廣泛應(yīng)用，提升各類(lèi)長(zhǎng)文本處理任務(wù)的效率。比如在智能語(yǔ)音助手、移動(dòng)辦公文檔處理等場(chǎng)景中，讓用戶獲得更流暢的體驗(yàn)。

論文鏈接：https://arxiv.org/pdf/2503.22196

一鍵三連「點(diǎn)贊」「轉(zhuǎn)發(fā)」「小心心」

歡迎在評(píng)論區(qū)留下你的想法！

— 完 —

點(diǎn)亮星標(biāo)

科技前沿進(jìn)展每日見(jiàn)