Mamba可以替代Transformer，但它们也能组合起来使用

1+1＞2。

Transformer 很厉害，但并不完美，尤其是在处理长序列方面。而状态空间模型（SSM）则在长序列上的表现相当不俗。早在去年就有研究者提出可使用 SSM 替代 Transformer，参见文章《 预训练无需注意力，扩展到4096个token不成问题，与BERT相当》，前些天基于 SSM 方法的 Mamba 更是异军突起，推理吞吐量达到了 Transformer 的五倍之多，参阅《 五倍吞吐量，性能全面包围Transformer：新架构Mamba引爆AI圈》。

但实际上，SSM 和 Transformer 并不是非此即彼的两种架构，它们完全可以组合起来！

近日公布的一篇 NeurIPS 2023 论文《Block-State Transformers》就采用了这种做法，其不仅能轻松支持 65k token 长度的超长输入，而且计算效率还非常高，速度相比使用循环单元的 Transformer 足可提升十倍之多！这篇论文也得到了 Mamba 作者 Tri Dao 的点赞，他表示：「SSM 和Transformer 似乎可以互补。」

但在我们介绍这种新方法前，先简单说说 Transformer。在许多不同的 自然语言处理 （NLP）任务上，Transformer 的表现都非常出色。可以说 Transformer 已经很大相当程度上替代了循环神经网络。不仅如此，它也正在图像和视频等 NLP 之外的领域大展拳脚。

其成功的原因有很多，包括计算效率和架构层面的归纳偏差，这让它们非常适合在自然语言任务进行大规模训练。在计算方面，Transformer 能以并行方式处理输入序列的 token，从而使其能充分利用现代加速器硬件。此外， 注意力机制 让 Transformer 可以找到更长序列之间的关系，其方式是在推断下一个 token 时读取从过去 token 提取的所有信息。相比于 RNN 和 LSTM， 自注意力 有两个优势：(1) 存储信息以及将这些信息直接用作上下文的能力得到了极大提升，(2) 在更长序列上能更稳定地训练。

尽管 Transformer 相比 RNN 有很多优势，但它在输入序列长度的扩展上依然存在问题，其中涉及计算性能和质量等方面的原因。更进一步说，Transformer 的运行时间会随输入序列长度的增长成二次方增长，这会让训练这些模型的成本越来越高。

此外，众所周知使用注意力的 Transformer 在长输入分类任务上表现不佳。最基本的 Transformer 在长序列上训练时可能不稳定，而且其 token 重要度聚焦在当前时间步骤周围约 50 个 token 的局部 感受野 中。

近来，越来越多的研究表明状态空间模型（SSM）可以替代 Transformer，因为 SSM 可以捕获极长序列之中的依赖关系，同时还有更高的计算效率和更好的并行化能力。

尽管 SSM 依然属于自回归序列模型，但其底层的线性时间不变式动态系统可使用基于快速傅立叶变换（FFT）的可并行化卷积算子来高效地处理序列，而且这个过程的复杂度仅为