论文
MSA:内存稀疏注意力,可实现高效的端到端内存模型扩展到 100M 词元
MSA: Memory Sparse Attention for Efficient End-to-End Memory Model Scaling to 100M Tokens
摘要
长期记忆是人类智力的基石。让人工智能能够处理生命周期的信息仍然是该领域的长期追求。由于全注意力架构的限制,大语言模型 (LLM) 的有效上下文长度通常限制为 1M 个词元。现有的方法,例如混合线性注意力、固定大小的记忆状态(例如 RNN)以及 RAG 或代理系统等外部存储方法,试图扩展这一限制。然而,它们通常会遭受严重的精度下降和随着上下文长度的增长而迅速增加的延迟、无法动态修改内存内容或缺乏端到端优化的困扰。这些瓶颈阻碍了大型语料库摘要、数字孪生和长历史代理推理等复杂场景,同时限制了内存容量并减慢了推理速度。我们提出了内存稀疏注意力(MSA),这是一种端到端的可训练、高效且可大规模扩展的内存模型框架。通过可扩展的稀疏注意力和文档级 RoPE 等核心创新,MSA 在训练和推理方面实现了线性复杂性,同时保持卓越的稳定性,从 16K 到 100M 词元扩展时,性能下降不到 9%。此外,KV 缓存压缩与内存并行相结合,可在 2xA800 GPU 上实现 100M 词元推理。我们还提出内存交错来促进跨分散内存段的复杂多跳推理。 MSA 在长上下文基准测试中显着超越了前沿 LLM、最先进的 RAG 系统和领先的内存代理。这些结果表明,通过将内存容量与推理分离,MSA 提供了一个可扩展的基础,为通用模型赋予内在的、生命周期规模的内存。