Files
myWiki/concepts/mamba-2.md

1.6 KiB
Raw Blame History

title, created, updated, type, tags, sources
title created updated type tags sources
Mamba-2 2026-06-18 2026-06-18 concept
ssm
architecture
mamba
efficiency
dao-transformers-are-ssms-2024

Mamba-2

Mamba-2 是 Dao & Gu (2024) 基于 structured-state-space-duality 设计的新架构——核心层是 mamba-ssm 选择性 SSM 的改进版,2-8x 更快

相对于 Mamba 的改进

架构层面

组件 Mamba (2023) Mamba-2 (2024)
A 矩阵 对角矩阵 标量 × 单位矩阵
Head 维度 P 1 64/128
Head 结构 多输入 SSM (MIS) 分组值注意力 (GVA)
并行性 不支持 TP 原生 Tensor Parallelism

效率层面

  • SSD 算法:利用 semiseparable-matrices 的块分解部分用循环O(T)、部分用矩阵乘法GPU 优化)
  • 比 Mamba 的 selective scan 快 2-8x
  • 支持 8x 更大的状态大小N几乎无减速
  • 序列长度 16K 时比 FlashAttention-2 快 6x

Chinchilla 缩放定律

在 Pile 数据集的 Chinchilla 设置下Mamba-2 Pareto 支配 Mamba 和 Transformer++

  • 2.7B 参数 / 300B tokens 训练 → 超越 Pythia-2.8B 和 Pythia-6.9B

关键设计决策

  1. 张量并行友好:将所有数据依赖投影移到块开头并行执行,减少同步点
  2. GVA Head 结构:分组值注意力 — 介于 MHA 和 MQA 之间
  3. 变长序列支持:无需 padding tokens通过传递循环状态实现

参考