[论文解读] The Hidden Attention of Mamba Models

Ameen Ali, Itamar Zimerman+1 more|arXiv (Cornell University)|Mar 3, 2024

Explainable Artificial Intelligence (XAI)被引用 5

一句话总结

本文将 Mamba 选择性状态空间层重新表述为自注意力的模仿，推导隐藏的注意力矩阵，并开发可解释性工具以将 Mamba 与 transformers 进行比较并实现 XAI 应用。

ABSTRACT

The Mamba layer offers an efficient selective state space model (SSM) that is highly effective in modeling multiple domains, including NLP, long-range sequence processing, and computer vision. Selective SSMs are viewed as dual models, in which one trains in parallel on the entire sequence via an IO-aware parallel scan, and deploys in an autoregressive manner. We add a third view and show that such models can be viewed as attention-driven models. This new perspective enables us to empirically and theoretically compare the underlying mechanisms to that of the self-attention layers in transformers and allows us to peer inside the inner workings of the Mamba model with explainability methods. Our code is publicly available.

研究动机与目标

• 揭示 Mamba 选择性状态空间层的内部信息流动动力学。
• 证明 Mamba 可以被视为一种隐性、因果的自注意力机制。
• 基于隐藏的注意力矩阵，为 Mamba 开发和改编可解释性工具。
• 在可解释性和表达能力方面，将基于 Mamba 的注意力与 transformer 注意力进行比较。
• 提供关于状态空间模型中注意力的发展与能力的理论与实证见解。

提出的方法

• 将选择性状态空间层重构为数据控制的线性算子，以提取隐藏的注意力矩阵。
• 推导逐通道的隐藏注意力矩阵，并通过数据相关算子将其与常规注意力联系起来，同时给出便于解释的简化形式。
• 将 Attention Rollout 适配到 Mamba，以实现跨层与跨通道的类别无关可解释性。
• 通过把 Transformer-Attribution 调整为使用 Mamba 注意力矩阵和门控信号，开发面向 Mamba 的特定注意力归因方法。
• 在视觉和语言任务中可视化并比较隐藏的注意力矩阵。
• 在 ImageNet 及相关数据集上，利用扰动和分段度量评估可解释性方法。

实验结果

研究问题

• RQ1 选择性状态空间层（Mamba）是否可以被解释为一种隐式注意力机制，类似于 transformers 中的自注意力？
• RQ2 Mamba 中隐藏注意力的特性与尺度是什么，它们与 transformer 注意力有何比较？
• RQ3 我们如何推导并利用隐藏的注意力矩阵来实现 Mamba 的类别无关和类别特定的可解释性？
• RQ4 在视觉和自然语言处理任务中，以 Mamba 为基础的可解释性方法是否达到与基于 transformer 的方法相当的可解释性？
• RQ5 哪些理论见解解释了状态空间模型中注意力的表达能力及其演变？

主要发现

• Mamba 层可以通过数据控制的线性算子重构为一种隐式因果自注意力机制。
• 单个 Mamba 通道会产生许多隐藏注意力矩阵，导致注意力图数量显著多于 transformers。
• 隐藏注意力矩阵使新型可解释性工具成为可能，在某些情形下可与基于 transformer 的可解释性相媲美。
• 在相近模型规模下，基于 Mamba 的注意力在可解释性指标上与 transformers 具有竞争力。
• 理论分析表明，数据控制的非对角混合器支撑着 Mamba 的富表达注意力能力及其情境学习潜力。