0%
💡 왼쪽 원문을 읽으면서 오른쪽에 따라 써보세요. Tab 키로 힌트를 받을 수 있습니다.
概述
自 2017 年 "Attention is All You Need" 论文发表以来，Transformer 几乎统治了从自然语言处理到图像、代码、音频在内的所有序列建模任务。然而 Transformer 有一个根本性的弱点，那就是自注意力的二次复杂度(O(n^2))问题。
2023 年，Albert Gu 与 Tri Dao 发表的 Mamba 优雅地解决了这一问题，给深度学习社区带来了巨大的冲击。Mamba 把状态空间模型(State Space Model, SSM)这一源自经典控制理论的概念，嫁接到现代深度学习之中，使得以线性时间复杂度处理长序列成为可能。
本指南将完整覆盖从 SSM 的数学基础、Mamba 的核心创新，到实战实现的全部内容。
1. Transformer 的局限与 SSM 的兴起
自注意力的 O(n^2) 复杂度问题
Transformer 的核心是自注意力(Self-Attention)机制。对于长度为 n 的序列，计算注意力矩阵时需要计算所有 token 两两之间的关系，因此时间与内存复杂度都是 O(n^2)。
注意力矩阵大小：n × n
= 1,000 个 token → 1,000,000 个元素
= 10,000 个 token → 100,000,000 个元素
= 100,000 个 token → 10,000,000,000 个元素(不可行！)
正因如此，实际的 LLM 都受到上下文窗口大小的制约。GPT-4 支持 128K token，但为此需要庞大的计算成本与优化技巧。
长序列处理的难题
长文档摘要、代码库理解、长时间对话维持等实际应用场景，都需要非常长的上下文。但 Transformer 的二次复杂度让实用化变得困难。
举例来说，若要一次性处理一整本书(约 100,000 词 = 130,000 token)，标准 Transformer 在当前硬件条件下所需的内存几乎不可能满足。
循环模型(RNN、LSTM)的局限
循环神经网络(RNN)与 LSTM 在理论上具有 O(n) 复杂度，因为它们在每个时间步都只维护一个固定大小的隐藏状态(hidden state)。但 RNN/LSTM 存在以下问题：
- 梯度消失/爆炸(Vanishing/Exploding Gradient)：长序列上反向传播时梯度会消失或爆炸
- 无法并行化：每个时间步都依赖前一个状态，导致 GPU 并行化困难
- 难以捕捉长程依赖：难以连接相距很远的信息
SSM 给出的解决方案
状态空间模型(SSM)结合了这两种方法的优点：
- 训练时：通过卷积(Convolution)实现完全并行化 → 训练效率可媲美 Transformer
- 推理时：通过循环(Recurrence)实现 O(1) 内存与 O(n) 计算 → 像 RNN 一样高效推理
- 长序列：线性复杂度使得处理非常长的序列成为可能
2. 状态空间模型(SSM)的数学基础
连续时间 SSM
SSM 的起源可以追溯到 1960 年代 Rudolf Kalman 的控制理论。连续时间线性系统可以表示为：
x'(t) = A·x(t) + B·u(t)
y(t) = C·x(t) + D·u(t)
其中：
u(t)：输入信号(input)x(t)：状态向量(latent state)，维度为 Ny(t)：输出信号(output)A：N×N 状态转移矩阵(state transition matrix)B：N×1 输入矩阵(input projection)C：1×N 输出矩阵(output projection)D：直接传递项(skip connection，通常设为 0)
该系统接收输入 u(t)，更新内部状态 x(t)，并生成输出 y(t)。状态 x(t) 可以看作是过去信息的"摘要"。
离散化(Discretization)
要在数字计算机上使用连续时间 SSM，需要进行离散化(discretization)。设采样间隔为 Delta，主要有两种离散化方法。
Zero-Order Hold (ZOH)：
A_bar = exp(Delta · A)
B_bar = (Delta · A)^(-1) · (exp(Delta · A) - I) · Delta · B
Bilinear (Tustin) 变换：
A_bar = (I - Delta/2 · A)^(-1) · (I + Delta/2 · A)
B_bar = (I - Delta/2 · A)^(-1) · Delta · B
离散化之后系统变为：
x[t] = A_bar · x[t-1] + B_bar · u[t]
y[t] = C · x[t]
这一形式就是线性循环(Linear Recurrence)，在每个时间步更新状态。
作为卷积核的 SSM
SSM 的强大之处在于训练时可以用卷积来计算。若以初始状态 x[0] = 0 开始：
y[0] = C · B_bar · u[0]
y[1] = C · A_bar · B_bar · u[0] + C · B_bar · u[1]
y[2] = C · A_bar^2 · B_bar · u[0] + C · A_bar · B_bar · u[1] + C · B_bar · u[2]
...
把它表示为卷积核 K：
K = (C·B_bar, C·A_bar·B_bar, C·A_bar^2·B_bar, ...)
y = K * u (卷积)
这个核 K 可以并行计算，若使用 FFT，则可以做到 O(n log n)。
这就是 SSM 的核心二元性：
- 推理：通过循环(Recurrence)实现 O(1) 内存
- 训练：通过卷积(Convolution)实现 O(n log n) 并行计算
3. S4 (Structured State Space Sequence Model)
S4 是 Albert Gu 等人于 2021 年发表的论文，是把 SSM 实用化地引入深度学习的第一项重要工作。
HiPPO 矩阵初始化
S4 的核心贡献之一是 HiPPO(High-order Polynomial Projection Operators)矩阵初始化。若简单地随机初始化 A，会出现梯度消失问题。HiPPO 提供了一个专门设计的 A 矩阵，使其能以多项式近似过去的输入。
HiPPO-LegS(基于 Legendre 多项式)：
A[n,k] = -sqrt((2n+1)(2k+1)) if n > k
A[n,k] = -(n+1) if n == k
A[n,k] = 0 if n < k
这一初始化保证状态 x[t] 始终维持对全部过去输入 u[0..t] 的最优多项式近似。
结构化矩阵 A (DPLR)
为了高效计算 A 矩阵的卷积核，S4 把 A 表示为 Diagonal Plus Low-Rank (DPLR) 形式：
A = Λ - P·Q^T
其中 Λ 是对角矩阵，P、Q 是低秩向量。利用这一结构，核计算可以降到 O(N)。
高效计算
import torch
import torch.nn as nn
import numpy as np
class S4Layer(nn.Module):
"""S4 层的简化实现"""
def __init__(self, d_model, d_state=64, dropout=0.0):
super().__init__()
self.d_model = d_model
self.d_state = d_state
# HiPPO-LegS 初始化
A = self._make_hippo(d_state)
# DPLR 分解
self.A = nn.Parameter(torch.tensor(A, dtype=torch.float32))
self.B = nn.Parameter(torch.randn(d_state, 1) * 0.01)
self.C = nn.Parameter(torch.randn(1, d_state) * 0.01)
self.log_delta = nn.Parameter(torch.zeros(d_model))
self.D = nn.Parameter(torch.ones(d_model))
def _make_hippo(self, N):
"""生成 HiPPO-LegS 矩阵"""
P = np.sqrt(1 + 2 * np.arange(N))
A = P[:, np.newaxis] * P[np.newaxis, :]
A = np.tril(A) - np.diag(np.arange(N))
return -A
def discretize(self, A, B, C, delta):
"""ZOH 离散化"""
dA = torch.matrix_exp(delta.unsqueeze(-1) * A)
dB = torch.linalg.solve(A, (dA - torch.eye(A.shape[0])) @ B)
return dA, dB, C
def forward(self, u):
# u: (batch, seq_len, d_model)
delta = torch.exp(self.log_delta)
# 离散化及卷积计算
# 实际实现更复杂，这里仅作概念说明
return u # 简化
4. H3 (Hungry Hungry Hippos)
2022 年发表的 H3，是把 S4 改进得更适合语言建模的模型。这个名字是对 HiPPO 缩写的幽默化表达。
与 S4 的区别
S4 能很好地捕捉长程依赖，但在语言建模中重要的token 间交互却有所欠缺。注意力可以直接计算"这个词与那个词有多相关"，但纯粹依靠循环的 S4 很难做到这一点。
门控机制
H3 使用两个 SSM，并在两者之间加入门控：
class H3Layer(nn.Module):
"""H3 层"""
def __init__(self, d_model, d_state=64):
super().__init__()
# 两个 SSM(shift SSM + diagonal SSM)
self.shift_ssm = S4Layer(d_model, d_state)
self.diag_ssm = S4Layer(d_model, d_state)
# 投影
self.Q_proj = nn.Linear(d_model, d_model)
self.K_proj = nn.Linear(d_model, d_model)
self.V_proj = nn.Linear(d_model, d_model)
self.out_proj = nn.Linear(d_model, d_model)
def forward(self, x):
# Q、K、V 投影
q = self.Q_proj(x)
k = self.K_proj(x)
v = self.V_proj(x)
# SSM 1：对 K 应用 Shift SSM
k_ssm = self.shift_ssm(k)
# 乘法门控：Q ⊙ K_ssm
gated = q * k_ssm
# SSM 2：对 V 应用 Diagonal SSM，再与 gated 相乘
v_ssm = self.diag_ssm(v)
output = gated * v_ssm
return self.out_proj(output)
语言建模的改进
H3 在 GPT-2 规模的语言模型上取得了接近 Transformer 的性能，同时推理速度快得多。这是证明 SSM 在语言建模中具有实用性的重要里程碑。
5. Mamba (Selective State Space Model)
2023 年 12 月，Albert Gu 与 Tri Dao 发表的 Mamba，是 SSM 研究中最重要的进展。Mamba 的核心创新是选择性机制(Selective Mechanism)。
核心创新：选择性机制
现有 S4 与 H3 的根本局限在于，矩阵 A、B、C 都是与输入无关的(input-independent)。因为要用卷积进行训练，核就必须是固定的。
Mamba 打破了这一约束，引入了 S6(Selective SSM)：把 B、C、Delta 矩阵变成输入 u(t) 的函数。
B(t) = Linear_B(u(t)) # 随输入变化的 B
C(t) = Linear_C(u(t)) # 随输入变化的 C
Delta(t) = softplus(Linear_Delta(u(t))) # 随输入变化的 Delta
这样一来，模型就能根据输入动态地决定把哪些信息存入状态、又忽略哪些信息。
直观地说：
- Delta 较大 → 当前输入被强烈地反映到状态中(记住信息)
- Delta 较小 → 状态变化很小(忽略信息)
这与 LSTM 的门(输入门、遗忘门)所起的作用类似，但它是在连续时间 SSM 框架内完成的。
Hardware-Aware Algorithm
引入选择性机制之后，由于 B、C 依赖于输入，就不能再用卷积来计算，这可能带来严重的计算低效问题。
Mamba 为此使用了 Hardware-Aware Algorithm，这是一种利用 GPU 内存层级结构(HBM vs SRAM)的核融合(Kernel Fusion)技巧：
问题：如果每个时间步循环计算的中间状态都存放在 HBM(慢速内存)中，会产生内存带宽瓶颈
解决方案：
- 所有中间计算都在快速的 SRAM(片上内存)中完成
- 使用并行扫描(Parallel Scan)算法，而非完全卷积
- 只把最终输出写入 HBM
# 并行扫描的概念(实际实现使用 CUDA)
def parallel_scan(gates, tokens):
"""
并行计算线性循环
x[t] = gates[t] * x[t-1] + tokens[t]
通过二叉树结构实现 O(log n) 的并行深度
"""
n = len(tokens)
# 上扫(up-sweep)阶段
log_n = int(np.log2(n))
for d in range(log_n):
step = 2 ** (d + 1)
for i in range(step - 1, n, step):
gates[i] = gates[i] * gates[i - 2**d]
tokens[i] = gates[i - 2**d] * tokens[i - 2**d] + tokens[i]
# 下扫(down-sweep)阶段(省略)
return tokens
Mamba 块结构
输入 x (B, L, D)
│
├──────────────────────┐
│ │
Linear(D→ED) Linear(D→ED)
+ SiLU activation │
│ SSM (S6)
│ │
└─────── ⊙ ────────────┘
(element-wise 乘法)
│
Linear(ED→D)
│
输出 y (B, L, D)
这里 E 是扩展比率(通常为 2)，D 是模型维度。
完整的 Mamba 实现
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, repeat
class MambaBlock(nn.Module):
"""
Mamba Block 实现
论文："Mamba: Linear-Time Sequence Modeling with Selective State Spaces"
"""
def __init__(
self,
d_model, # 模型维度 D
d_state=16, # SSM 状态维度 N
d_conv=4, # 卷积核大小
expand=2, # 扩展比率 E
dt_rank="auto",
dt_min=0.001,
dt_max=0.1,
dt_init="random",
dt_scale=1.0,
dt_init_floor=1e-4,
bias=False,
):
super().__init__()
self.d_model = d_model
self.d_state = d_state
self.d_conv = d_conv
self.expand = expand
self.d_inner = int(self.expand * self.d_model)
if dt_rank == "auto":
self.dt_rank = max(1, int(d_model / 16))
else:
self.dt_rank = dt_rank
# 输入投影(D → 2*ED，一次计算两个分支)
self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=bias)
# 局部卷积(depthwise conv)
self.conv1d = nn.Conv1d(
in_channels=self.d_inner,
out_channels=self.d_inner,
bias=True,
kernel_size=d_conv,
groups=self.d_inner,
padding=d_conv - 1, # causal padding
)
# SSM 参数投影
self.x_proj = nn.Linear(
self.d_inner, self.dt_rank + self.d_state * 2, bias=False
)
self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True)
# A 初始化(基于 HiPPO)
A = repeat(
torch.arange(1, self.d_state + 1),
"n -> d n",
d=self.d_inner
)
self.A_log = nn.Parameter(torch.log(A))
self.A_log._no_weight_decay = True
# D(skip connection)
self.D = nn.Parameter(torch.ones(self.d_inner))
self.D._no_weight_decay = True
# 输出投影
self.out_proj = nn.Linear(self.d_inner, d_model, bias=bias)
def forward(self, hidden_states):
"""
hidden_states: (B, L, D)
Returns: (B, L, D)
"""
batch, seqlen, dim = hidden_states.shape
# 输入投影
xz = self.in_proj(hidden_states)
x, z = xz.chunk(2, dim=-1) # 各为 (B, L, ED)
# 卷积(causal 1D conv)
x = rearrange(x, "b l d -> b d l")
x = self.conv1d(x)[:, :, :seqlen] # causal trimming
x = rearrange(x, "b d l -> b l d")
x = F.silu(x)
# SSM
y = self.ssm(x)
# 门控(与 z 分支逐元素相乘)
y = y * F.silu(z)
# 输出投影
output = self.out_proj(y)
return output
def ssm(self, x):
"""选择性 SSM (S6) 计算"""
d_in, n = self.d_inner, self.d_state
# A 矩阵(用 -exp(-A_log) 始终保证为负数 = 稳定性)
A = -torch.exp(self.A_log.float()) # (ED, N)
# x_proj：计算 Delta、B、C
x_dbl = self.x_proj(x) # (B, L, dt_rank + 2N)
delta, B, C = x_dbl.split(
[self.dt_rank, n, n], dim=-1
)
# Delta：用 softplus 保证为正数
delta = F.softplus(self.dt_proj(delta)) # (B, L, ED)
# ZOH 离散化及选择性扫描
y = self.selective_scan(x, delta, A, B, C, self.D)
return y
def selective_scan(self, u, delta, A, B, C, D):
"""
选择性扫描算法
实际实现中使用 mamba_ssm 的 CUDA 核
这里用纯 PyTorch 说明概念
"""
b, l, d_in = u.shape
n = A.shape[1]
# 离散化
# delta: (B, L, ED), A: (ED, N) → dA: (B, L, ED, N)
deltaA = torch.exp(
torch.einsum("bld,dn->bldn", delta, A)
)
# delta: (B, L, ED), B: (B, L, N), u: (B, L, ED)
# → deltaB_u: (B, L, ED, N)
deltaB_u = torch.einsum("bld,bln,bld->bldn", delta, B, u)
# 循环扫描
x = torch.zeros(b, d_in, n, device=u.device, dtype=u.dtype)
ys = []
for i in range(l):
x = deltaA[:, i] * x + deltaB_u[:, i]
y = torch.einsum("bdn,bn->bd", x, C[:, i, :])
ys.append(y)
y = torch.stack(ys, dim=1) # (B, L, ED)
# D(skip connection)
y = y + u * D
return y
class MambaModel(nn.Module):
"""由若干堆叠的 Mamba 块组成的序列模型"""
def __init__(self, d_model, n_layers, d_state=16, expand=2):
super().__init__()
self.layers = nn.ModuleList([
MambaBlock(d_model, d_state=d_state, expand=expand)
for _ in range(n_layers)
])
self.norm = nn.LayerNorm(d_model)
def forward(self, x):
"""x: (B, L, D)"""
for layer in self.layers:
x = x + layer(self.norm(x)) # Pre-norm residual
return x
6. Mamba 2
2024 年 5 月，Tri Dao 与 Albert Gu 发表了 Mamba 2，进一步强化了理论基础。
与 Mamba 1 的区别
Mamba 1 的选择性扫描需要在每个时间步进行顺序计算。Mamba 2 发现了 State Space Duality(SSD)，让计算变得更加高效。
State Space Duality (SSD)
Mamba 2 的核心洞见是：具有特定结构的 SSM 可以表示为半可分矩阵(Semi-Separable Matrix)，而这与特定形式的注意力是等价的。
通过这一数学等价关系(duality)：
- SSM 计算可以表示为矩阵乘法形式
- 可以利用高度优化的 BLAS 库
- 借助 Tensor Core 把 GPU 效率发挥到极致
# SSD 运算的概念
# SSM 循环：x[t] = A[t] * x[t-1] + B[t] * u[t]
# y[t] = C[t]^T * x[t]
#
# 按 chunk(块)单位处理：
# - chunk 内部：用矩阵乘法并行计算
# - chunk 之间：用循环传播状态
class Mamba2Block(nn.Module):
"""Mamba 2 块"""
def __init__(self, d_model, d_state=64, n_heads=8, chunk_size=64):
super().__init__()
self.d_model = d_model
self.d_state = d_state
self.n_heads = n_heads
self.chunk_size = chunk_size
self.d_head = d_model // n_heads
# 多头结构
self.norm = nn.LayerNorm(d_model)
self.in_proj = nn.Linear(d_model, d_model * 2 + d_state * 2 + n_heads)
self.out_proj = nn.Linear(d_model, d_model)
# A 参数(按头划分)
self.A_log = nn.Parameter(torch.randn(n_heads))
def forward(self, x):
"""x: (B, L, D)"""
# 实际实现使用 mamba_ssm 包的 CUDA 核
return x
多头结构
Mamba 2 引入了多头 SSM，使其在结构上与 Transformer 的 Multi-Head Attention 更加相似。这带来了：
- 更强的表示能力
- 与注意力之间的理论关联
- 更易于设计混合架构
与 Transformer 融合的可能性
SSD 理论表明，特定的 SSM 与特定的注意力是等价的。这为把 Mamba 与 Transformer 混合起来的混合架构提供了理论依据。
7. 混合架构
MambaFormer
MambaFormer 把 Mamba 块与 Transformer 块交替(interleaving)排列：
第 1 层：Mamba Block (捕捉局部模式)
第 2 层：Attention (捕捉全局依赖)
第 3 层：Mamba Block
第 4 层：Attention
...
这一结构利用了各组件的优势：
- Mamba：高效的序列处理，局部模式
- Attention：选择性信息检索，全局依赖
Jamba (SSM + Transformer + MoE)
2024 年 AI21 Labs 发布的 Jamba，是更加复杂的混合模型：
Jamba = Mamba + Transformer + Mixture of Experts
架构：
- 52B 参数(激活参数：12B)
- 层比例：Attention 1 : Mamba 7
- 部分层应用 MoE
- 支持 256K 上下文窗口
与同等规模的 Transformer 相比，Jamba：
- 推理吞吐量提升 3 倍
- 长上下文下的内存效率大幅改善
# Jamba 风格的混合块(概念)
class JambaLayer(nn.Module):
def __init__(self, d_model, layer_idx, attn_every_n=8):
super().__init__()
self.use_attention = (layer_idx % attn_every_n == 0)
if self.use_attention:
self.mixer = nn.MultiheadAttention(d_model, num_heads=8)
else:
self.mixer = MambaBlock(d_model)
# MoE(仅部分层)
self.use_moe = (layer_idx % 2 == 0)
if self.use_moe:
self.ffn = MixtureOfExperts(d_model, n_experts=16)
else:
self.ffn = nn.Sequential(
nn.Linear(d_model, d_model * 4),
nn.SiLU(),
nn.Linear(d_model * 4, d_model)
)
def forward(self, x):
x = x + self.mixer(x)
x = x + self.ffn(x)
return x
RWKV(线性注意力)
RWKV 是 Transformer 与 RNN 的混合体，是 "Receptance Weighted Key Value" 的缩写：
- 训练：像 Transformer 一样并行(矩阵形式)
- 推理：像 RNN 一样循环(O(1) 状态)
- 不使用注意力的 token 混合
最新版本(RWKV-6、RWKV-7)展现出与 Mamba 相竞争的性能。
RetNet (Retentive Network)
微软的 RetNet 旨在同时实现 "Training Parallelism, Inference Efficiency, Competitive Performance"：
3 种计算范式：
1. Parallel：训练时 O(n^2) 并行计算(相比 Transformer 常数更低)
2. Recurrent：推理时 O(1) 内存
3. Chunkwise Recurrent：兼顾两者的折中方法
8. Mamba 性能对比
推理速度(线性扩展)
Mamba 最大的优势在于随序列长度线性扩展。
以序列长度 2K 时的性能归一化为 1：
序列长度 Transformer Mamba
1K 0.5x 0.5x
2K 1.0x 1.0x
4K 3.5x 2.0x ← 比 Transformer 快 1.75 倍
8K 13x 4.0x ← 比 Transformer 快 3.25 倍
16K 50x 8.0x ← 比 Transformer 快 6.25 倍
100K ~1800x ~50x ← 比 Transformer 快 36 倍
内存效率
推理时状态(state)大小的对比：
模型 推理时内存(每 1K token)
Transformer O(n) KV Cache
Mamba O(1) SSM 状态(固定大小)
例：130M 参数模型，1M token 序列
- Transformer：约 16GB KV Cache
- Mamba：约 1MB 状态(常数！)
长序列任务
Long Range Arena (LRA) 基准测试(序列长度 1K-16K)：
模型 ListOps Text Retrieval Image Path-X 平均
Transformer 36.4 65.0 57.5 42.4 0.0 40.3
LSTM 35.9 63.7 65.0 43.3 0.0 41.6
S4 59.6 86.8 90.9 88.7 86.1 82.4
Mamba ~达到 S4 水准的性能，处理速度更快
Mamba 在选择性复制(Selective Copy)、归纳头(Induction Head)等合成基准测试中明显领先 S4，而这类任务与真实的语言建模能力关系更为密切。
9. Mamba 的应用领域
自然语言处理
基于 Mamba 的语言模型正不断涌现：
- MambaChat：对话式 AI 助手
- Falcon Mamba：TII 以开源形式发布的 7B Mamba 模型
- CodeMamba：专注于代码生成
在长文档处理、摘要、翻译方面，相比 Transformer 更加高效。
生物信息学
基因组序列分析需要处理非常长的序列(数百万 bp)，因此 Mamba 尤其具有优势：
- Caduceus：长 DNA 序列建模
- Hyena：长序列 DNA/蛋白质分析
- 在蛋白质结构预测中的应用潜力
# 生物序列的 Mamba 应用示例
from mamba_ssm import MambaLMHeadModel
# DNA 序列建模(A、T、G、C、N token)
DNA_VOCAB = {'A': 0, 'T': 1, 'G': 2, 'C': 3, 'N': 4}
VOCAB_SIZE = len(DNA_VOCAB)
# 即便是非常长的序列也能高效处理
model = MambaLMHeadModel.from_pretrained(
"state-spaces/mamba-130m",
device="cuda",
dtype=torch.float16
)
时间序列分析
在金融、气象、IoT 传感器数据等长时间序列场景中，Mamba 展现出优势：
- TimeMamba：长时间序列预测
- MambaMixer：多变量时间序列建模
- 基于 S4/S5 的时间序列模型的发展
图像处理(VMamba)
VMamba 把 Mamba 扩展到二维图像处理：
# VMamba 的核心：2D 选择性扫描
# 沿 4 个方向扫描图像，以捕捉 2D 结构
# 方向：
# 1. 左→右，上→下(常规光栅扫描)
# 2. 右→左，下→上(反方向)
# 3. 上→下，左→右(列优先)
# 4. 下→上，右→左(列优先反方向)
class VMambaBlock(nn.Module):
"""VMamba: Visual Mamba Block"""
def __init__(self, d_model, d_state=16):
super().__init__()
self.norm = nn.LayerNorm(d_model)
# 4 方向 SSM
self.ssms = nn.ModuleList([
MambaBlock(d_model, d_state) for _ in range(4)
])
self.out_proj = nn.Linear(d_model * 4, d_model)
def forward(self, x):
"""x: (B, H, W, D) - 图像 patch 嵌入"""
b, h, w, d = x.shape
x_flat = x.view(b, h*w, d)
outputs = []
# 4 方向扫描
for i, ssm in enumerate(self.ssms):
if i == 0: # 正向
seq = x_flat
elif i == 1: # 反向
seq = x_flat.flip(1)
elif i == 2: # 列优先
seq = x.permute(0, 2, 1, 3).reshape(b, h*w, d)
else: # 列优先反向
seq = x.permute(0, 2, 1, 3).reshape(b, h*w, d).flip(1)
out = ssm(seq)
if i % 2 == 1:
out = out.flip(1)
outputs.append(out)
# 合并 4 个方向的结果
combined = torch.cat(outputs, dim=-1)
return self.out_proj(combined).view(b, h, w, d)
10. 实战使用
安装 mamba-ssm 包
# 所需依赖包
pip install torch torchvision torchaudio
pip install causal-conv1d>=1.2.0
pip install mamba-ssm
# 或从源码安装(获取最新特性)
git clone https://github.com/state-spaces/mamba
cd mamba
pip install -e ".[dev]"
使用 MambaLMHeadModel
import torch
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
from transformers import AutoTokenizer
# 加载预训练模型
model_name = "state-spaces/mamba-2.8b-slimpj"
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
model = MambaLMHeadModel.from_pretrained(
model_name,
device="cuda",
dtype=torch.bfloat16
)
model.eval()
# 文本生成
def generate_text(prompt, max_new_tokens=200, temperature=0.7):
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")
with torch.no_grad():
output = model.generate(
input_ids=input_ids,
max_new_tokens=max_new_tokens,
temperature=temperature,
top_p=0.9,
do_sample=True,
eos_token_id=tokenizer.eos_token_id,
)
generated = output[0][input_ids.shape[1]:]
return tokenizer.decode(generated, skip_special_tokens=True)
# 示例
prompt = "State Space Models are powerful because"
result = generate_text(prompt)
print(result)
微调示例
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
from transformers import AutoTokenizer, get_linear_schedule_with_warmup
class TextDataset(Dataset):
def __init__(self, texts, tokenizer, max_length=512):
self.encodings = tokenizer(
texts,
truncation=True,
max_length=max_length,
padding="max_length",
return_tensors="pt"
)
def __len__(self):
return len(self.encodings['input_ids'])
def __getitem__(self, idx):
return {
'input_ids': self.encodings['input_ids'][idx],
'attention_mask': self.encodings['attention_mask'][idx],
}
def finetune_mamba(
model_name="state-spaces/mamba-130m",
texts=None,
num_epochs=3,
learning_rate=1e-4,
batch_size=8,
):
"""微调 Mamba"""
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
tokenizer.pad_token = tokenizer.eos_token
# 加载模型
model = MambaLMHeadModel.from_pretrained(
model_name,
device=device,
dtype=torch.bfloat16
)
# 数据集
dataset = TextDataset(texts, tokenizer)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
# 优化器(Mamba 推荐使用 AdamW)
optimizer = torch.optim.AdamW(
model.parameters(),
lr=learning_rate,
weight_decay=0.1
)
# 调度器
total_steps = len(dataloader) * num_epochs
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=total_steps // 10,
num_training_steps=total_steps
)
# 训练循环
model.train()
for epoch in range(num_epochs):
total_loss = 0
for batch_idx, batch in enumerate(dataloader):
input_ids = batch['input_ids'].to(device)
# 语言建模：预测下一个 token
outputs = model(input_ids)
logits = outputs.logits
# Shift for next token prediction
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = input_ids[..., 1:].contiguous()
loss = nn.CrossEntropyLoss()(
shift_logits.view(-1, shift_logits.size(-1)),
shift_labels.view(-1)
)
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
scheduler.step()
total_loss += loss.item()
if batch_idx % 100 == 0:
print(f"Epoch {epoch+1}, Step {batch_idx}: Loss = {loss.item():.4f}")
avg_loss = total_loss / len(dataloader)
print(f"Epoch {epoch+1} 完成：平均 Loss = {avg_loss:.4f}")
return model
自定义 Mamba 模型配置
from mamba_ssm import Mamba
from mamba_ssm.models.config_mamba import MambaConfig
# 自定义配置
config = MambaConfig(
d_model=1024,
n_layer=48,
vocab_size=50280,
d_state=16,
d_conv=4,
expand=2,
dt_rank="auto",
dt_min=0.001,
dt_max=0.1,
dt_init="random",
dt_scale=1.0,
dt_init_floor=1e-4,
rms_norm=True,
residual_in_fp32=True,
fused_add_norm=True,
pad_vocab_size_multiple=8,
)
# 创建模型(约 1.4B 参数)
model = MambaLMHeadModel(config, device="cuda", dtype=torch.bfloat16)
print(f"参数数量：{sum(p.numel() for p in model.parameters()):,}")
结语：Mamba 的未来
Mamba 给深度学习序列建模领域带来了革命性的变化。归纳其核心贡献：
- 选择性机制：根据输入动态变化的 SSM 参数
- Hardware-Aware 设计：利用 GPU 内存层级实现的高效计算
- 双重表示：训练时并行化，推理时循环化，两端都得到优化
- 线性复杂度：随序列长度线性增长的计算与内存复杂度
仍然存在有待解决的课题：
- 在 In-context Learning 上的表现相比 Transformer 略显不足
- 需要在非常大的规模上进行验证
- 需要更明确地证明相对于基于注意力的模型的优势
不过，Mamba 与 SSM 系列模型，将在长序列处理、实时推理、边缘设备部署等 Transformer 力有不逮的领域，扮演越来越重要的角色。
参考资料
- Mamba 论文："Mamba: Linear-Time Sequence Modeling with Selective State Spaces" (Gu & Dao, 2023)
- Mamba 2 论文："Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality" (Dao & Gu, 2024)
- S4 论文："Efficiently Modeling Long Sequences with Structured State Spaces" (Gu et al., 2021)
- H3 论文："Hungry Hungry Hippos: Towards Language Modeling with State Space Models" (Fu et al., 2022)
- Mamba GitHub：https://github.com/state-spaces/mamba
- Jamba 论文："Jamba: A Hybrid Transformer-Mamba Language Model" (AI21 Labs, 2024)
현재 단락 (1/569)
自 2017 年 "Attention is All You Need" 论文发表以来，Transformer 几乎统治了从自然语言处理到图像、代码、音频在内的所有序列建模任务。然而 Transfor...
작성 글자: 0자원문 글자: 17,759자작성 단락: 0/569