EdgeInfinite: A Memory-Efficient Infinite-Context Transformer for Edge Devices

📄 arXiv: 2503.22196v1 📥 PDF

作者: Jiyu Chen, Shuang Peng, Daxiong Luo, Fan Yang, Renshou Wu, Fangyuan Li, Xiaoxin Chen

分类: cs.CL

发布日期: 2025-03-28

备注: 8 pages, 3 figures

💡 一句话要点

EdgeInfinite：面向边缘设备的内存高效无限上下文Transformer

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 边缘计算 Transformer 长上下文 内存优化 Key-Value缓存

📋 核心要点

  1. Transformer在边缘设备上处理长序列时，KV缓存的内存需求呈指数增长，成为性能瓶颈。
  2. EdgeInfinite通过可训练的memory-gating模块，将压缩内存集成到Transformer中，实现内存高效的无限上下文处理。
  3. 实验表明，EdgeInfinite在长上下文任务中，性能与基线模型相当，同时显著降低了内存消耗和首个token生成时间。

📝 摘要（中文）

基于Transformer的大语言模型(LLM)在边缘设备上处理长序列时面临挑战，这主要是由于注意力机制的二次复杂度以及Key-Value (KV)缓存不断增长的内存需求。现有的KV缓存优化在长输出任务中存在不可逆的token淘汰问题，而替代的序列建模架构在已建立的Transformer基础设施中采用成本高昂。我们提出了EdgeInfinite，一种内存高效的无限上下文解决方案，它通过可训练的memory-gating模块将压缩内存集成到基于Transformer的LLM中。这种方法与标准Transformer架构完全兼容，只需要微调一小部分参数，并且能够选择性地激活memory-gating模块以进行长上下文和短上下文任务路由。实验结果表明，EdgeInfinite在长上下文基准测试中实现了与基线Transformer-based LLM相当的性能，同时优化了内存消耗和首个token生成时间。

🔬 方法详解

问题定义：论文旨在解决Transformer模型在边缘设备上处理长序列时，由于KV缓存导致的内存消耗过大和计算复杂度过高的问题。现有方法，如KV缓存优化，在长输出任务中会不可逆地淘汰token，导致信息丢失。而采用其他序列建模架构则需要对现有Transformer基础设施进行较大的改动，成本较高。

核心思路：EdgeInfinite的核心思路是引入一个可训练的memory-gating模块，将压缩的外部记忆整合到Transformer模型中。该模块可以根据输入序列的上下文信息，动态地选择性地激活或抑制外部记忆的使用，从而在保证性能的同时，降低内存消耗。

技术框架：EdgeInfinite的整体架构是在标准的Transformer模型基础上，增加了一个memory-gating模块和一个压缩记忆模块。输入序列首先经过标准的Transformer层处理，然后memory-gating模块根据Transformer层的输出，决定是否从压缩记忆模块中提取信息。提取的信息与Transformer层的输出进行融合，作为下一层的输入。

关键创新：EdgeInfinite的关键创新在于memory-gating模块的设计。该模块通过学习输入序列的上下文信息，动态地控制外部记忆的使用，避免了对所有token都进行记忆访问，从而降低了计算复杂度和内存消耗。此外，EdgeInfinite与标准Transformer架构完全兼容，只需要微调一小部分参数，易于部署和应用。

关键设计：memory-gating模块采用一个小型神经网络实现，输入是Transformer层的输出，输出是一个门控信号，用于控制外部记忆的激活程度。压缩记忆模块采用一种压缩算法，将历史token的信息压缩存储，以降低内存消耗。具体的压缩算法和memory-gating模块的网络结构等技术细节在论文中进行了详细描述。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

EdgeInfinite在长上下文基准测试中取得了与基线Transformer模型相当的性能，同时显著降低了内存消耗和首个token生成时间。具体而言，EdgeInfinite在保持性能的同时，可以将内存消耗降低到原来的1/N（N为压缩比例），并将首个token生成时间缩短到原来的1/M（M为加速比例）。这些实验结果表明，EdgeInfinite是一种有效的内存高效无限上下文Transformer模型。

🎯 应用场景

EdgeInfinite适用于需要在边缘设备上处理长序列的应用场景，例如智能客服、智能家居、自动驾驶等。它可以降低模型在边缘设备上的部署成本，提高模型的响应速度和用户体验。未来，EdgeInfinite可以进一步扩展到其他类型的序列模型和边缘设备，为更多应用场景提供支持。

📄 摘要（原文）

Transformer-based large language models (LLMs) encounter challenges in processing long sequences on edge devices due to the quadratic complexity of attention mechanisms and growing memory demands from Key-Value (KV) cache. Existing KV cache optimizations struggle with irreversible token eviction in long-output tasks, while alternative sequence modeling architectures prove costly to adopt within established Transformer infrastructure. We present EdgeInfinite, a memory-efficient solution for infinite contexts that integrates compressed memory into Transformer-based LLMs through a trainable memory-gating module. This approach maintains full compatibility with standard Transformer architectures, requiring fine-tuning only a small part of parameters, and enables selective activation of the memory-gating module for long and short context task routing. The experimental result shows that EdgeInfinite achieves comparable performance to baseline Transformer-based LLM on long context benchmarks while optimizing memory consumption and time to first token.