正文
然后计算归一化的注意力分数 S 并计算注意力权重。
除了可以执行单个注意力函数,我们还可以执行多头注意力。这让模型可以捕获不同类型的关系,并从多个视角理解输入序列。多头注意力会使用多组自注意力模块并行地处理输入序列。其中每个头都独立运作,执行的计算与标准自注意力机制一样。
之后,将每个头的注意力权重汇聚组合,得到值向量的加权和。这个聚合步骤可让模型使用来自多个头的信息并捕获输入序列中的多种不同模式和关系。
状态空间
状态空间模型(SSM)是一种传统的数学框架,可用于描述系统随时间变化的动态行为。近些年来,人们已将 SSM 广泛应用于控制论、机器人学和经济学等多个不同领域。
究其核心,SSM 是通过一组名为「状态」的隐藏变量来体现系统的行为,使其能有效捕获时间数据的依赖关系。不同于 RNN,SSM 是一种具有关联(associative)属性的线性模型。具体来说,经典的状态空间模型会构建两个关键方程(状态方程和观察方程),以通过一个 N 维的隐藏状态 h (t) 建模当前时间 t 时输入 x 与输出 y 之间的关系。
为了满足机器学习的需求,SSM 必需经历一个离散化过程 —— 将连续参数转变成离散参数。通常来说,离散化方法的目标是将连续时间划分为具有尽可能相等积分面积的 K 个离散区间。为了实现这一目标,SSM 采用的最具代表性的解决方案之一是 Zero-Order Hold(ZOH),其假设区间 Δ = [𝑡_{𝑘−1}, 𝑡_𝑘 ] 上的函数值保持不变。离散 SSM 与循环神经网络结构相似,因此离散 SSM 能比基于 Transformer 的模型更高效地执行推理过程。
离散 SSM 是一个具有结合属性的线性系统,因此可以与卷积计算无缝整合。
RNN、Transformer 和 SSM 之间的关系
图 2 展示了 RNN、Transformer 和 SSM 的计算算法。
一方面,常规 RNN 的运作基于一种非线性的循环框架,其中每个计算都仅依赖于之前的隐藏状态和当前输入。
尽管这种形式可让 RNN 在自回归推理时快速生成输出,但它也让 RNN 难以充分利用 GPU 的并行计算能力,导致模型训练速度变慢。
另一方面,Transformer 架构是在多个「查询 - 键」对上并行执行矩阵乘法,而矩阵乘法可以高效地分配给硬件资源,从而更快地训练基于注意力的模型。但是,如果要让基于 Transformer 的模型生成响应或预测,则推理过程会非常耗时。
不同于仅支持一类计算的 RNN 和 Transformer,离散 SSM 灵活性很高;得益于其线性性质,它既能支持循环计算,也可支持卷积计算。这种特性让 SSM 不仅能实现高效推理,也能实现并行训练。但是,需要指出,最常规的 SSM 是时不变的,也就是说其 A、B、C 和 Δ 与模型输入 x 无关。这会限制其上下文感知型建模的能力,导致 SSM 在选择性复制等一些特定任务上表现不佳。
Mamba
为了解决上述传统 SSM 的缺点,实现上下文感知型建模,Albert Gu 和 Tri Dao 提出了可用作通用序列基础模型主干网络的 Mamba,参阅报道《五倍吞吐量,性能全面包围 Transformer:新架构 Mamba 引爆 AI 圈》。
之后,他们俩又进一步提出了 Mamba-2,其中的结构化空间状态对偶(SSD/Structured Space-State Duality)构建了一个将结构化 SSM 与多种形式的注意力连接起来的稳健的理论框架,让我们可将原本为 Transformer 开发的算法和系统优化技术迁移用于 SSM,也可参阅报道《再战 Transformer!原作者带队的 Mamba 2 来了,新架构训练效率大幅提升》。
Mamba-1:使用硬件感知型算法的选择式状态空间模型
Mamba-1 基于结构化状态空间模型引入了三大创新技术,即基于高阶多项式投影算子(HiPPO)的内存初始化、选择机制和硬件感知型计算。如图 3 所示。这些技术的目标是提升 SSM 的长程线性时间序列建模能力。
具体来说,其中的初始化策略可构建一个连贯的隐藏状态矩阵,以有效地促进长程记忆。
然后,选择机制可让 SSM 有能力获取可感知内容的表征。
最后,为了提升训练效率,Mamba 还包含两种硬件感知型计算算法:Parallel Associative Scan(并行关联扫描)和 Memory Recomputation(内存重新计算)。
Mamba-2:状态空间对偶
Transformer 启发了多种不同技术的发展,比如参数高效型微调、灾难性遗忘缓解、模型量化。为了让状态空间模型也能受益于这些原本为 Transformer 开发的技术,Mamba-2 引入了一个新框架:结构化状态空间对偶(SSD)。该框架在理论上将 SSM 和不同形式的注意力连接到了一起。
本质上讲,SSD 表明,Transformer 使用的注意力机制和 SSM 中使用的线性时不变系统都可被视为半可分离的矩阵变换。