Mamba 中隐藏注意力

Phoenixtree_DongZhao

已于 2024-08-10 11:26:38 修改

阅读量595

点赞数 23

分类专栏： Mamba MyDLNote-Attention 文章标签： Mamba 注意力机制

于 2024-08-10 11:22:51 首次发布

本文链接：https://blog.csdn.net/u014546828/article/details/141087836

版权

MyDLNote-Attention 同时被 2 个专栏收录

40 篇文章 7 订阅

订阅专栏

Mamba

21 篇文章 1 订阅

订阅专栏

The Hidden Attention of Mamba Models

Arxiv

GitHub

Ameen Ali ∗ , Itamar Zimerman ∗ , and Lior Wolf

School of Computer Science, Tel Aviv University

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.

Mamba层提供了一种高效的选择性状态空间模型（SSM），该模型在建模多个领域方面表现优异，包括自然语言处理（NLP）、长距离序列处理和计算机视觉。

选择性SSM被视为双重模型，其中一个模型通过IO感知的并行扫描对整个序列进行并行训练，并以自回归方式进行部署。

本文增加了第三种视角，并表明此类模型可以视为注意力驱动模型。这一新视角使能够从经验上和理论上将其基本机制与Transformer中的自注意力层进行比较，并允许使用可解释性方法深入了解Mamba模型的内部工作原理。

这一新的视角不仅增强了我们对Mamba模型根本性质的理解，还通过展示它们实际上是通过一个独特的数据控制线性运算符来实现隐式注意力，揭示了它们隐藏的注意力机制。

Introduction

Mamba模型的概述：

核心概念：Mamba层提供了一个高效的选择性状态空间模型（SSM），该模型在包括自然语言处理（NLP）、长序列处理和计算机视觉在内的多个领域表现优异。
训练与部署：SSM被视为双模型，并行于整个序列进行训练，通过IO-aware并行扫描，并在自回归方式下部署。

研究动机：

新视角：文章提出了一种新的视角，即Mamba模型可以视为注意力驱动模型，这使得与Transformer的自注意力层进行理论和实证比较成为可能。
现有研究不足：尽管Mamba模型在多领域取得了显著成功，但其token间的信息流动态和学习机制尚未得到充分探索。
应用限制：Mamba模型缺乏互操作性方法，这可能会阻碍其调试并降低在需要可解释性的社会敏感领域中的应用性。

研究目标：

解析机制：通过研究旨在提供对Mamba模型动态的见解，并开发对其进行解释的方法。
视角转换：从传统的卷积或循环层视角转变为注意力模型视角，揭示了Mamba计算中的隐藏注意力矩阵。

主要贡献：

隐性注意力：揭示了Mamba模型依赖于隐性注意力，通过独特的数据控制线性运算符实现。
注意力矩阵数量：分析表明，Mamba模型产生的注意力矩阵数量比Transformer多三个数量级。
解释性工具：基于这些隐藏注意力矩阵，提供了一套解释性和可解释性工具。
解释性指标：在可比模型大小下，Mamba模型的注意力展示出了与Transformer相当的解释性指标结果。
理论分析：对状态空间模型中注意力能力的演变和表达能力进行了理论分析，为理解Mamba的有效性提供了更深入的理解。

Method

文档首先指出了Mamba模型在多个领域（如NLP、长序列处理和计算机视觉）中的高效性和有效性，但关于其信息流动态和学习机制的研究尚不充分。因此，本文旨在揭示Mamba模型的动态特性，并开发解释性方法。

Mamba模型的基本结构

Mamba模型通过选择性状态空间模型（SSM）实现高效处理。每个Mamba块包含多个通道，每个通道独立处理，并通过数据控制线性算子实现隐式注意力机制。

隐式注意力矩阵的揭示

数据控制线性算子：通过重新表述Mamba计算，使用数据控制线性算子来揭示Mamba层内的隐式注意力矩阵。
公式推导：文档详细推导了如何从Mamba模型的数学表达中提取隐式注意力矩阵。这包括从系统矩阵（如¯A, ¯B, C）和输入信号（如ˆx）中推导出注意力矩阵˜α。

注意力矩阵的生成

注意力矩阵的公式：文档给出了注意力矩阵˜α的详细计算公式（如Eq. 11），该公式通过考虑输入信号的非线性变换和状态空间模型的动态特性来计算注意力权重。
多头和多通道：由于Mamba模型中的通道数（D）远大于Transformer中的头数（H），Mamba层生成的注意力矩阵数量是Transformer的多个数量级。

解释性工具的开发

Attention-Rollout的改编：基于Attention-Rollout方法，为Mamba模型开发了相应的解释性工具（Mamba-Attribution）。该方法通过累加各通道的注意力矩阵来生成整体的注意力图。
Transformer-Attribution：提出了一个针对Mamba模型的Transformer-Attribution方法，该方法结合了LRP分数和注意力梯度来生成相关性分数。

实验与分析

注意力矩阵的可视化：通过可视化Mamba和Transformer在NLP和视觉任务中的注意力矩阵，展示了两者在注意力分布上的差异。
解释性评估：使用多种解释性方法（如Raw-Attention、Attn-Rollout、Transformer-Attr、Mamba-Attr）对Mamba模型进行解释性评估，并通过定量指标（如AUC、像素准确率、mAP、mIoU）来比较不同方法的性能。