架构演进 · 注意力之外

Mamba

Mamba-130m · 选择性状态空间模型

对比
Mamba-130m vs GPT-2 small
层
24 GPT-2 12
dmodel
768 GPT-2 768
混合器
扫描 · 状态 16 GPT-2 注意力 · 12 个头
内存
1.3 MiB，任意长度 GPT-2 每词元 36 KiB
参数
130M GPT-2 124M

全部步骤

  1. GPT-2 vs Mamba

    Mamba-130m 与 GPT-2 small 规模相当，但没有注意力。它的 24 个块每个都是一个混合器：一个加宽投影、一个短的因果卷积、一次选择性状态空间扫描和一个门控。词元之间只能通过一个从左往右传递的状态交换信息。点击标签跳过去。

    24 × 768 · 130M · 无注意力

  2. 用一个滚动状态代替缓存

    为了预测下一个词元，注意力要把它和之前的每个词元比较，并保存它们所有的键和值，所以 KV 缓存随每个词元增长。Mamba 则携带一个固定大小的状态：每个新词元的开销都一样，超过大约 38 个词元后，它的状态就比 GPT-2 的缓存还小。

    KV 缓存 n × 36 KiB · 状态 1.3 MiB

  3. 扫描

    状态是每个通道的一个小向量（Mamba 中是 16 个数）。每来一个词元，它先按 Ā 衰减，再通过 B̄ 写入输入；C 读出输出。这里是一个玩具通道，状态有 4 个数，每个词元的步长 Δ 都相同，就像早期的 S4 模型。

    ht = Ā ht−1 + B̄ xt · yt = C ht

  4. 选择性

    Mamba 让步长 Δ 以及 B 和 C 依赖于当前词元。Δ 大，就把这个词元强力写入并忘掉更多过去；Δ 小，就让它过去而保留记忆。右侧：Mamba-130m 的真实 Δ；在中间层，人名得到的步长最大。

    Δt = softplus(W xt) · Ā_t = exp(Δt A)

  5. 同样的规模，同样的任务

    对同一个提示词，有注意力的 GPT-2 small 和没有注意力的 Mamba-130m 给出的真实下一词元猜测：两个规模相当的模型，在不同的网页文本上训练。两者都想到了猫可能坐的地方。

    同一个提示词，两种架构

代码

# the Mamba block: one mixer where GPT-2 has attention and an MLP
x, z = in_proj(rms_norm(u)).chunk(2, dim=-1)       # 768 → 2 × 1,536
x = F.silu(causal_conv1d(x))                       # depthwise, width 4
dt, B, C = x_proj(x).split([48, 16, 16], dim=-1)   # all three depend on the token
dt = F.softplus(dt_proj(dt))                       # Δ: a step size per channel
A = -torch.exp(A_log)                              # (1536, 16) learned decay rates
for t in range(T):                                 # the selective scan
    h = torch.exp(dt[t, :, None] * A) * h + dt[t, :, None] * B[t] * x[t, :, None]
    y[t] = (h * C[t]).sum(-1) + D * x[t]
u = u + out_proj(y * F.silu(z))

延伸阅读