深度学习新架构 Mamba
论文介绍 Mamba: Linear-Time Sequence Modeling with Selective State Spaces
关注微信公众号: DeepGoAI
我的项目地址：https://github.com/state-spaces/mamba（曾经 6.3k+）
论文地址：https://arxiv.org/abs/2312.00752
本文介绍了一种新的序列模型架构，名为 Mamba，它通过选择性状态空间模型（Selective State Space Models, SSMs）来改良传统的状态空间模型。Mamba 通过输出依赖的形式调整 SSM 参数，容许模型依据以后的数据选择性地传递或忘记信息，从而解决了以前模型在解决离散和信息密集型数据（如文本）时的有余。此外，只管这种扭转使得模型不能应用高效的卷积计算，研究者设计了一种硬件感知的并行算法，以递归模式运行，使得 Mamba 在推理速度上比传统的 Transformer 快 5 倍，并且在序列长度上实现线性缩放。在语言模型、音频和基因组数据模型等多个畛域，Mamba 都获得了最先进的性能。
上图提供了结构化状态空间模型（SSM）如何解决输出数据的概述。具体来说，它阐明了输出 $x$ 的每个通道是如何通过更高维度的潜在状态 $h$ 独立映射到输入 $y$ 的。例如，如果输出有 5 个通道（$D = 5$），则通过维度为 $N = 4$ 的潜在状态解决它们。先前实现的 SSM 因为其在解决大无效状态时的效率而被留神到，这个大无效状态是维度 $D$ 和 $N$，以及批量大小 $B$ 和序列长度 $L$ 的乘积。这种效率是通过防止这个大状态的实体化（materialization of the large state）并改用依赖工夫不变性的代替计算门路来实现的——意味着参数（$\Delta$，$A$，$B$，$C$）随工夫放弃不变。文本提到了一个“抉择机制”，在此上下文中引入了输出依赖的动静。这意味着参数当初可能会依据输出变动，引入了更动静且可能更精确的工夫数据建模。然而，这也须要一个“审慎的硬件感知算法”，因为减少的复杂性和计算要求。为确保模型能够在特定的硬件配置上无效执行而做的优化或思考，例如提到的 GPU 及其 SRAM 和 HBM（高带宽内存）。
办法概述
结构化状态空间序列模型（S4）
结构化状态空间序列模型（Structured State Space Sequence Models，简称 S4）是一种新兴的深度学习序列模型，它与循环神经网络（RNNs）、卷积神经网络（CNNs）及经典的状态空间模型相关联。这些模型受到了一种将一维函数或序列映射到另一个一维函数或序列通过隐式潜在状态的特定间断零碎的启发。
具体来说，S4 模型应用四个参数（$\Delta A$, $B$, $C$）定义序列到序列的转换过程，分为两个阶段：
第一阶段（离散化）：将“间断参数”（$\Delta A$, $B$）转换为“离散参数”（$\bar{A}$, $\bar{B}$）。
第二阶段：通过离散化后的参数计算序列转换，能够通过线性递归或全局卷积两种形式实现
S4 模型 外围公式
首先形容 S4 模型的外围动静，即如何从一个工夫步骤到下一个工夫步骤更新潜在状态，并如何从这个潜在状态生成输入序列：
公式(1a): $h'(t) = A h(t-1) + B x(t)$
这个公式形容了潜在状态 $h(t)$ 的更新规定，其中 $A$ 是状态转换矩阵，$B$ 是输出到状态的转换矩阵，$x(t)$ 是以后工夫步的输出。这表明以后的潜在状态是由前一个工夫步的潜在状态和以后输出的组合决定的。
公式(1b): $y(t) = C h(t)$
这个公式形容了如何从潜在状态 $h(t)$ 生成输入 $y(t)$，其中 $C$ 是状态到输入的转换矩阵。这表明输入间接依赖于以后的潜在状态。
进一步，作者提供这两个公式 (1a) 和(1b)的另一种示意模式，强调了模型的递归个性：
公式(2a): $h_t = \bar{A} h_{t-1} + \bar{B} x_t$
这是对 (1a) 的简化示意，强调了潜在状态是通过递归形式更新的，其中 $h_{t-1}$ 代表前一时间步的潜在状态。
公式(2b): $y_t = C h_t$
这是对 (1b) 的简化示意，阐明了输入 $y_t$ 是以后潜在状态 $h_t$ 的间接函数。
以下两个公式提供了通过全局卷积实现序列到序列转换的视角：
公式(3a): $\bar{K} = (C\bar{B}, C\bar{AB}, \ldots, C\bar{A}^{k}\bar{B}, \ldots)$
公式(3b): $y = x \ast \bar{K}$
公式 (3a) 和(3b) 展现了如何通过状态空间模型将输出转换成输入。公式 (3a) 形容了通过状态转换矩阵 $\bar{A}$ 对输出矩阵 $\bar{B}$ 的间断利用，联合输入矩阵 $C$，结构出核 $\bar{K}$。这个过程无效地建设了一系列利用于输出的转换。公式 (3b) 展现了输出 $x$ 与核 $K$ 的卷积，阐明了如何综合变换后的输出产生输入 $y$。这种公式化的目标是为了模仿动静零碎在将输出转换为输入时的行为，捕获工夫依赖性和零碎参数对输入的影响。
离散化
首先将“间断参数”（$∆$、$A$、$B$）转换为“离散参数”（$\bar{A}$、$\bar{B}$）的过程，这一转换通过固定的公式实现，例如通过零阶放弃（ZOH）规定，能够失去：
$[A = e^{\Delta A}, \quad B = (\Delta A)^{-1}(e^{\Delta A} – I) \cdot \Delta B]$
这种离散化与间断工夫零碎有粗浅的分割，它为模型赋予了额定的属性，如分辨率不变性和主动的模型正则化等。此外，它还与 RNNs 的门控机制无关，但从机械角度来看，离散化能够简略视为 SSM 在前向流传中的计算图的第一步。
计算
参数从（$∆$、$A$、$B$、$C$）转换为（$\bar{A}$、$\bar{B}$、$C$）后，模型能够以两种形式计算：线性递归或全局卷积。模型通常在训练时应用卷积模式（以便于并行化解决），在自回归推理时切换到递归模式（逐工夫步解决输出）。
线性工夫不变性（LTI）
这些方程的一个重要属性是模型的动静性质在工夫上是恒定的，即 $∆$、$A$、$B$、$C$（包含 $\bar{A}$、$\bar{B}$）对所有工夫步都是固定的。这种性质称为线性工夫不变性，与递归和卷积深度相关联。从直观上讲，LTI SSMs 等价于任何线性递归或卷积，并应用 LTI 作为这些模型类别的统称。
然而，所有的结构化 SSMs 都是 LTI 的，即通过卷积计算，这是因为根本的效率束缚。但本介绍的工作的一个外围洞见是，LTI 模型在模仿某些类型的数据时存在根本性限度，作者的技术奉献包含在克服效率瓶颈的同时移除 LTI 束缚。
构造和维度
高效计算结构化 SSMs 也须要对 $A$ 矩阵施加构造，最常见的构造模式是对角线。在这种状况下，$A \in \mathbb{R}^{N \times N}, \quad B \in \mathbb{R}^{N \times 1}, \quad C \in \mathbb{R}^{l \times N}
$ 矩阵能够用数字示意。为了解决具备通道数的输出序列，SSM 独立利用于每个通道。在这种状况下，总的暗藏状态具备每个输出的维度，其计算跨序列长度须要的工夫和内存是效率瓶颈的本源。
Selective State Space Models
文中进一步提出一种改良的状态空间模型（SSM），这种模型通过引入抉择机制来加强其对输出序列的解决能力。这个抉择机制可能容许模型依据输出的相关性来动静调整其关注点，从而进步解决长序列数据的效率和准确性。
上图形容了几个不同的工作，用于展现工夫不变模型如线性递归和全局卷积的局限性，以及选择性模型在解决这些限度方面的劣势。这些工作包含：规范复制工作（左图）：这个工作在输出和输入元素之间有恒定的间距，能够通过工夫不变模型，如线性递归和全局卷积轻松解决。
选择性复制工作（右上图）：这个工作在输出之间有随机间距，须要工夫变动模型可能依据内容选择性地记住或疏忽输出。
演绎头工作（右下图）：这是一种基于上下文检索答案的关联回顾示例，对于大型语言模型来说是一个要害能力。
从中能够发现在设计序列模型时，引入可能依据输出序列的内容和上下文进行动静调整的机制的重要性。这种机制不仅进步了模型在解决简单序列时的灵活性和效率，而且对于晋升模型在执行基于上下文的高级认知工作方面的能力也是至关重要的。
算法 1 和算法 2 展现了作者应用的次要抉择机制。它们的要害区别在于，使得一些参数（$\Delta$, $B$, $C$）成为输出的函数，并随同整个过程中张量形态的相干变动。具体而言，这些参数当初具备长度维度，意味着模型从工夫不变（time-invariant）变为工夫变动（time-varying）。这一变动失去了与卷积（公式 3）的等效性，对其效率有所影响。
算法 1（SSM（S4））：展现了一个结构化的状态空间模型（SSM）的构建过程，它不包含抉择机制。这是一个工夫不变模型，能够通过递归或卷积的形式实现。
算法 2（SSM + 抉择（S6））：在算法 1 的根底上增加了抉择机制，通过使 $\Delta$, $B$, $C$ 成为输出的函数，引入了工夫变动个性。这种机制容许模型依据输出动静调整其行为，使其可能更灵便地解决简单序列数据。
通过引入抉择机制，算法 2 可能针对每个工夫步的特定输出调整模型参数，从而使模型可能依据上下文内容进行适应性变动。这种工夫变动的个性使得模型可能在序列建模中更无效地压缩上下文信息，进步了解决长序列和简单依赖关系的能力。
改良状态空间模型
通过使模型的要害参数（如 RNN 的循环动静或 CNN 的卷积核）成为输出的函数，能够实现抉择机制的集成。这种办法使得模型可能依据输出动静调整其行为，进而在序列中更无效地流传或疏忽信息。
具体来说，算法 1 和算法 2 展现了采纳的次要抉择机制。要害的变动在于，通过使某些参数（$\Delta$、$B$、$C$）成为输出的函数，以及相干的张量尺寸变动，模型从工夫不变转变为工夫变动。这意味着，与应用固定卷积核的模型相比，该办法容许模型依据输出内容的不同动静调整其外部参数，从而在解决序列数据时提供更大的灵活性。然而，这也意味着与卷积操作的等效性丢失，对模型的效率产生影响。
作者特地抉择将参数化投影到维度 $d$ 的函数作为 $\Delta B$ 和 $\Delta C$，因为它与 RNN 的门控机制相干。通过使模型可能基于输出内容调整其状态转换的动态性，抉择机制进步了模型的灵活性和解决序列数据的能力。
这些形容展现了 SSMs 在架构设计方面的翻新和多样性，每种办法都试图通过不同的技术改良来进步模型的性能和效率。这里介绍了它是如何通过在 SSM 之间引入门控连贯和部分卷积来加强模型解决能力的。这些架构的共同点在于它们都试图通过结构化的办法来捕捉和解决序列数据的依赖关系，同时进步计算效率。
总结
这篇文章次要围绕结构化状态空间模型（Structured State Space Models, SSMs）开展，奉献能够从以下几个方面进行总结：
- 提出新的架构和办法：文章通过介绍不同的 SSM 架构（如 H3、Hyena、RetNet）和办法（如线性注意力），展现了在解决序列数据和工夫序列预测方面的翻新。这些办法试图通过特定的设计来进步模型的计算效率和性能，特地是在解决大规模数据集时。
- 硬件感知算法：文章强调了为实现高效计算而采纳的硬件感知算法的重要性。通过思考特定硬件（如 GPU）的个性，文章探讨了如何优化算法来利用这些硬件的高性能计算能力。这种办法有助于减速模型的训练和推理过程，使得简单的模型在理论利用中更加可行。
- 输出依赖的动静：与传统的 SSMs 不同，这篇文章通过引入抉择机制来增加输出依赖的动静，从而容许模型捕捉和利用输出数据中的工夫变动信息。这种动态性的引入旨在进步模型对工夫序列数据的了解和预测能力，使模型可能更好地适应数据的变动和发展趋势。
- 进步模型性能和效率：通过上述翻新，文章展现了如何在放弃或进步模型性能的同时，缩小计算和存储需要。这对于解决大规模数据集和长时间序列特地重要，因为这些场景通常须要大量的计算资源和工夫。
总的来说，这篇文章通过引入新的架构设计、硬件感知算法优化和输出依赖的动静，为进步结构化状态空间模型在序列数据处理和工夫序列预测方面的性能和效率做出了奉献。这些翻新使得 SSMs 在面对大规模和简单数据时更加高效和无效，拓宽了它们在不同畛域利用的可能性。
鉴于篇幅，试验和其余更多细节请参阅论文原文
关注微信公众号: DeepGoAI
本文由 mdnice 多平台公布