MSA (Memory Sparse Attention) 项目笔记
notes/research/github/msa-memory-sparse-attention-2026-03-19.md
Port 8777
# MSA (Memory Sparse Attention) 项目笔记
**记录时间**: 2026-03-19 13:52 (Asia/Shanghai)
**项目来源**: https://github.com/EverMind-AI/MSA
**分类**: GitHub 项目/长上下文/注意力机制
---
## 📌 项目概览
| 指标 | 信息 |
|------|------|
| **项目名称** | MSA (Memory Sparse Attention) |
| **组织** | EverMind-AI |
| **GitHub** | https://github.com/EverMind-AI/MSA |
| **主页** | https://evermind.ai |
| **创建时间** | 2025-10-29 |
| **Stars** | 58 |
| **Forks** | 6 |
| **许可证** | MIT |
| **代码状态** | Coming Soon |
| **模型状态** | Coming Soon |
---
## 🎯 核心定位
**MSA 是一个可扩展、端到端可训练的潜在记忆框架，支持 1 亿 (100M) token 上下文。**
> 长期记忆对通用智能至关重要，但全注意力瓶颈将大多数 LLM 的有效上下文长度限制在 128K-1M。
### 解决的问题
现有方案的局限性：
- **混合线性注意力**：在极端尺度下精度快速衰减、延迟增长
- **固定大小状态记忆 (如 RNNs)**：缺乏端到端可微性或动态记忆维护
- **外部存储 (如 RAG/Agents)**：需要复杂管道，缺乏端到端可微性
---
## ✨ 核心创新
### 1. Memory-Sparse Attention (MSA)
**端到端可训练的稀疏注意力层**
- **可扩展稀疏注意力** + **文档级 RoPE** (并行/全局)
- **近线性复杂度** O(L) (训练和推理)
- **端到端可微**：将检索和生成融合到单个可微循环中
- **分层路由**：路由仅应用于上层，下层保持独立文档处理
### 2. KV Cache 压缩 + Memory Parallel
**100M token 推理的 tiered 存储方案**
- **GPU 常驻路由键**：K̄ᵣ 分片跨 GPU (查询广播 → 本地评分 → 全局 reduce)
- **CPU 内容 K/V**：K̄/V̄保留在主机 DRAM，选中时异步获取
- **平衡 VRAM 和吞吐量**：在 2×A800 GPU 上实现 100M token 推理
### 3. Document-wise RoPE
**解决训练短 - 推理长的位置漂移问题**
- **并行 RoPE**：每个文档从 0 重置位置
- **全局 RoPE**：查询起始索引偏移 k (Top-k 检索块)，保持因果顺序
- **效果**：64K 训练可外推到 100M
### 4. Memory Interleave
**多轮多跳推理增强**
- 自适应交替"生成式检索 → 上下文扩展 → 生成"
- 跨分散记忆片段的多跳推理
---
## 🏗️ 架构设计
### MSA Layer 工作流程
```
1. 文档潜在状态 (K/V/Kᵣ) → chunk-mean pooling 压缩
2. 路由器投影器计算相关性 (余弦相似度)
- 头维度 mean-pool
- token-wise max
3. 选择 Top-k 文档
4. 拼接选中文档的压缩 K/V + 查询的本地 K/V
5. 自回归解码
```
### 三阶段推理管道
```
┌─────────────────────────────────────────────────────────┐
│ Stage 1: Global Memory Encoding (offline) │
│ 对整个语料库前向传播，缓存 chunk-pooled (K̄, V̄, K̄ᵣ) │
└─────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────┐
│ Stage 2: Online Routing & Context Assembly │
│ 投影查询到 Qᵣ，与 K̄ᵣ匹配选择 Top-k，加载选中 K̄/V̄ │
└─────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────┐
│ Stage 3: Sparse Generation │
│ 在稀疏上下文上自回归生成 │
└─────────────────────────────────────────────────────────┘
```
---
## 📊 实验结果
### 评估设置
**QA 任务** (9 个数据集):
- MS MARCO v1, 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
- 指标：平均准确率
**骨干模型**: Qwen3-4B-Instruct-2507
### 核心性能
| 指标 | 结果 |
|------|------|
| **16K→100M 退化率** | < 9% |
| **MS MARCO** | < 9% 退化，强外推能力 |
| **对比基线** | 超越同骨干 RAG、最佳 RAG 栈、领先长上下文模型 |
### 对比对象
- 同骨干 RAG (same-backbone RAG)
- 最佳 RAG 栈 (best-of-breed RAG stacks, 如 KaLMv2)
- 领先长上下文模型
---
## 🔑 技术亮点
### 1. 端到端可微性
MSA 将检索和生成融合到单个可微循环中，避免了传统 RAG 的分离管道问题。
### 2. 层级存储设计
```
GPU VRAM: 路由键 K̄ᵣ (分片跨 GPU)
CPU DRAM: 内容 K̄/V̄ (异步获取)
```
这种设计平衡了显存占用和吞吐量。
### 3. 文档解耦
- 训练时：独立文档处理
- 推理时：Top-k 选择 + 稀疏注意力
- 保持端到端可微的同时实现文档解耦
### 4. 位置编码创新
| RoPE 类型 | 作用 |
|-----------|------|
| **Parallel RoPE** | 每个文档重置位置，防止位置漂移 |
| **Global RoPE** | 查询索引偏移 k，保持因果顺序 |
---
## 🚀 潜在应用
1. **超长文档理解**：法律文档、学术论文、技术手册
2. **多轮对话记忆**：跨会话的长期记忆维护
3. **知识库问答**：10M+ token 级别的知识库检索
4. **多跳推理**：跨文档的复杂推理任务
5. **个人 AI 助手**：终身记忆存储和检索
---
## 🔗 相关资源
- **论文**: [MSA__Memory_Sparse_Attention_for_Efficient_End_to_End_Memory_Model_Scaling_to_100M_Tokens.pdf](https://github.com/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
- **组织主页**: https://evermind.ai
---
## 📝 个人思考
### 与现有方案对比
| 方案 | 上下文长度 | 端到端可微 | 动态记忆 | 复杂度 |
|------|------------|------------|----------|--------|
| **标准 Attention** | 128K-1M | ✅ | ✅ | O(L²) |
| **线性 Attention** | 1M+ | ✅ | ✅ | O(L) 但精度衰减 |
| **RAG** | 无限 | ❌ | ✅ | 复杂管道 |
| **RNN/状态空间** | 固定 | ✅ | ❌ | O(1) 但容量有限 |
| **MSA** | **100M** | ✅ | ✅ | **O(L)** |
### 潜在挑战
1. **代码未开源**：目前代码和模型都标记为"Coming Soon"
2. **实现复杂度**：三阶段管道 + tiered 存储需要精细工程优化
3. **硬件依赖**：Memory Parallel 需要多 GPU 配置才能发挥优势
4. **训练成本**：100M token 训练的计算和存储需求
### 跟进事项
- [ ] 关注代码开源时间
- [ ] 复现 16K→100M 外推实验
- [ ] 对比与 FlashAttention-4 的性能
- [ ] 评估在 Qwen3.5 上的集成可能性
---
*笔记创建时间：2026-03-19 13:53 (Asia/Shanghai)*