MSA (Memory Sparse Attention)

--- title: "MSA - Memory Sparse Attention for 100M Token Contexts" category: "research" source_type: "github" created_by: "小美虾" status: "new" tags: ["long-context", "attention", "memory", "LLM", "sparse-attention"] --- # MSA (Memory Sparse Attention) **来源**: https://github.com/EverMind-AI/MSA **核心亮点**: 可扩展的端到端可训练潜在记忆框架，支持 **1 亿 token 上下文** --- ## 问题背景 - 传统全注意力机制瓶颈限制 LLM 有效上下文长度为 128K-1M - 现有方案（混合线性注意力、固定大小状态记忆如 RNN、外部存储如 RAG/Agent）的问题： - 在极端尺度下精度快速衰减、延迟增长 - 缺乏端到端可微分性或动态记忆维护 - 需要复杂流水线 --- ## MSA 核心创新 ### 1. Memory Sparse Attention (MSA Layer) - 端到端可训练的稀疏注意力层 - 文档级 RoPE（并行/全局），实现 O(L) 线性复杂度 - 将检索和生成整合到单个可微分循环中 ### 2. KV Cache 压缩 - 文档潜在状态 (K/V/Kᵣ) 通过 chunk-mean pooling 压缩 - Memory Parallel 推理引擎：在 2×A800 GPU 上实现 100M token 吞吐量 ### 3. Memory Interleave - 支持跨分散记忆片段的多轮、多跳推理 ### 4. Parallel (Document-wise) RoPE - 每个文档从 0 重置位置，防止训练短上下文和推理长上下文之间的位置漂移 - 64K 训练可外推到 100M ### 5. Global RoPE (Active Context) - 查询起始索引偏移 Top-k 检索块，保持因果顺序：background → query → generation --- ## 三阶段推理流程 1. **Global Memory Encoding (离线)**: 对语料库前向传播，缓存 chunk-pooled (K̄, V̄, K̄ᵣ) 2. **Online Routing & Context Assembly (在线)**: - 查询投影到 Qᵣ - 与 K̄ᵣ 匹配选择 Top-k - 加载选中的 K̄/V̄ 并与局部上下文拼接 3. **Sparse Generation**: 在稀疏上下文上自回归生成 **Memory Parallel**: K̄ᵣ 跨 GPU 分片（查询广播 → 本地评分 → 全局 reduce），内容 K̄/V̄ 保留在主机 DRAM 中异步获取 --- ## 性能表现 ### Long-Context QA (9 个数据集) - MS MARCO, NQ, DuReader, TriviaQA(10M), NarrativeQA, PopQA, 2WikiMultiHopQA, HotpotQA, MuSiQue - 记忆库：277K → 10M tokens - 指标：LLM judge (0-5) ### NIAH (Needle-in-a-Haystack / RULER) - 8 个子任务，32K → 1M tokens - 报告平均准确率 ### 关键结果 - **16K → 100M token 范围内，性能衰减 < 9%** - 超越同 backbone 的 RAG、最佳 RAG 栈和领先的长上下文模型 - 在 MS MARCO 上展现强大的外推能力 --- ## 技术细节 - **Backbone**: Qwen3-4B-Instruct-2507 - **路由策略**: 仅应用于上层，下层保持独立文档处理以实现层次对齐 - **路由计算**: Router projector 通过余弦相似度计算相关性（head 间 mean-pooled，然后 token-wise max），选择 Top-k 文档 --- ## 资源链接 - [论文 PDF](/EverMind-AI/MSA/blob/main/paper/MSA__Memory_Sparse_Attention_for_Efficient_End_to_End_Memory_Model_Scaling_to_100M_Tokens.pdf) - [代码](Coming Soon) - [模型](Coming Soon) --- **记录时间**: 2026-03-20 08:18 **记录者**: 小美虾 🦐