【大模型】Mamba模型架构
前言:有望替代Transformer进行序列处理的算法框架。
目录
- 1. 简介
- 2. 状态空间模型SSMs
- 3. Mamba
1. 简介
Mamba = 有选择处理信息 + 硬件感知算法 + 更简单的SSM架构
和transformer对比:
2. 状态空间模型SSMs
SSM是一种线性模型,具有关联属性。SSM结构如下所示:
简化表示为:
其中,A、B、C、D四个矩阵是可学习参数;x(t)是前一个状态
步骤:
(1)核心:
* 状态方程:描述了系统状态随时间的演化
* 输出方程:描述了系统的输出如何依赖于系统状态和控制输入
(2)离散化:将连续的参数离散化处理,零阶保持技术
(3)卷积计算:作为线性系统,离散SSM具有关联属性,因此可以无缝集成到卷积计算中。更具体地说,它可以独立地计算每个时间步的输出。
序列模型区别:
(1)RNNs在非线性递归框架内运作,便于在自回归推理过程中快速输出。
(2)Transformers在多个查询-键对上并行执行矩阵乘法,便于并行训练。
(3)SSMs通过线性属性适应递归和卷积计算,融合了RNNs和Transformers的优势,允许SSMs进行递归推理和并行训练。
传统的时间不变SSM在上下文感知建模方面不足,即A、B、C的静态导致内容感知方面的问题,导致在特定任务中的性能下降:
(1)选择性复制:目标是复制输入的部分内容并按顺序输出它们。由于(循环/卷积)SSM 是线性时间不变的,因此在这项任务中表现不佳。正如我们之前看到的,对于 SSM生成的每个 token,矩阵 A、B 和 C 都是相同的。因此,由于固定的 A、B 和 C 矩阵,SSM 无法执行内容感知推理,因为它对每个 token 都一视同仁。
(2)归纳头:目标是重现输入中发现的模式。由于 SSM 是时间不变的,它无法选择从历史中回忆哪些之前的 token。
3. Mamba
HiPPO:主要为了解决如何在有限的存储空间中有效地解决序列建模的长距离依赖问题。HiPPO 通过函数逼近产生状态矩阵 A 的最优解,有效的解决了长距离依赖问题。
为了解决时间不变SSM的问题,提出了选择性SSM(S6或Mamba)。通过让 SSM 的矩阵 A、B、C 依赖于输入数据,从而实现了选择性。这意味着模型可以根据当前的输入动态地调整其状态,选择性地传播或忽略信息。
由于 A、B、C 都是输入相关了,不再是线性时间不变系统,也就失去了卷积的性质,不能用 FFT来进行高效训练了,Mamba 作者采用了一种称为硬件感知的算法解决这个问题(内核融合,并行扫描,重计算)