计算机视觉基础与实践

从“注意力”到“状态空间”：探索序列建模的新范式

发布日期: 2026年02月05日 作者: 杨其臻 阅读时间: 12分钟

摘要

本文探讨了超越Transformer的序列建模新方法——状态空间模型（SSM），特别是Mamba架构。我们将解释其核心原理、如何克服注意力机制的固有缺陷，以及在长序列处理上的潜力，为读者提供一个理解下一代序列模型的基础视角。

引言：Transformer的瓶颈

Transformer及其核心的注意力机制无疑是过去十年AI领域最成功的架构之一，驱动了从GPT到BERT等一系列革命性模型。然而，随着我们对模型能力要求的不断提高，其固有缺陷也日益凸显。

注意力机制的主要挑战在于：

  • 二次复杂度：计算输入序列中所有元素两两之间的关联，导致时间和内存开销随序列长度呈平方级（O(n²)）增长。
  • 上下文窗口限制：处理超长文本（如整本书、长视频）时极为低效，通常需要复杂的工程技巧进行“分块”。
  • ：传统注意力在推理时对所有输入“一视同仁”，缺乏根据内容动态调整信息流的能力。

这些瓶颈催生了研究者对“后Transformer”时代架构的探索，而状态空间模型（State Space Models, SSMs），特别是其最新代表Mamba，正成为最有希望的竞争者之一。

什么是状态空间模型？

状态空间模型并非全新概念，它起源于控制理论和信号处理领域，用于描述动态系统的输入、输出和内部状态之间的关系。近年来，研究者将其引入深度学习，用于序列建模。

一个连续的线性时不变（LTI）SSM可以用以下方程描述：

\[ \begin{aligned} h'(t) &= \mathbf{A}h(t) + \mathbf{B}x(t) \\ y(t) &= \mathbf{C}h(t) + \mathbf{D}x(t) \end{aligned} \]

其中：

  • \( x(t) \) 是输入信号（如一个词嵌入随时间的变化）。
  • \( h(t) \) 是系统的隐藏状态，它像一个记忆单元，汇总了到当前时刻为止的历史信息。
  • \( y(t) \) 是输出信号。
  • \( \mathbf{A}, \mathbf{B}, \mathbf{C}, \mathbf{D} \) 是可学习的参数矩阵。

直观理解：SSM将序列处理看作一个动态系统。输入 \( x(t) \) 不断更新内部状态 \( h(t) \)（由矩阵 \( \mathbf{A} \) 和 \( \mathbf{B} \) 控制），而输出 \( y(t) \) 则由当前状态 \( h(t) \) 和当前输入 \( x(t) \) 共同决定（由矩阵 \( \mathbf{C} \) 和 \( \mathbf{D} \) 控制）。

核心优势

  • 线性复杂度：通过递归计算，处理长度为 \( n \) 的序列仅需 \( O(n) \) 时间，远优于注意力的 \( O(n^2) \)。
  • 理论上无限上下文：状态 \( h(t) \) 理论上可以压缩无限长的历史信息。
  • 可并行训练：通过特定的数学变换（如HiPPO初始化、结构化矩阵和并行扫描算法），可以在训练时实现高效的并行计算。

从S4到Mamba的进化

早期的深度学习SSM，如S4 (Structured State Space Sequence Model)，证明了SSM在长序列任务（如语音、基因组学）上可以媲美甚至超越Transformer。然而，S4有一个关键限制：它是一个线性时不变（LTI）系统。

这意味着参数 \( \mathbf{A}, \mathbf{B}, \mathbf{C} \) 在序列处理过程中是固定的，与输入内容无关。这限制了模型根据上下文进行“选择性”关注或遗忘的能力，而这正是语言理解等任务的核心。

S4与Mamba对比示意图

图1: S4（时不变）与Mamba（时变/选择性）的对比。Mamba的参数（Δ, B, C）根据输入动态变化，实现了内容感知的信息处理。

Mamba的突破性贡献在于，它将SSM从“时不变”升级为“时变”或“选择性”系统。简单说，Mamba让SSM的参数能够根据当前的输入 \( x(t) \) 动态变化，从而实现了类似注意力的内容感知能力，同时又保持了线性复杂度。

选择性：Mamba的核心创新

Mamba实现“选择性”的核心机制是引入一个基于输入的离散化参数。在传统的S4中，连续参数 \( \mathbf{A}, \mathbf{B} \) 被离散化为 \( \bar{\mathbf{A}}, \bar{\mathbf{B}} \) 的步骤是固定的。而在Mamba中，这一步变得动态。

具体来说，Mamba模型会：

  • 从当前输入 \( x(t) \) 通过一个线性层投影，得到一组控制参数（特别是时间步长参数 \( \Delta_t \)）。
  • 使用这个 \( \Delta_t \) 来动态地离散化 \( \mathbf{A} \) 和 \( \mathbf{B} \)： \( \bar{\mathbf{A}}_t = \exp(\Delta_t \mathbf{A}) \), \( \bar{\mathbf{B}}_t = \Delta_t \mathbf{B} \)。
\[ \begin{aligned} \bar{\mathbf{A}}_t &= \exp(\Delta_t \mathbf{A}) \\ \bar{\mathbf{B}}_t &= \Delta_t \mathbf{B} \\ h_t &= \bar{\mathbf{A}}_t h_{t-1} + \bar{\mathbf{B}}_t x_t \\ y_t &= \mathbf{C}_t h_t \end{aligned} \]

这里的精妙之处在于 \( \Delta_t \)：

  • 当 \( \Delta_t \) 较大时， \( \bar{\mathbf{A}}_t \) 衰减慢，模型倾向于保留更多历史信息（记住）。
  • 当 \( \Delta_t \) 较小时， \( \bar{\mathbf{A}}_t \) 衰减快，模型倾向于忽略遥远历史，聚焦当前输入（遗忘/关注）。
  • 参数 \( \mathbf{B}_t \) 和 \( \mathbf{C}_t \) 也变为输入依赖，决定了当前输入如何影响状态以及状态如何贡献给输出。

这使Mamba能够像人阅读一样，根据读到的是关键名词还是无关虚词，动态决定是扩大“记忆窗口”还是快速略过。

效率优势与硬件友好性

Mamba不仅在理论上优雅，在工程实践上也极具优势。其效率核心源于对现代硬件（尤其是GPU）特性的深度利用。

1. 线性序列长度缩放

这是最显著的优势。在处理长达100万token的序列时，Mamba的内存消耗增长是线性的，而标准Transformer是爆炸性的平方级增长。

2. 无需键值（KV）缓存

Transformer在自回归生成（如逐字生成文本）时，需要缓存之前所有时间步的键（Key）和值（Value）张量，这带来了巨大的内存开销。Mamba的递归特性使其只需要维护一个固定大小的隐藏状态 \( h_t \)，内存占用恒定。

3. 硬件感知算法

Mamba论文提出了一个关键见解：选择性SSM的递归计算虽然破坏了传统S4的卷积并行性，但可以通过并行扫描（Parallel Scan）算法在GPU上高效实现。结合对GPU内存层次（SRAM vs HBM）的精心管理，Mamba实现了极高的硬件利用率，训练和推理速度远超同等规模的Transformer。

Mamba块结构图

图2: Mamba块的结构。它结合了选择性SSM、残差连接和门控机制，构成了一个强大的序列建模基本单元。

潜在应用与未来展望

Mamba架构为多个领域带来了新的可能性：

  • 超长文本建模：处理书籍、长代码库、法律文档、学术论文，无需复杂的分段和上下文管理。
  • 高分辨率视觉：将图像视为极长的像素序列，应用于视频理解、医学影像分析。
  • 多模态与强化学习：在需要长期记忆和快速决策的任务中，如机器人控制、游戏AI。
  • 边缘设备部署：低内存占用的特性使其更适合在手机、IoT设备上运行大型语言模型。

目前，社区已经出现了基于Mamba的各类模型，如语言模型Mamba-2.8B、视觉模型VMamba、多模态模型MambaByte等。未来的方向可能包括：

  • 与注意力机制或MoE（专家混合）结合，形成混合架构。
  • 探索更高效的选择性机制和状态压缩方法。
  • 将其确立为下一代大语言模型的基础骨干网络。

概念性代码演示

以下是一个高度简化的、用于理解Mamba选择性SSM核心循环的概念性PyTorch代码。实际实现（如`mamba-ssm`库）使用了更复杂的并行扫描和硬件优化。

import torch
import torch.nn as nn

class SimplifiedSelectiveSSM(nn.Module):
    """
    一个极度简化的选择性SSM演示，用于理解其递归逻辑。
    忽略离散化、并行化、结构化矩阵等复杂细节。
    """
    def __init__(self, state_dim, input_dim):
        super().__init__()
        self.state_dim = state_dim
        # 可学习的全局参数（实际Mamba中A是结构化矩阵）
        self.A = nn.Parameter(torch.randn(state_dim, state_dim))
        # 投影层：从输入生成时变参数 B_t, C_t, Δ_t
        self.proj = nn.Linear(input_dim, input_dim + 2 * state_dim)

    def forward(self, x):
        # x 形状: (batch_size, seq_len, input_dim)
        batch_size, seq_len, _ = x.shape
        # 初始化隐藏状态
        h = torch.zeros(batch_size, self.state_dim, device=x.device)
        outputs = []

        for t in range(seq_len):
            x_t = x[:, t, :]  # 当前时间步的输入
            # 1. 根据输入生成时变参数
            proj_result = self.proj(x_t)  # (batch, input_dim + 2*state_dim)
            delta_t = torch.softplus(proj_result[:, :self.state_dim])  # Δ_t, 控制时间尺度
            B_t = proj_result[:, self.state_dim:2*self.state_dim]  # B_t
            C_t = proj_result[:, 2*self.state_dim:]  # C_t

            # 2. 简化的选择性离散化（此处仅为示意，非精确公式）
            A_bar_t = torch.matrix_exp(delta_t.unsqueeze(-1) * self.A)  # 动态的A_bar
            B_bar_t = delta_t.unsqueeze(-1) * B_t.unsqueeze(-1)  # 动态的B_bar

            # 3. 递归更新状态 (核心SSM方程)
            h = torch.bmm(A_bar_t, h.unsqueeze(-1)).squeeze(-1) + B_bar_t.squeeze() * x_t
            # 4. 计算输出
            y_t = (C_t * h).sum(dim=-1, keepdim=True)
            outputs.append(y_t)

        return torch.stack(outputs, dim=1)

# 示意性使用
model = SimplifiedSelectiveSSM(state_dim=16, input_dim=32)
dummy_input = torch.randn(2, 100, 32)  # (batch, seq_len, dim)
output = model(dummy_input)
print(f"输入形状: {dummy_input.shape}")
print(f"输出形状: {output.shape}")  # 应为 (2, 100, 1)

这段代码清晰地展示了“选择性”的核心：参数 `Δ_t, B_t, C_t` 来源于当前输入 `x_t`，从而实现了信息处理的动态化。真正的Mamba实现会使用并行扫描来替代这个for循环，以实现训练时的高效并行。

结论

<