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

要約

状態空間モデル（SSM）、特にMambaは最近、トランスの有望な代替品として浮上しています。
Mambaは、SSM層（S6）に入力選択性を導入し、畳み込みとゲーティングをブロック定義に組み込みます。
これらの変更は、SSMの前任者に対するMambaのパフォーマンスを改善しますが、MAMBAが入力選択性によって提供される追加の機能をどのように活用するか、およびこれらがMAMBAアーキテクチャの他の操作とどのように相互作用するかはほとんど不明のままです。
この作業では、マンバにおける入力選択性の役割を分かり、関数近似の出力、長期的な暗記、および連想リコール能力への影響を調査します。
特に：（i）マンバのS6層がHaarウェーブレットへの投影を表すことができることを証明し、実際に発生する不連続機能を近似する斜めのSSM（S4D）の前身にエッジを提供します。
（ii）S6層がどのようにメモリ減衰に動的に対抗できるかを示します。
（iii）Mamba、Mamba-2、およびS4Dのさまざまなミキサーを使用したMAMBAアーキテクチャを使用して、MQAR連合リコールタスクに分析ソリューションを提供します。
具体的なタスクに関する経験的な結果を伴う理論的構成の緊密性を示します。
私たちの調査結果は、マンバの機械的な理解を提供し、改善の機会を明らかにします。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google