解析Mamba输入选择性如何提升函数逼近与记忆能力
Understanding Input Selectivity in Mamba: Impact on Approximation Power, Memorization, and Associative Recall Capacity
- 通过证明S6层可表示哈尔小波投影,揭示其对不连续函数的逼近优势
- 发现S6层能动态抑制长期记忆衰减,增强模型持久记忆能力
- 首次给出Mamba系列在联想回忆任务中的理论解,适合模型机制研究者
状态空间模型(SSMs)特别是Mamba,已成为Transformer的有前途替代方案。Mamba在其SSM层(S6)中引入输入选择性,并将卷积与门控机制融入模块设计。尽管这些改进提升了性能,但输入选择性如何与架构中其他操作协同作用仍不明确。本文深入剖析输入选择性在Mamba中的作用,研究其对函数逼近能力、长期记忆与联想回忆的影响:(i) 证明了Mamba的S6层可表示哈尔小波投影,在逼近实际中常见的不连续函数方面优于其前代模型S4D;(ii) 展示了S6层可动态对抗记忆衰减;(iii) 针对MQAR联想回忆任务,给出了Mamba、Mamba-2和S4D不同混合器的解析解,并通过具体任务的实证结果验证了理论构造的紧致性。研究为Mamba提供了机制层面的理解,也揭示了改进方向。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。