Mamba
要理解 Mamba,我们需要将思维从 Transformer 的“全局注意力”切回到 RNN 的“时序状态机”,但这次,我们要解决 Vanilla RNN 和 LSTM 共同的工程死穴:训练阶段无法并行。 Mamba 的核心目标非常明确:既要拥有 Transformer 级别的高效并行训练能力,又要保持 RNN $O(1)$ 的推理复杂度和无限上下文潜力。 它主要通过整合状态空间模型(State Space Model, SSM)、**选择性机制(Selective Mechanism)和底层硬件优化(Hardware-aware Algorithm)**来实现这一目标。 以下从数学和工程的视角详细拆解 Mamba 架构。 数学基础:从连续到离散的状态空间模型 (SSM) Mamba 的前身是 S4 等状态空间模型。SSM 的思想源于控制论,它假设系统的状态演化可以用一组连续的微分方程来描述: $$h'(t) = A h(t) + B x(t)$$$$y(t) = C h(t)$$ $x(t)$ 是输入。 $h(t)$ 是隐藏状态(类比 RNN 的 $h_t$ 或 LSTM 的 $c_t$)。 $A$ 是状态转移矩阵,描述系统本身的演化规律。 $B$ 和 $C$ 是输入输出投影矩阵。 为了在深度学习中使用(处理离散的 token),必须对连续系统进行离散化(Discretization)。通过引入一个步长参数 $\Delta$(Delta),使用零阶保持(Zero-Order Hold)等数学技巧,上述方程可以转换为离散的递归形式: $$h_t = \bar{A} h_{t-1} + \bar{B} x_t$$$$y_t = C h_t$$其中 $\bar{A} = \exp(\Delta A)$。 你可以看到,离散化后的公式与 Vanilla RNN 非常相似:当前状态 $h_t$ 是前一个状态 $h_{t-1}$ 的线性变换加上当前输入 $x_t$ 的线性变换。因为它是纯线性的(没有 LSTM 那样的 $\tanh$ 或 Sigmoid 阻断),这为后续的数学化简和并行计算奠定了基础。 ...