Mamba: Linear-Time Sequence Modeling with Selective State Spaces

基本信息

• 标题：Mamba: Linear-Time Sequence Modeling with Selective State Spaces
• 作者：Albert Gu, Tri Dao
• 提交日期：2023年12月1日（v1），最后修订：2024年5月31日（v2）
• arXiv 编号：arXiv:2312.00752 [cs.LG]
• DOI：https://doi.org/10.48550/arXiv.2312.00752

主题分类

• 主分类：Machine Learning (cs.LG)
• 其他分类：Artificial Intelligence (cs.AI)

摘要

基础模型（Foundation models）目前支撑着深度学习领域大多数令人瞩目的应用，它们几乎普遍基于 Transformer 架构及其核心注意力模块。为解决 Transformer 在长序列上的计算效率问题，研究者提出了许多次二次时间复杂度（subquadratic-time）的架构，如线性注意力（linear attention）、门控卷积与循环模型（gated convolution and recurrent models），以及结构化状态空间模型（structured state space models, SSMs）。然而，这些模型在语言等重要模态上的表现仍不及注意力机制。

本文指出，这类模型的一个关键缺陷是无法进行基于内容的推理（content-based reasoning），并据此提出了若干改进：

1. 输入依赖的 SSM 参数：让 SSM 参数成为输入的函数，从而解决其在离散模态上的弱点，使模型能够根据当前 token 选择性地沿序列长度维度传播或遗忘信息。
2. 硬件感知的并行算法：尽管上述改动阻碍了高效卷积的使用，作者设计了一种基于循环模式的硬件感知并行算法。

作者将这些选择性 SSM（selective SSMs）集成到一个简化的端到端神经网络架构中，不包含注意力机制，甚至不含 MLP 块，即 Mamba。

Mamba 具备以下特性：
- 快速推理：吞吐量比 Transformer 高 5 倍
- 线性序列长度扩展：在真实数据上可处理长达百万级长度的序列
- 作为通用序列模型骨干，在语言、音频、基因组学等多种模态上取得最先进（state-of-the-art）性能
- 在语言建模上，Mamba-3B 模型在预训练与下游评估中均超越同等规模的 Transformer，并匹敌两倍规模的 Transformer

关键要点

• 提出选择性状态空间模型，解决现有 SSM 无法进行内容推理的核心问题
• 设计硬件感知并行扫描算法，兼顾训练效率与推理速度
• 在多个模态上验证了通用序列建模能力，成为 Transformer 之外的重要替代方案