[Sequence Modeling 03] Mamba

要理解 Mamba，我们需要将思维从 Transformer 的“全局注意力”切回到 RNN 的“时序状态机”，但这次，我们要解决 Vanilla RNN 和 LSTM 共同的工程死穴：训练阶段无法并行。

Mamba 的核心目标非常明确：既要拥有 Transformer 级别的高效并行训练能力，又要保持 RNN 的推理复杂度和无限上下文潜力。

它主要通过整合状态空间模型（State Space Model, SSM）、**选择性机制（Selective Mechanism）和底层硬件优化（Hardware-aware Algorithm）**来实现这一目标。

以下从数学和工程的视角详细拆解 Mamba 架构。

01数学基础：从连续到离散的状态空间模型 (SSM)#

Mamba 的前身是 S4 等状态空间模型。SSM 的思想源于控制论，它假设系统的状态演化可以用一组连续的微分方程来描述：

• 是输入。
• 是隐藏状态（类比 RNN 的 或 LSTM 的 ）。
• 是状态转移矩阵，描述系统本身的演化规律。
• 和 是输入输出投影矩阵。

为了在深度学习中使用（处理离散的 token），必须对连续系统进行离散化（Discretization）。通过引入一个步长参数 （Delta），使用零阶保持（Zero-Order Hold）等数学技巧，上述方程可以转换为离散的递归形式：

其中 。 你可以看到，离散化后的公式与 Vanilla RNN 非常相似：当前状态 是前一个状态 的线性变换加上当前输入 的线性变换。因为它是纯线性的（没有 LSTM 那样的 或 Sigmoid 阻断），这为后续的数学化简和并行计算奠定了基础。

1.1为什么需要离散化#

状态空间模型 (SSM) 最初用于控制论，描述的是连续时间内的动态系统（例如电路的电压变化、流体的速度）。连续系统的数学表达是常微分方程 (ODE)：

这里 是一个连续的实数。

但在自然语言处理 (NLP) 中，输入数据是文本 Token。文本在时间维度上是高度离散的（第 1 个 Token，第 2 个 Token……不存在第 1.5 个 Token）。因此，我们无法直接将离散的文本序列输入到连续的微分方程中。

我们需要一种数学映射，把描述连续曲线演化的系统，转换为一步一步跳跃迭代的系统。这就是“离散化”。

为了进行离散化，我们需要定义两个相邻 Token 之间的“虚拟时间间隔”，这就是步长 。

在连续系统中，我们要计算 时刻到 时刻的状态变化，本质上是对微分方程求积分。为了简化积分过程，工程上通常引入零阶保持 (ZOH) 假设：

• ZOH 假设：假设在时间区间 内，输入信号 保持不变（恒定为当前的 Token 向量）。

基于这个假设，解上述线性常微分方程的精确解析解为：

我们将连续时间点映射到离散的步数上（令 , ），并定义：

• (这是矩阵指数)

这样，复杂的连续积分就被转换成了极其简单的离散代数递推公式：

1.2为什么“纯线性”能实现并行计算？#

Vanilla RNN 和离散化后的 SSM 在公式上非常像：

• RNN:
• SSM:

RNN 无法并行的罪魁祸首就是那个非线性激活函数 。因为 的存在，你无法将公式展开。要算 ，必须先算完 ，因为 ，这构成了严格的时序数据依赖。

而 SSM 的离散形式是纯线性的。纯线性的最大优势在于它满足结合律 (Associativity)。我们可以通过代数展开来看：

假设我们要直接计算 （假设初始状态 ）：

• 展开结果:

计算 不再依赖于 和 的中间计算结果了！

只要我们在 GPU 内存中提前算好 的幂次（即 ），我们就可以把序列中所有的 与对应的转换矩阵相乘，然后一次性求和。

在实际的硬件底层（CUDA 层面），这种利用结合律将串行累加转化为树状并行计算的算法被称为 并行前缀和 (Parallel Prefix Sum) 或 并行扫描 (Parallel Scan)。

02Mamba 的核心创新：选择性机制 (Selective SSM)#

传统的 SSM（如 S4）有一个致命缺陷：它的矩阵 是静态的（Linear Time-Invariant, LTI）。这意味着系统对待所有的输入 token（不管它是关键的动词还是无用的标点符号）都在使用一套固定的规则进行特征提取。它缺乏 LSTM 中“门（Gate）”那种根据当前输入动态决定“记住什么、遗忘什么”的能力。

Mamba 通过引入选择性机制（Selective Mechanism）打破了 LTI 假设

它强制让 、 和 成为当前输入 的函数：

这意味着对于序列中的每一个 token：

• （步长）：类似于遗忘门。如果当前输入是无用信息，网络可以输出一个很小的 ，使得 （保留原状态），从而忽略当前输入；如果遇到关键信息，则输出较大的 以更新状态。
• 和 ：动态决定如何将当前输入写入隐藏状态，以及如何从隐藏状态中读取输出。

通过这种方式，Mamba 具备了基于内容进行推理（Content-based Reasoning）的能力，能够像 LSTM 一样过滤信息，解决了传统 SSM 难以处理离散语言信息（如精准复制、忽略填充词）的问题。

03工程魔法：硬件感知与并行扫描 (Parallel Scan)#

既然 Mamba 让 和 随时间变化，它就退化成了一个普通的时变 RNN。由于 依赖于 ，这看似又回到了必须用 for 循环按顺序计算的死胡同。

Mamba 的工程神来之笔在于：利用并行前缀和（Parallel Associative Scan）结合 GPU 内存层级优化，强行把线性的 RNN 依赖并行化。

因为 这个操作满足结合律（Associativity），这使得我们可以像计算并行前缀和（Prefix Sum）一样，通过二叉树状的规约算法在 GPU 的多个线程中同时计算出所有的 。

此外，Mamba 的作者编写了定制的 CUDA Kernel：

1. SRAM 驻留： 避免了在极其缓慢的 GPU HBM（高带宽内存）和计算单元之间来回搬运庞大的隐藏状态。
2. 将离散化（计算 ）和并行扫描融合（Kernel Fusion）在 GPU 极速但容量极小的 SRAM 中一气呵成完成。
3. 最终只将输出 写回 HBM。

这种设计使得 Mamba 在训练阶段无需显式的 for 循环，速度与 Transformer 相当；而在推理阶段，又可以切换回标准的 RNN 步进模式，无需缓存历史 token。

3.1并行前缀和#

假设我们有一个数组 ，我们需要计算它的前缀和数组 ，即 。

输出应该是：

串行计算（Vanilla RNN 的方式）： 必须按顺序计算：。计算 必须等待 完成。对于长度为 的序列，这需要 步（时间步）。

并行计算（利用结合律）： 加法满足结合律：。这意味着我们可以改变计算的优先级，将计算任务分配给 GPU 的多个线程并行执行。

并行前缀和通常通过构建一个二叉树（Binary Tree）的 Reduction 过程来实现：

1. 第一层（并发执行）： 线程 1 算 ，线程 2 算 ，线程 3 算 ，线程 4 算 。
2. 第二层（并发执行）： 线程 1 算 ，线程 2 算 。
3. 第三层： 线程 1 算所有结果的总和。

通过这种树状的合并，长度为 的序列，原本需要 步，现在只需要 步。这就是并行化的本质：用多余的计算单元（GPU 核心）换取时间复杂度的指数级下降。

3.2Mamba 的公式如何套用并行前缀和#

前缀和算的是“加法”，但 Mamba 的公式 既有乘法又有加法，它怎么算前缀和？

这就需要进行代数抽象。前缀和算法不仅适用于加法，它适用于任何满足结合律的二元运算符。

我们将 Mamba 的状态更新公式进行变量替换。令 ，公式变为纯递推形式：

我们定义一个新的操作数元组 。 接下来，我们定义一个自定义的二元运算符 ，用于连接相邻的两个元组 和 ：

为什么这样定义？ 我们来看看连续更新两步会发生什么：

你会发现，状态从 转移到 的等效转移矩阵，恰好就是 的结果。

验证结合律： 纯线性结构保证了这个算子 严格满足结合律：

既然算子 满足结合律，我们就可以完全抛弃 的串行 for 循环，直接把包含矩阵乘法和加法的算子 扔进刚才的二叉树并行前缀和算法中。GPU 的各个线程块会并发地计算不同时间段的等效转移矩阵 ，最终在 的时间步内解算出所有 时刻的隐藏状态 。

04代码与工作流解析#

以下是 Mamba 核心处理模块的伪代码，展示了参数动态生成和离散化的过程（省略了底层的定制化 CUDA Scan 实现，用概念上的接口代替）。先看整体数据流：

import torch
import torch.nn as nn
import torch.nn.functional as F

class MambaBlock(nn.Module):
    def __init__(self, d_model: int, d_state: int, d_conv: int):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        
        # 1D Convolution applied to input before SSM mapping
        # This helps capture local token context (like an n-gram)
        self.conv1d = nn.Conv1d(
            in_channels=d_model, 
            out_channels=d_model, 
            kernel_size=d_conv, 
            padding=d_conv - 1,
            groups=d_model
        )
        
        # Linear projections for data-dependent parameters (The Selective part)
        # B and C depend on the input, mapping from d_model to d_state
        self.x_proj = nn.Linear(d_model, d_state * 2 + 1)
        
        # Projection for Delta (step size), which also depends on input
        self.dt_proj = nn.Linear(1, d_model)
        
        # The A matrix is an internal learned parameter, initialized specifically
        # Shape: (d_model, d_state)
        A = torch.arange(1, d_state + 1).float().unsqueeze(0).repeat(d_model, 1)
        self.A_log = nn.Parameter(torch.log(A)) 
        
        # D parameter for a residual skip connection
        self.D = nn.Parameter(torch.ones(d_model))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x shape: (batch_size, seq_len, d_model)
        batch_size, seq_len, _ = x.shape
        
        # Step 1: Local sequence convolution
        # Transpose for PyTorch Conv1d format: (batch, channels, length)
        x_conv = x.transpose(1, 2)
        x_conv = self.conv1d(x_conv)[:, :, :seq_len] 
        x_conv = x_conv.transpose(1, 2)
        x_conv = F.silu(x_conv) # Activation
        
        # Step 2: Project input to Selective Parameters (B, C, and raw delta)
        # x_proj outputs (batch, seq_len, 2 * d_state + 1)
        x_proj_out = self.x_proj(x_conv)
        
        # Split into delta, B, and C
        delta_raw, B, C = torch.split(
            x_proj_out, 
            [1, self.d_state, self.d_state], 
            dim=-1
        )
        
        # Process delta: apply softplus to make it positive
        delta = F.softplus(self.dt_proj(delta_raw)) # (batch, seq_len, d_model)
        
        # Step 3: Discretization (Zero-order hold approximations)
        A = -torch.exp(self.A_log) # Shape: (d_model, d_state)
        
        # delta * A
        delta_A = torch.einsum('b l d, d n -> b l d n', delta, A)
        # bar_A = exp(delta * A)
        bar_A = torch.exp(delta_A) 
        
        # bar_B = delta * B
        bar_B = torch.einsum('b l d, b l n -> b l d n', delta, B)
        
        # Step 4: The Core SSM computation (Selective Scan)
        # In actual implementation, this part is replaced by a highly optimized 
        # custom CUDA kernel to perform parallel associative scan in SRAM.
        # Here we show the conceptual unrolled loop for clarity.
        
        y = self._conceptual_selective_scan(x_conv, bar_A, bar_B, C)
        
        # Step 5: Add skip connection (D)
        output = y + x * self.D
        
        return output
        
    def _conceptual_selective_scan(self, x, bar_A, bar_B, C):
        """
        Conceptual implementation of the sequential dependency.
        In reality, this is computed in parallel $O(log N)$ using prefix sums.
        """
        batch, seq_len, d_model = x.shape
        d_state = self.d_state
        
        h_t = torch.zeros(batch, d_model, d_state, device=x.device)
        outputs = []
        
        for t in range(seq_len):
            x_t = x[:, t, :] # (batch, d_model)
            
            # Update hidden state: h_t = bar_A * h_{t-1} + bar_B * x_t
            # This is linear!
            h_t = bar_A[:, t, :, :] * h_t + bar_B[:, t, :, :] * x_t.unsqueeze(-1)
            
            # Compute output: y_t = C * h_t
            y_t = torch.einsum('b d n, b n -> b d', h_t, C[:, t, :])
            outputs.append(y_t)
            
        return torch.stack(outputs, dim=1)

05Mamba vs RNN#

事实上，在推理阶段，Mamba 的本质就是一个 RNN。 它们都是处理时序数据的状态机，都具有 的推理复杂度和无限长度的理论上下文。

但如果看它的网络拓扑结构和数学表达，Mamba 和传统的 RNN/LSTM 在架构上有三个根本性的差异。正是这三个差异，不仅让 Mamba 可以使用并行前缀和，也彻底改变了它的特征提取能力。

5.1核心差异一：非线性激活函数的位置（纯线性 vs 非线性递推）#

这是两者在架构上最底层的分歧。

• RNN/LSTM：非线性发生在“循环内部”（Inside the Recurrence）

在 RNN 中，每一个时间步的递推公式是：。

注意这个 ，隐藏状态 在传递给 时，必须穿过一个非线性激活函数。这种设计虽然增加了单步表达能力，但导致了严重的信号衰减（梯度消失），并且彻底破坏了代数上的结合律。

• Mamba：非线性发生在“循环外部”（Outside the Recurrence）

Mamba 的递推公式是：。

这里的递推是严格的纯线性操作，没有 ，没有 Sigmoid。 到 是一条毫无阻挡的线性高速公路。

5.2核心差异二：隐藏状态的维度映射（单通道压缩 vs 多通道展开）#

这决定了模型的“记忆容量”和信息瓶颈。

• RNN/LSTM：低维度的全局状态

假设输入 的特征维度是 （即 d_model，例如 1024）。RNN 通常维护一个维度也为 的隐藏状态 。 所有的 个特征通道的数据，在进行矩阵乘法 时是完全混合在一起的。这就像把所有的信息压缩到一个固定大小的向量里，很容易发生“信息覆盖”。

• Mamba：高维度的结构化状态空间（Structured State Space）

Mamba 引入了状态维度（d_state，通常设为 16）。 对于输入 ，Mamba 并不是把它压缩成一个 维的 。相反，Mamba 为输入特征的每一个通道，都分配了一个独立的 维（d_state）记忆空间。 因此，Mamba 实际的隐藏状态 的形状是 （即 [d_model, d_state]，例如 ）。

并且，Mamba 的状态转移矩阵 是对角矩阵（Diagonal Matrix）。这意味着在递推时，特征的各个通道之间是独立演化的，互不干扰。通道间的交互交由循环外部的其它线性层来完成。这种极高的隐藏维度赋予了 Mamba 远超传统 RNN 的记忆容量。

5.3核心差异三：转移矩阵的动态性（静态权重 vs 数据依赖）#

这是 Mamba 被称为 Selective SSM（选择性状态空间模型）的原因，也是它区别于早期纯线性 RNN 的关键。

• RNN：静态转移矩阵（Static Transition）

RNN 的权重矩阵 在训练完成后就是固定的。无论当前输入的 token 是至关重要的名词，还是毫无意义的标点符号，RNN 总是使用同一个 去乘以前一个状态 。

• Mamba：动态生成转移矩阵（Data-Dependent Transition）

Mamba 中的 和 带有下标 。这意味着对于序列中的每一个 token，网络都会实时计算出一套全新的转移参数。

# Vanilla RNN (Weights are parameters)
# self.W_hh is an nn.Parameter
h_t = torch.tanh(self.W_hh @ h_prev + self.W_xh @ x_t)

# Mamba (Transition matrices are dynamic outputs of the input)
# delta, B, C are generated dynamically via linear layers from x_t
delta_t = F.softplus(self.proj_delta(x_t))
B_t = self.proj_B(x_t)

# A_bar changes for every single token based on delta_t
A_bar_t = torch.exp(delta_t * A_parameter) 

# The recurrence happens with the dynamic A_bar_t and B_t
h_t = A_bar_t * h_prev + (delta_t * B_t) * x_t 

正是这种动态生成 的能力，让 Mamba 拥有了类似于 LSTM 中“遗忘门”的功能：如果 是无关信息，网络会让生成的 趋近于 0，使得 ，从而 ，完美保留了历史记忆，忽略了当前垃圾信息。

5.4总结对比表#