[Paper] 使用 Auto-Encoders 对 Mamba 的选择性记忆进行表征

Source: arXiv - 2512.15653v1

概述

本文研究了 what 种类的信息会在 Mamba 系列状态空间语言模型（SSMs）处理更长文本流时被遗忘。通过训练一个自动编码器从 Mamba 的隐藏状态重建原始输入，作者揭示了系统性的偏差——例如数学符号、组织名称以及非标准方言——这些信息更容易丢失。了解这些盲点对于将 SSM 视为生产系统中内存高效替代 Transformer 的开发者至关重要。

关键贡献

• Token‑level forgetting analysis: 量化不同词性、命名实体类型和语言变体的遗忘率。
• Sequence‑type profiling: 显示整个领域（数学表达式、代码片段等）信息丢失更为严重。
• Auto‑encoder probing framework: 引入一种简单、可复现的方法，在不修改原始 SSM 的前提下测量隐藏状态的保真度。
• Empirical study on Mamba models: 评估参数规模从 130 M 到 1.4 B、窗口大小为 4–256 token 的模型。
• Link to pre‑training frequency: 证明了预训练语料中 token 稀有度与遗忘倾向之间的强相关性。

方法论

1. 数据准备: 作者抽样了一套多样化的句子，涵盖自然语言、代码、数学问题以及方言变体。
2. 隐藏状态提取: 将每个 token 序列输入冻结的 Mamba 模型；记录最终隐藏状态（即“记忆向量”）。
3. 自编码器训练: 一个轻量级的编码器‑解码器网络仅通过隐藏状态学习重建原始 token 序列。重建损失作为 SSM 所保留信息量的代理指标。
4. 错误分析: 将重建错误按 token 类型（词性标签、命名实体类别、方言标记）以及整体序列领域进行拆分。
5. 频率相关性: 计算原始 Mamba 预训练语料库中 token 的出现频率，并使用统计检验评估稀有度与遗忘率之间的关系。

该方法刻意保持模型无关性：任何固定记忆的语言模型都可以使用相同的自编码器流水线进行探测。

结果与发现

在预训练数据中，标记频率与其遗忘率之间发现了强负相关（Pearson r ≈ ‑0.73）。更大的模型（1.4 B）表现出更低的整体损失，但保持相同的相对偏差模式。

实际意义

• 为特定领域应用选择合适的模型： 如果你的产品处理方程式、金融数据或专业术语，普通的 Mamba 模型可能会悄悄丢失关键 token。考虑通过领域特定的微调或混合架构（例如，为高精度 token 使用小型 Transformer 缓存）来增强模型。
• 设计内存高效的流水线： 自动编码器探针可以集成到 CI 测试中，以标记新 SSM 版本开始忘记目标 token 集的情况，从而在部署前实现早期检测。
• 数据收集策略： “频率‑遗忘”关联表明，在预训练语料库中加入缺乏代表性的 token（数学符号、方言文本）可以直接提升记忆保持能力，为数据策划预算提供指导。
• 混合推理系统： 开发者可以保留一个轻量级的“token 监控器”，在检测到高风险 token 时强制执行重新编码步骤（例如，将该段落重新送入小型 Transformer），以确保准确性。
• 可解释性工具： 自动编码器生成的重建误差热图可作为调试叠加层，帮助构建 LLM 驱动助手的开发者识别模型记忆不足的区域。

限制与未来工作

• 固定窗口大小： 实验仅在 256 个 token 处停止；在真正的长文档（数千个 token）上的行为仍未经过测试。
• 自编码器容量： 探针本身可能引入偏差；更具表达能力的解码器可能会掩盖遗忘而非揭示它。
• 模型范围： 仅研究了 Mamba 系列；尚不清楚观察到的模式是否能推广到其他 SSM 变体（例如 S4、Hyena）。
• 缓解策略： 论文指出了问题，但未提出具体的架构改动或训练目标来降低选择性遗忘。未来工作可以探索记忆增强技术、基于课程的预训练或 token 感知正则化等方向。

作者

• Tamanna Hossain
• Robert L. Logan
• Ganesh Jagadeesan
• Sameer Singh
• Joel Tetreault
• Alejandro Jaimes

论文信息

• arXiv ID: 2512.15653v1
• 分类: cs.CL
• 发表时间: 2025年12月17日
• PDF: 下载 PDF