By Unbug 简介
在人工智能的世界里，序列建模就像是一场”记忆马拉松”——从传统的循环神经网络（RNN）到革命性的 Transformer，再到近年来异军突起的状态空间模型（SSMs），研究者们一直在寻找既能高效处理超长序列，又能牢牢记住关键信息的完美架构。
2025 年，一篇题为《MemMamba: Rethinking Memory Patterns in State Space Model》的论文被 ICLR 2026 接收，这篇论文提出了一个突破性的架构，成功地解决了 Mamba 等状态空间模型在长序列建模时面临的”记忆衰退”问题。
背景：序列建模的困境
让我们先回顾一下序列建模领域面临的核心挑战：
1. 传统 RNN 和 LSTM 的局限性
想象一下，你在读一本长达 1000 页的小说，但每翻一页，你就忘记了前面 999 页的内容——这就是传统 RNN 和 LSTM 面临的困境！
- 梯度消失/爆炸问题：在处理超长序列时，早期的信息几乎完全丢失
- 计算效率：只能顺序处理，无法并行化，就像只能一个字一个字地看书
- 实际限制：通常只能有效处理几千个 token
2. Transformer 的优势与劣势
Transformer 的出现就像是给 AI 装上了”全景望远镜”——通过自注意力机制，它可以同时看到整个序列的所有内容。但这个望远镜也有一个致命的缺点：
- 优势：能够全局建模依赖关系，性能出色
- 劣势：计算复杂度是 O(n²)，就像每次看东西都要把整个宇宙扫描一遍
- 实际限制：即使有稀疏注意力等优化，也很难处理超过 10 万 token 的序列
3. Mamba 的突破与局限
2023 年，Mamba 架构的出现带来了新的希望——它就像是一辆”高速列车”，既快又稳：
- 突破：线性时间复杂度 O(n)，推理效率高，在语言建模等任务上超越了同规模的 Transformer
- 局限：虽然计算效率高，但 Mamba 的长程记忆会指数级衰减——就像你有一个记忆力超强但只能记住最近 10 分钟事情的助手！
问题的根源：Mamba 的记忆衰减机制
论文的第一部分通过数学推导和信息论分析，系统性地揭示了 Mamba 的记忆衰减机制。
数学分析
Mamba 的状态更新方程为：
h_t = A · h_{t-1} + B · x_t
y_t = C · h_t
| 其中 A 是状态转移矩阵，满足 | A | < 1 以保证稳定性。 |
对于 k 步之前的输入 x_{t-k}，它对当前状态 h_t 的贡献可以表示为：
Contribution(x_{t-k} → h_t) = |A^k · B · x_{t-k}| ≤ |A|^k · |B| · |x_{t-k}|
关键洞察：随着 k 增大（即输入越久远），A^k 会指数级衰减，导致早期输入几乎完全被遗忘！
横向-纵向记忆保真度框架
为了量化关键信息的损失，论文引入了一个创新的框架：
- 预期 Token 记忆保真度（ETMF）：衡量 token 语义信息在层内水平传播时的保留程度
- 预期跨层记忆保真度（ECLMF）：衡量信息在层间垂直传输时的保留程度
这两个指标从互补的角度揭示了 Mamba 中记忆衰减的双重挑战。
MemMamba 的创新设计
受到人类阅读长文档时记笔记的启发，MemMamba 提出了一个全新的架构框架。
核心思想：像人类一样”记笔记”
当人类阅读长文档时，我们不会试图记住每一个字，而是会：
- 提取关键信息
- 做笔记
- 在需要时回顾笔记
MemMamba 正是模仿了这个过程！
架构组件
MemMamba 由 n 个堆叠的 MemMamba 块层组成，每个层集成了三个核心组件：
1. Note Block（笔记模块）
这个模块动态识别和提取关键信息，模仿人类的记笔记过程：
- 重要性评分：使用评分函数 I_token 评估每个 token 的重要性
- 阈值判断：如果评分超过阈值 τ1，就执行”记笔记”操作
- 状态池：将压缩后的摘要插入到状态池中（采用 FIFO 或优先级替换策略）
2. 跨 Token 注意力
在每个层中，如果检测到遗忘，就在状态池和当前输入之间执行跨 Token 注意力：
if I_state(z_{t-1}^l) > τ2 ⇒
c_{token,t}^l = Attention(Q = x_t^l, K = ˜s_{t-1}^l, V = ˜s_{t-1}^l)
3. 跨层注意力
每隔 p 层触发一次跨层注意力，聚合之前层的状态池：
if l mod p = 0 ⇒
c_{layer,t}^l = Attention(Q = x_t^l, K = s_R(l), V = s_R(l))
设计优势
这种双阈值和稀疏跨层机制确保了：
- 跨 Token 补充在每个层都发生
- 跨层记忆交互被稀疏激活
- 在记忆保留和计算效率之间取得平衡
实验结果
论文在多个长序列基准测试上进行了全面的实验评估，结果令人印象深刻。
1. PG19 语言建模任务
PG19 数据集包含约 1 亿 token 的英语小说，平均长度为 69k token。
困惑度（PPL）对比：
| 模型 | 参数 | 1K | 2K | 4K | 10K | 20K | 30K | 40K | 50K | 60K |
|---|---|---|---|---|---|---|---|---|---|---|
| Mamba | 130M | 21.00 | 19.60 | 18.77 | 19.29 | 31.63 | INF | INF | INF | INF |
| DeciMamba | 150M | 21.90 | 20.06 | 18.55 | 21.98 | 23.15 | 27.05 | 40.48 | INF | INF |
| MemMamba | 200M | 19.35 | 18.23 | 17.52 | 17.71 | 18.25 | 17.33 | 17.54 | 17.97 | 17.35 |
惊人的发现：
- 在 30k-60k token 的超长序列中，Mamba 和 DeciMamba 完全崩溃（PPL > 100）
- MemMamba 保持稳定的 PPL（17.33-18.25）
- 即使在 60k token 时，PPL 仍然只有 17.35！
2. Passkey Retrieval 任务
在这个任务中，一个目标 token 被随机插入到极长的输入序列中，模型需要在预测时精确检索这个信息。
检索准确率对比：
| 模型 | 1K | 2K | 4K | 8K | 16K | 32K | 64K | 128K | 256K | 400K |
|---|---|---|---|---|---|---|---|---|---|---|
| Pythia-160M | 1.0 | 1.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 |
| Mamba-130M | 1.0 | 1.0 | 1.0 | 1.0 | 0.8 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 |
| DeciMamba-130M | 1.0 | 1.0 | 1.0 | 1.0 | 1.0 | 1.0 | 1.0 | 1.0 | 1.0 | 0.6 |
| MemMamba | 1.0 | 1.0 | 1.0 | 1.0 | 1.0 | 1.0 | 1.0 | 1.0 | 1.0 | 0.9 |
令人印象深刻：
- 在 400k token 时，MemMamba 仍然保持 90% 的检索准确率！
- 相比之下，DeciMamba 只有 60%，其他模型完全失败
3. 文档检索任务（含噪声）
在这个任务中，随着噪声文档数量的增加，模型需要在干扰中找到正确的信息。
| 模型 | 10 | 20 | 120 | 160 | 200 |
|---|---|---|---|---|---|
| Mamba | 0.68 | 0.71 | 0.01 | 0 | 0 |
| DeciMamba | 0.72 | 0.74 | 0.48 | 0.19 | 0.12 |
| MemMamba | 0.80 | 0.66 | 0.52 | 0.44 | 0.24 |
关键优势：在高噪声条件下，MemMamba 始终保持更高的分数，展示了跨文档和跨域推理任务的优势。
4. 效率评估
尽管引入了额外的计算来增强建模能力，MemMamba 的端到端延迟仅为 Transformer 的 0.52 倍——即48% 的推理加速！
理论贡献
论文不仅提供了实证结果，还进行了深入的理论分析：
线性复杂度证明
尽管引入了状态摘要和跨层注意力，MemMamba 仍然保持线性的时间和空间复杂度 O(n·d)，而 Transformer 是 O(n²d)。
BIBO 稳定性
论文证明了 MemMamba 具有有界输入有界输出（BIBO）稳定性，不会出现发散或病态衰减。
长序列召回保证
理论分析表明，MemMamba 的长序列召回率 ≥ 90%，远超过 Mamba 和 Transformer。
影响与意义
MemMamba 的突破具有重要的理论和实践意义：
1. 理论意义
- 首次系统性地分析了 Mamba 的记忆机制
- 提出了横向-纵向记忆保真度框架，为理解长序列模型提供了新视角
- 展示了如何在不牺牲效率的情况下增强记忆能力
2. 实践意义
- 超长文本理解：可以处理整本书籍、长文档等
- 基因组分析：DNA 序列通常非常长，MemMamba 可以高效处理
- 多模态长序列：视频、音频等时序数据
- 实时应用：48% 的推理加速使其适合实时应用
3. 未来方向
论文指出了几个有前景的未来方向：
- 多模态扩展：将 MemMamba 扩展到多模态设置
- 检索增强系统：与检索增强系统集成
- 大规模扩展：将 MemMamba 扩展为高效、高保真记忆的基础模型
总结
MemMamba 代表了向新一代以记忆为中心的神经架构迈出的重要一步，这种架构将记忆保留和推理视为一等公民。
核心要点：
- 问题诊断：通过数学推导揭示了 Mamba 的记忆衰减机制
- 创新设计：模仿人类记笔记的过程，引入状态摘要和跨层/跨 Token 注意力
- 实证突破：在多个基准测试上取得 state-of-the-art 结果
- 理论保证：证明了线性复杂度、稳定性和高召回率
MemMamba 的成功表明，通过巧妙的架构设计，我们可以打破”复杂度-记忆”的权衡，为超长序列建模开辟新的可能性。
论文链接：https://arxiv.org/pdf/2510.03279
作者：Youjin Wang, Yangjingyi Chen, Jiahao Yan, Jiaxuan Lu, Xiao Sun
机构：中国人民大学、上海财经大学、上海人工智能实验室
世界模型统一框架：突破任务特定知识注入的局限性