Understanding Input Selectivity in Mamba: Impact on Approximation Power, Memorization, and Associative Recall Capacity

📄 arXiv: 2506.11891v1 📥 PDF

作者: Ningyuan Huang, Miguel Sarabia, Abhinav Moudgil, Pau Rodriguez, Luca Zappella, Federico Danieli

分类: cs.LG

发布日期: 2025-06-13

💡 一句话要点

揭示Mamba中的输入选择性对近似能力和记忆的影响

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 状态空间模型 输入选择性 函数近似 长期记忆 联想回忆 Mamba架构 深度学习

📋 核心要点

  1. 现有的状态空间模型在处理复杂函数和记忆保持方面存在局限性，尤其是在不连续函数的近似和长期记忆的衰退问题上。
  2. 论文提出了Mamba架构，通过引入输入选择性和卷积、门控机制，增强了模型的近似能力和记忆保持能力。
  3. 实验结果表明，Mamba在近似不连续函数和联想回忆任务上优于其前身S4D，验证了理论构建的紧密性。

📝 摘要（中文）

状态空间模型（SSMs），尤其是Mamba，最近成为Transformer的有力替代方案。Mamba在其SSM层（S6）中引入了输入选择性，并在其模块定义中结合了卷积和门控机制。尽管这些改进提升了Mamba的性能，但其如何利用输入选择性的额外功能仍不清晰。本文探讨了输入选择性在Mamba中的作用，分析其对函数近似能力、长期记忆和联想回忆能力的影响。我们证明了Mamba的S6层能够表示Haar小波的投影，显示出其在近似实际中常见的不连续函数方面的优势；同时，我们展示了S6层如何动态抵消记忆衰退；最后，我们提供了使用不同混合器的Mamba架构在MQAR联想回忆任务上的解析解。我们的发现为Mamba提供了机制理解，并揭示了改进的机会。

🔬 方法详解

问题定义：本文旨在解决Mamba架构在函数近似、长期记忆和联想回忆能力方面的不足，特别是如何有效利用输入选择性。

核心思路：通过引入输入选择性，Mamba能够更灵活地处理输入信息，从而提升模型在复杂任务中的表现，尤其是在记忆保持和函数近似方面。

技术框架：Mamba架构由多个模块组成，包括S6层、卷积层和门控机制。S6层负责输入选择性，卷积层用于特征提取，而门控机制则调节信息流动。

关键创新：Mamba的S6层能够表示Haar小波的投影，这是其相较于S4D的主要创新，使其在处理不连续函数时表现更佳。

关键设计：在设计中，Mamba的损失函数和网络结构经过优化，以适应输入选择性带来的新特性，确保模型在动态记忆和近似能力上的提升。具体参数设置和网络层次结构在实验中进行了详细验证。

📊 实验亮点

实验结果显示，Mamba在MQAR联想回忆任务中表现出色，尤其是与S4D相比，近似不连续函数的能力提升显著，验证了理论分析的有效性。具体性能数据表明，Mamba在多个任务上均优于传统模型，展现出强大的应用潜力。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、图像识别和时间序列预测等。Mamba架构的输入选择性和记忆能力使其在处理复杂数据时具有更高的灵活性和准确性，未来可能在智能助手和自动化系统中发挥重要作用。

📄 摘要（原文）

State-Space Models (SSMs), and particularly Mamba, have recently emerged as a promising alternative to Transformers. Mamba introduces input selectivity to its SSM layer (S6) and incorporates convolution and gating into its block definition. While these modifications do improve Mamba's performance over its SSM predecessors, it remains largely unclear how Mamba leverages the additional functionalities provided by input selectivity, and how these interact with the other operations in the Mamba architecture. In this work, we demystify the role of input selectivity in Mamba, investigating its impact on function approximation power, long-term memorization, and associative recall capabilities. In particular: (i) we prove that the S6 layer of Mamba can represent projections onto Haar wavelets, providing an edge over its Diagonal SSM (S4D) predecessor in approximating discontinuous functions commonly arising in practice; (ii) we show how the S6 layer can dynamically counteract memory decay; (iii) we provide analytical solutions to the MQAR associative recall task using the Mamba architecture with different mixers -- Mamba, Mamba-2, and S4D. We demonstrate the tightness of our theoretical constructions with empirical results on concrete tasks. Our findings offer a mechanistic understanding of Mamba and reveal opportunities for improvement.