Mamba论文

题目

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Mamba：具有选择性状态空间的线性时间序列建模

作者简介

Albert Gu

Tri Dao

状态空间模型SSM

详细网站

https://blog.csdn.net/weixin_39699362/article/details/136295477

定义

结构状态空间模型（Structural State Space Model, SSM）是一种用于描述和分析时间序列数据的数学模型。它通过显式地表示时间序列数据的潜在结构（即状态）以及这些状态随时间的演变方式，来捕捉数据的动态特性和复杂的依赖关系。

1. 基本概念

• 状态变量（State Variables）: 这些是隐藏的（不可观测的）变量，代表系统的潜在状态。它们随时间变化，反映了系统内部的动态特性。
• 观测变量（Observation Variables）: 这些是可观测的变量，通常是我们直接可以测量或看到的数据。观测变量是状态变量的某种函数加上噪声。
• 状态转移方程（State Transition Equation）: 描述状态变量如何随时间变化，通常是前一时刻状态和噪声的函数。
• 观测方程（Observation Equation）: 描述观测变量如何由状态变量生成，通常包括状态变量和观测噪声。
• 

3. 应用

结构状态空间模型广泛应用于各种领域，包括但不限于：

• 经济学和金融: 用于经济指标的建模和预测、资产价格建模等。
• 工程和控制系统: 用于系统状态估计、故障检测、自动控制等。
• 生物医学: 用于心脏信号建模、流行病传播建模等。
• 气象学: 用于天气预报和气候建模。

4. 估计和推断

为了使用状态空间模型进行分析，我们通常需要进行以下步骤：

1. 参数估计: 估计状态转移矩阵、观测矩阵等参数。这通常通过最大似然估计或贝叶斯方法来完成。
2. 状态估计: 给定观测数据，估计状态变量。这通常通过卡尔曼滤波、粒子滤波等方法来完成。
3. 预测: 基于估计的状态变量和模型参数，对未来进行预测。

5. 优点

• 灵活性: 可以处理非平稳时间序列，适应不同类型的数据结构。
• 系统性: 能够显式地描述系统的动态特性和观测过程。
• 可解释性: 状态变量和模型参数具有明确的物理意义
• 

图1  选择性SSM概述

涉及到一种高级的结构化状态空间模型（SSM），以及它如何通过高维潜在状态和巧妙的计算路径来处理多通道输入数据并映射到输出。

离散化方法

定义

符号解释

零阶保持器

Computation计算

主要内容

一种混合训练和推理策略，该策略在训练和推理过程中使用不同的模式来提高效率。具体来说，这种策略在训练过程中使用并行的卷积模式，而在推理过程中使用递归模式。这种方法结合了卷积神经网络（CNN）和递归神经网络（RNN）的优点，以提高模型的训练效率和推理性能。

卷积如何进行并行操作的

卷积算法并行编程原理是通过将卷积操作划分为多个小的并行任务，利用多个处理单元同时计算，以加速卷积计算的过程。

并行编程的原理可以基于多种不同的架构和编程模型。以下是几种常见的卷积算法并行编程原理：

1. 数据并行：将输入数据划分为多个部分，每个处理单元负责处理一部分数据。这种方式适用于多核CPU或GPU等多处理器架构。每个处理单元独立计算局部卷积，最后将结果合并得到最终的卷积结果。

2. 线程并行：使用线程来并行执行卷积操作。这种方式适用于支持线程并行的多核CPU或GPU。每个线程负责处理输入数据的一部分，并独立计算局部卷积。通过合并各个线程的计算结果，得到最终的卷积结果。

3. 图像并行：将输入图像划分为多个较小的块，并将每个块分配给不同的处理核心进行独立计算。这种方式适用于处理大尺寸的图像。每个处理核心独立计算局部卷积，并将结果合并得到最终的卷积结果。

4. 指令级并行：利用SIMD（单指令多数据）指令集来实现并行计算。通过同时处理多个数据元素，以提高运算速度。这种方式适用于支持SIMD指令集的处理器。通过对输入数据进行向量化操作，实现并行计算。

在实际应用中，我们可以根据具体的硬件架构和编程模型选择适合的并行编程原理，并结合优化技术，如数据重用、内存局部性等，进一步提高卷积算法的并行性能。

卷积和递归有什么不同？

递归概念及优缺点

卷积概念及优缺点

不同之处

推理模式使用递归的好处

线性时间不变性(LTI)

探讨了模型动力学随时间保持恒定的重要性质，即线性时不变性（LTI）

线性时不变性（LTI）：

方程1-3

LTI在本文概念、LTI模型的局限性和作者对于LIT模型的改进

结构及尺寸

讨论了结构化SSM（结构化状态空间模型）的命名原因和其特定的结构化形式，重点是如何有效地计算这些模型

一般空间状态模型

讨论了“状态空间模型”（State Space Model, SSM）这一术语在不同领域中的广泛含义，以及该论文中对这一术语的具体使用

S4模型

SSM体系结构

• 介绍SSM及其在神经网络中的应用，强调其独特性和灵活性。
• 讨论几种著名的SSM架构及其特点，展示SSM在不同任务中的应用。
• 提到其他密切相关的SSM和架构，并将在扩展的相关工作中进一步讨论这些方法。

解释了状态空间模型（SSM）及其在现代神经网络架构中的应用。

1.讲述了 SSM的独立性和端到端性

2.SSNN又可以叫SSM，结构化状态空间神经网络，可以多层处理

介绍了一些著名的SSM架构，如线性注意力、H3结构、Hyena结构、Retnet和RWKV

3 Selective State Space Models 

讨论了论文中的研究流程和主要贡献，尤其是选择机制如何被纳入状态空间模型（SSM），并解决了由此带来的计算挑战。

对每一个章节的介绍

第3.1节和3.2节内容

3.1节：通过合成任务来激励理解选择机制

3.2节：解释如何将选择机制纳入状态空间模型以及引入了时变SSM（动态参数在时间上是变化的）

第3.3节、3.4节和3.5节内容 

3.3节：因为时变SSM不能卷积，改进硬件提出硬件感知算法

3.4节：一种无需注意力机制和MLP块的SSM架构

3.5节：讨论文鲁棒性、可解释性和泛化性，以及对模型进行优化

3.1 Motivation:Selection as a Means of Compression

动机：选择作为一种压缩手段

下面一段探讨了序列建模中的一个基本问题，即如何将上下文信息压缩到较小的状态。这种压缩在不同的序列模型中呈现出不同的权衡和挑战。

序列模型上下文压缩问题

序列建模中的一个重要权衡在于如何有效地压缩和存储上下文信息。

注意力机制和递归模型各有优缺点：

1.注意力机制：能够灵活捕捉上下文，但需要大量内存和计算资源。

2.递归模型：在效率上有优势，但其有限的状态容量可能限制其对长序列的处理能力。

合成任务（为理解压缩推理原理）

解释了为了理解上下文压缩和推理的原理，研究者们使用了两个合成任务的运行示例。

这些任务被设计来测试和展示模型在不同情况下处理和记住上下文信息的能力。

合成任务解释

合成任务的目的和意义

序列模型局限性特殊情况（失效模式）

深入探讨了序列模型的局限性，特别是线性时不变（LTI）模型在处理某些任务时的失效模式，并引出了一种新的构建原则，即选择性。

递归模型和卷积模型的局限性

递归模型：状态更新规则不变，难以根据动态调整

卷积模型：在选择任务中不佳，缺乏内容感知，输入到输出间隔是变化的，无法有效建模

状态压缩的权衡及选择性机制的提出

状态压缩：为了高效，有效处理记住上下文信息。需要较小的状态及包含所有必要信息的状态

选择性机制：为了克服上述局限性，提出选择性机制

序列维度的传播与交互定义

信息在递归时候，通过时间步传播。在卷积模型，通过卷积进行传播

序列的交互通过自注意机制去聚合或通过门控机制决定哪些信息可以传递或丢弃

3.2 Improving SSMs with Selection通过选择改进SSM

解释了如何将选择机制整合到模型中，并指出了这种整合对模型效率的影响。

选择机制的整合

在RNN中动态调整动力学参数，在CNN中动态调整卷积核参数

从时不变到时变性的转换

改为时变性，模型可以动态调整参数，但是失去卷积效率。

描述参数及函数

描述了选择机制中的具体参数化方法，详细说明了如何使用线性投影和激活函数来实现这些选择机制

设计三个线性投影和一个激活函数

一个参数化的线性投影函数，一个广播操作，一个激活函数

四个选择机制，两个将x投影到维度N上得到B和C，两个机制将x投影到维度1，然后广播到D上。

RNN的门控机制GRU如下网址所示

9.1. 门控循环单元（GRU） — 动手学深度学习 2.0.0 documentation

图2 

Figure 2 指的是三个合成任务示例，这些任务展示了不同类型的序列建模问题以及不同模型在解决这些问题时的挑战和要求。

复制任务

选择性复制任务

归纳头任务

算法流程图

这里给出了两个算法，分别是 SSM (S4) 和 SSM + Selection (S6)。它们描述了结构化状态空间模型（SSM）在处理输入序列 𝑥 后生成输出序列 𝑦 的过程

算法1 SSM(S4)算法

算法2 SSM+S6算法

两个算法之间不同之处

3.3 Efficient Implementation of Selective SSMs

引文

讨论了在现代硬件（如GPU）上高效实现选择性状态空间模型（SSM）的挑战和方法

3.3.1解决SSM局限性方法（表达与效率的平衡）

探讨了在实现状态空间模型（SSM）时，如何在表达能力和计算效率之间取得平衡，并介绍了克服现有方法局限性的方法。

1.模型隐状态维度越大，表达能力越强，计算速度越慢。希望在不增加计算速度下，最大化隐藏状态维度

2.循环比卷积更灵活，需要计算具体化潜在状态h的形状

3.引入高效的卷积模式，对卷积核进行计算

4.LTI状态空间模型利用双循环卷积形式，有效维度增到10-100倍

1）递归比卷积更灵活的原因

2）LTI状态空间模型

3.3.2选择性扫描概述:硬件感知状态扩展

详细探讨了克服线性时不变（LTI）模型的局限性以及优化选择性状态空间模型（SSM）的计算问题的具体方法。

优化SSM计算效率的技术

通过观察FLOPs，得到在长序列和较小维度时候，循环模式可能更高效。

使用递归时候，使用更快内存来重复计算。通过HBM到SRAM，在SRAM进行离散化递归计算，减少内存带宽限制。

使用并行扫描算法并行化递归。不进行存储中间状态，而是利用重复计算得到中间状态。

并行扫描算法

如下为详细介绍

CUDA 并行算法Scan、Reduce 图像直方图均衡 - 简书

HBM和SRAM

HBM高带宽存储与传统SRAM、SDRAM的特点_sram与hbm速度-CSDN博客

内存带宽

内存条带宽的换算方法以及DDR4内存理论带宽_ddr4内存带宽计算公式-CSDN博客

卷积为什么比循环要快

LTI模型

SSM模型

SSM模型与LTI模型区别与联系

LTI模型为SSM模型特例，LTI模型为线性时不变，SSM为非线性时变。

3.4简化的SSM架构 

详细描述了选择性状态空间模型（Selective State-Space Models，SSM）如何集成到神经网络架构中，并介绍了一种被称为Mamba的新架构

H3架构，ssm架构一种，线性注意力和MLP交织一起，与GAU相关。

参数集中在线性投影，加入可控扩展因子控制模型维度，SSM参数少。

Mamba架构多个块和规范层残差层链接。

描述了一种简化且高效的架构设计，融合了SSM的优点和现代神经网络的技术，通过合理的参数设置和模块设计，实现了更高效的序列处理能力。Mamba架构的设计灵感来自多种已有的技术，如GAU、SwiGLU和RetNet，并在此基础上进行了创新和优化。

线性注意力

传统注意力机制

线性注意力的原理

elu函数

函数表达式

ELU函数特性

ELU 是一种在深度神经网络中常用的激活函数，它通过在负值区域的平滑指数函数形式，减缓了梯度消失和死亡神经元问题，从而加快了模型的收敛速度。与 ReLU、Leaky ReLU 等激活函数相比，ELU 提供了更稳定的训练过程和潜在的性能提升。

门控注意力单元GAU

用于图像语义分割的GAU与PPM_global attention upsample-CSDN博客

Silu/Swich激活函数

[杂记]激活函数SiLU和Swish的区别-CSDN博客

3.5选择机制的性质

选择机制（Selection Mechanism）是一个广泛的概念，可以应用于多种神经网络架构和参数调控，以下是它在不同情况下的应用示例和解释：

图3

该架构通过将H3块与现代神经网络中的多层感知器（MLP）块结合，简化了SSM（状态空间模型）的设计。

3.5.1与门控机构的连接

探讨了选择机制（Selection Mechanism）在SSM（State Space Models）中的应用，并将其与RNN（Recurrent Neural Network）经典门控机制进行了对比

选择机制和RNN门控机制的关联

对定理1的解释

1.离散化和时间系统关系

可以对连续时间系统进行离散化

2.选择机制的实现操作

通过△，进行选择机制的实现

s△和τ△的区别

3.5.2选择机制的解释

深入探讨了选择机制在SSM（State Space Models）中的应用，并详细解释了它如何通过控制不同参数（如Δ、A、B和C）来实现过滤噪声、筛选上下文和重置边界等功能。

利用改变A、B、C和△来实现多种功能，可变间距，筛选上下文及边界重置操作。

3.6 其他型号模型的细节

探讨了SSM（State Space Models）在使用复数和实数时的表现差异，以及初始化方法对模型性能的影响。

实数和复数在SMM中应用

如何对SSM模型参数进行初始化

选择性SSM的△参数化

详细讨论了选择性SSM（Selective State Space Model）中的Δ参数化方法，并解释了其机制和实验中的应用。

△的初始化

S6模型各种方面优点：选择性S△，单维度到多维度，初始化△参数

4 试验评估

介绍了第4节中的实验和评估内容，涵盖了多个方面的任务和测试，以验证Mamba架构和选择性SSM的性能和效率。

1.合成任务效率

2.下游任务泛化性及表现

3.DNA训练表现及长序列分类任务微调

4.音频训练

5.评估计算效率

6.介绍Mamba架构各个组成部分

选择性复制任务

解释了选择性复制任务在序列建模中的重要性及其设计意图，同时比较了不同模型在该任务中的性能表现。

诱导头任务

表1 

Arch门控机制，layer层类型，model模型架构

图2

展示了“诱导头”（Induction Heads）任务的外推结果。模型在序列长度为256（2^8）上进行训练，并在增加的序列长度上进行测试，从64（2^6）到1048576（2^20）。