下载掌阅APP，畅读海量书库
立即打开
Mamba的基本架构基于状态空间模型，这是一种数学模型，它需要一个输入序列 x ，并利用学习到的参数 A、B、C 以及一个延迟参数 Δ 来产生输出 y ，用以描述和分析动态系统的行为。这一过程包括离散化参数（将连续函数转换为离散函数）并应用SSM运算，该运算是时不变的，即它不会随着不同的时间步而改变。
此外，还有一种选择状态空间模型（Selective State Space Model），它为参数添加了输入依赖性 B 和 C ，以及一个延迟参数 Δ 。这允许模型有选择地关注输入序列 x 的某些部分。在选择的基础上，对参数进行离散化，并使用扫描操作以随时间变化的方式应用SSM运算，该扫描操作顺序处理元素，并随时间动态调整焦点，如图3-1所示。
图3-1 经典状态空间模型与选择状态空间模型
状态空间模型（经典与选择）在多个领域都有广泛应用，包括控制理论、信号处理、经济学和机器学习等。在深度学习领域，状态空间模型被用来处理序列数据，如时间序列分析、自然语言处理和视频理解等。通过将序列数据映射到状态空间，模型可以更好地捕捉数据中的长期依赖关系。
作为SSM的核心概念，状态空间本质上是一个包含描述系统完整状态所需最小数量变量的集合。它提供了一种以数学方式表达问题的方法，即通过定义系统的所有可能状态来揭示其内在运行机制。以穿越迷宫为例，迷宫中的“状态空间”就好比是一幅详尽的地图，其中每个点代表迷宫中的一个独特位置，这些位置携带着特定的信息，如距离出口的具体距离、周围环境的布局等。
想象你在一个迷宫里，目标是从起点走到终点。在这个迷宫中，每个位置（如你现在站的格子）都可以看作一个“状态”。迷宫的每个格子都可能是一个状态，所以整个迷宫的布局就构成了一个巨大的状态空间。
在迷宫中移动时，你可以向上、向下、向左或向右走，但有时路可能被墙壁阻挡。每次从一个格子移动到另一个格子，就表示状态的改变。状态空间模型用于描述和跟踪你在迷宫中每个可能位置的状态。
简单来说，状态空间模型可以帮助我们理解如何从迷宫的任何一个点到达其他任何点。如果我们能清晰地想象出所有可能的行走路径，找到从起点到终点的路径将变得更加容易。
这就是迷宫问题中的状态空间模型——考虑所有可能的位置和移动，找出最优路径到达目的地。此模型不仅适用于实际的迷宫，也适用于任何需要从一个位置到另一个位置的场景，如图3-2所示。
图3-2 一个寻路的状态空间表示
“状态空间表示”是对系统动态行为的直观且简化的描述。它直观地展示了当前所处的位置（即当前状态）、未来可能到达的地点（潜在的未来状态），以及导致这些状态转移的具体变化（如向右或向左的移动）。
尽管状态空间模型依赖于复杂的方程和矩阵来精确追踪这种行为的演变，但其核心目的仍在于提供一种有效的方法来追踪位置、潜在路径以及到达目标的方式。在状态空间模型中，描述系统状态的变量（例如，在迷宫中可能是 X 和 Y 坐标以及到出口的距离）被抽象为“状态向量（state vectors）”。这些向量不仅捕捉了当前状态的关键信息，还揭示了状态之间的转换关系。
类似地，在自然语言处理领域，语言模型中的嵌入向量也扮演着描述输入序列“状态”的重要角色。这些嵌入向量通过捕捉单词或短语的语义信息，将文本数据转换为能够被神经网络理解和处理的高维空间中的点。通过这种方式，语言模型能够像状态空间模型一样，捕捉文本序列中的长期依赖关系，并据此生成准确且连贯的输出。因此，状态空间表示和语言模型中的嵌入向量，虽然在具体实现和应用上存在差异，但它们的本质都是对系统或数据状态的一种抽象和表示。例如，当前位置的向量（状态向量）如图3-3所示。
图3-3 迷宫中的状态向量
就神经网络而言，系统的“状态”通常是其隐藏状态，在大型语言模型的背景下，这是生成token的重要表示之一。
进一步来说，如果我们假设动态系统的状态空间模型（例如在3D空间中移动的物体），并用两个核心方程来预测其行为，这些方程基于系统在时间 t 的状态进行建模。这些方程通常包括一个状态转移方程和一个观测方程。
状态转移方程描述了系统从一个时间步到下一个时间步的状态变化。在3D空间中移动的物体，状态可能包括物体的位置（ x、y 坐标）、与出口的距离（distance）等。状态转移方程捕获了这些状态变量如何随时间变化，通常包括一些表示系统动态（如物理定律）的参数。一个简单的例子（不考虑外部力和加速度的简化情况）如下：
其中，( x , y ) t 和d t 分别是物体在时间 t 的位置和距离出口的距离，Δ t 是时间步长。
观测方程，又称为输出方程，描述了如何从系统的状态变量得到我们实际可以观测到的数据。在实际情况中，由于传感器噪声、测量误差等因素，我们观测到的数据可能与系统的真实状态存在偏差。为了模拟这种偏差，观测方程通常包括一个噪声项。
在状态空间中，如果我们有一个能够测量物体位置的传感器，观测方程可能会简单地表示为物体位置的带噪声版本：
z t = x t +noise
其中， z t 是我们在时间 t 时刻观测到的位置， x t 是物体的真实位置，noise是表示测量误差的随机变量。
通过这两个方程，我们可以使用“状态空间模型”来模拟“状态空间的转移和输出”，从而为我们提供一个框架，用以理解和预测动态系统的行为，如图3-4所示。
图3-4 状态方程与输出方程的整合
在给定系统当前状态（以及可能的观测数据）的情况下，我们可以使用状态转移方程来预测系统的未来状态，并使用观测方程来验证预测结果是否与观测数据一致。
状态空间模型（SSM）是一种功能强大的统计工具，旨在精准地描述和预测随时间演进的复杂系统的动态状态。通过将状态空间模型的核心思想与深度学习技术相结合，开创了一种前沿的神经网络架构，这一架构不仅继承了Transformer的高效并行训练能力，同时保持了RNN的线性推理速度，从而在性能上实现了质的飞跃。
状态空间模型用于描述系统状态的表示，并根据某些输入预测系统的下一个状态。在时间 t 时刻的状态下：
● 输入序列 x ( t )：例如，在迷宫中向左和向下移动。
● 潜在状态的表示 h ( t )：例如，距离出口的 x / y 坐标。
● 预测输出序列 y ( t )：例如，再次向左移动以更快到达出口。
然而，状态空间模型在处理系统动态时，并不局限于离散（discrete）序列事件（如向左移动一次）。相反，它采用了更为灵活和泛化的方法，即将连续序列（continuous sequence）作为输入，并据此预测输出序列。
这种方法使得状态空间模型能够处理更为复杂和细微的系统变化。例如，在迷宫问题中，状态空间模型不仅可以考虑“向左”或“向右”的离散动作，还可以处理诸如“以0.5的速度向左上方移动”这样的连续动作。通过捕捉这种连续变化，状态空间模型能够更精确地模拟系统的动态行为，并预测出更为准确的未来状态。
同样地，在自然语言处理中，连续序列的概念也尤为重要。与传统的基于离散词汇的模型相比，基于连续嵌入向量的语言模型能够更好地捕捉词汇之间的语义关系，并生成更为流畅和自然的文本。通过将文本序列视为连续空间中的点，这些模型能够利用深度学习的强大能力，从海量数据中学习出词汇和句子之间的复杂依赖关系，如图3-5所示。
图3-5 序列中的状态空间模型
它不使用离散序列（如向左移动一次），而是将连续序列作为输入并预测输出序列。因此，通过将连续序列作为输入和预测目标，状态空间模型不仅在理论上具有更高的泛化能力，而且在实际应用中能够处理更为复杂和细微的系统变化，为我们提供了更为精确和有效的建模工具。
根据我们在前面介绍的状态空间模型，对于完整的状态空间模型，我们可以通过假设一些可以与深度学习进行结合的模型公式来完成设计。具体如下：
● 状态方程： h '( t )= Ah ( t )+ Bx ( t )。
这一方程描述了系统状态 h 随时间 t 的演变规律，其中 A 是状态转移矩阵， B 是输入矩阵， x ( t )是外部输入。通过该方程，我们可以根据当前状态 h ( t )和外部输入 x ( t )来预测下一时刻的状态 h '( t )。
● 输出方程： y ( t )= Ch ( t )+ Dx ( t )。
这一方程建立了系统状态 h ( t )与外部观测 y ( t )之间的联系，其中 C 是输出矩阵， D 是直接传输矩阵。通过该方程，我们可以根据系统状态 h ( t )和外部输入 x ( t )来计算观测值 y ( t )，从而验证我们的状态预测是否准确。
因此，我们的目标是求解这些方程，之后便可以根据统计原理和观察到的数据（输入序列和先前状态）预测系统的状态，目标是找到从输入序列到输出序列的状态表示 h ( t )。状态空间模型的核心方程如图3-6所示。
图3-6 SSM核心公式
接下来，我们将逐步探讨状态空间模型中的各个组件如何影响学习过程。
状态方程在状态空间模型中扮演着至关重要的角色。它详细描述了系统状态如何随时间变化，以及外部输入如何影响这些状态。具体来说，状态方程包含以下两个部分。
● 状态转移矩阵 A ：这个矩阵描述了在没有外部输入的情况下，系统状态如何随时间自然演变。它捕捉了系统内部的动态特性，如惯性、阻尼等。在连续时间系统中，状态转移矩阵 A 通常与系统的微分方程或差分方程相关。
● 输入矩阵 B ：这个矩阵描述了外部输入如何影响系统状态。它建立了输入信号与系统状态之间的直接联系，使得我们可以通过控制输入来影响系统的行为。在控制系统中，输入矩阵 B 通常与系统的控制策略或外部激励相关。
结合这两个部分，状态方程可以表示为：
状态转换方程： h '( t )= Ah ( t )+ Bx ( t )。
其中， h ( t )是系统的状态向量， h '( t )是状态向量的导数（在连续时间系统中）或差分（在离散时间系统中）， x ( t )是外部输入向量， A 是状态转移矩阵， B 是输入矩阵，状态方程的变换如图3-7所示。
图3-7 状态方程的变换
其中， h ( t )指的是任何给定时间 t 的潜在状态表示，而 x ( t )指的是某个输入。
输出方程：
y ( t )= Ch ( t )+ Dx ( t )
该方程描述了状态如何转换为输出（通过矩阵 C ），以及输入如何影响输出（通过矩阵 D ），如图3-8所示。
图3-8 输出方程变换
在深度学习中，矩阵 A 、 B 、 C 和 D 通常也被称为可学习的参数。
下面用Python代码来模拟一个简单的状态空间模型。假设我们有一个二维状态向量 h ( t )=[ h 1( t ), h 2( t )]，以及一个一维输入 x ( t )。我们将使用简单的矩阵 A 和 B 来描述系统的动态。
以下是一个使用NumPy实现的简单状态转移函数的例子。该函数根据给定的状态转移矩阵 A 、输入矩阵 B 、当前状态 h 和外部输入 x 来计算下一个状态（或状态的导数，在连续时间系统中）。请注意，这个函数为连续时间系统设计的，并假设 h '( t )代表状态向量的导数。对于离散时间系统，需要实现一个不同的状态转移函数，可能需要使用矩阵指数或其他方法。
import numpy as np
def state_transition(A, B, h, x):
"""
计算下一个状态（或状态的导数），给定当前状态、输入和时间步长
参数:
A (np.ndarray): 状态转移矩阵
B (np.ndarray): 输入矩阵
h (np.ndarray): 当前状态向量
x (float or np.ndarray): 外部输入，如果是标量，则直接用于所有状态；如果是向量，则与
B相乘
dt (float, optional)
返回:
np.ndarray: 下一个状态（或状态的导数）的近似值
"""
# 如果x是标量，则将其转换为与B列数相同的向量
if np.isscalar(x):
x = np.array([x] * B.shape[1])
# 计算状态的导数（或下一个状态的近似值，取决于dt的使用）
h_prime = np.dot(A, h) + np.dot(B, x)
# 在连续时间系统中，我们通常计算状态的导数h'(t)
# 如果要近似下一个状态h(t+dt)，可以使用一个简单的欧拉方法
# h_next = h + dt * h_prime
# 但这里我们仅返回状态的导数
return h_prime
# 示例使用
# 定义状态转移矩阵A和输入矩阵B
A = np.array([[0, 1], [-2, -3]])
B = np.array([[0], [1]])
# 初始状态向量
h = np.array([1, 0])
# 外部输入，假设是一个标量
x = 2
# 调用状态转移函数
h_prime = state_transition(A, B, h, x)
print("状态的导数 h'(t):",h_prime)
# 如果要计算下一个状态的近似值（使用欧拉方法），可以这样做
dt = 0.01 # 时间步长，用于近似连续时间系统的状态转移。默认为0.01
h_next = h + dt * state_transition(A, B, h, x)
print("下一个状态的近似值 h(t+dt):",h_next)
打印结果如下：
状态的导数 h'(t): [0 0] 下一个状态的近似值 h(t+dt): [1. 0.]
回到对状态方程的可视化处理，结合以下两个方程：
状态转换方程： h '( t )= Ah ( t )+ Bx ( t )。
输出方程： y ( t )= Ch ( t )+ Dx ( t )。
我们将在同一个图形中整合这两个方程，如图3-9所示。
图3-9 空间状态模型整合
接下来，让我们更深入地探讨这些矩阵如何共同影响状态空间模型中的学习过程。想象一下，我们有一些输入信号，这些信号首先通过一个名为矩阵 B 的“转换器”。矩阵 B 的作用就像是一个翻译者，它告诉我们这些输入信号如何影响系统的状态，如图3-10所示。
图3-10 矩阵 B 的作用
当我们得到这些输入信号的影响后，它们与当前的系统状态相结合。但在此之前，我们需要了解系统状态是如何随时间演变的。这就是矩阵 A 发挥作用的地方。读者可以把矩阵 A 想象成一个连接所有内部状态的“纽带”，它描述了系统内部状态之间的相互作用和动态变化。当我们把当前状态与矩阵 A 相乘时，就得到了对系统未来状态变化的预测，如图3-11所示。
图3-11 矩阵 A 与 B 结合计算
矩阵 A 在状态更新之前被应用，它的计算结果与输入信号（经过矩阵 B 转换后）相加，共同决定了新的系统状态。这个过程就像是系统的“大脑”，根据过去的经验和当前的输入来预测未来的状态。
然后，当系统状态得到了更新时，我们想知道这个状态如何转换为可以观察到的输出。这就是矩阵 C 的任务。矩阵 C 就像一个“解码器”，它将系统的潜在状态转换为我们可以理解的输出信号，如图3-12所示。
图3-12 将转换结果输入矩阵 C
此外，矩阵 D 提供了一个从输入直接到输出的“快速通道”。这就像在神经网络中的跳跃连接（skip-connection），允许信息绕过一些处理层直接传播到输出。在SSM中，矩阵 D 允许输入信号直接对输出产生影响，而不需要经过系统状态的转换。
前面我们讲解了残差神经网络（ResNet）是跳跃连接的经典应用，它通过允许信息直接传播到网络的更深层，帮助解决了深度神经网络中的梯度消失问题。在SSM中，矩阵 D 也起到了类似的作用，允许输入信号直接对系统输出产生影响，如图3-13所示。
图3-13 矩阵 D 的作用
为了简化问题，我们可以暂时不考虑矩阵 D ，只关注矩阵 A 、 B 和 C 。通过这两个方程，我们可以根据观测数据来预测系统的状态，并根据系统状态来生成输出。由于输入信号是连续的，因此SSM通常采用连续时间表示来描述系统的动态变化，如图3-14所示。
图3-14 完整的状态空间模型
总的来说，矩阵 A 、 B 、 C 和 D 在SSM中扮演着不同的角色，它们共同决定了系统如何根据输入信号更新状态，并根据状态来生成输出。通过调整这些矩阵的参数，我们可以改变系统的行为，使其更好地适应不同的任务和环境。
最后，我们将如图3-14所示的完整空间状态模型进行简化，最终得到如图3-15所示的简化版。
图3-15 简化后的完整状态空间模型
之后，我们像之前一样更新原始方程来表示每个矩阵的用途，如图3-16所示。
图3-16 可更新的完整状态空间模型
这两个方程旨在根据观测数据预测系统的状态。由于输入是连续的，SSM的主要表示是连续时间表示（continuous-time representation）。因此，这就需要我们对SSM进行离散化，通过使用特殊的第4个参数Δ，将连续参数 A 、 B 和 C 转换为离散参数来表示SSM。
新状态转换方程：
。
新输出方程： y ( t )= Ch ( t )。
下面我们将继续介绍将连续信号转换成离散信号的方法。