手机流畅处理128K长文本,vivo端侧新算法突破内存限制
在端侧设备上处理长文本常常面临计算和内存瓶颈。
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
一键三连「点赞」「转发」「小心心」
欢迎在评论区留下你的想法!
— 完 —

点亮星标
科技前沿进展每日见