背景

Transformer:以其注意力机制而闻名，其中序列的任何部分都可以动态地与任何其他部分相互作用，特别是具有因果注意力机制的的Transformer，擅长处理序列中的单个元素。但是它们带来了显著的计算和内存成本，与序列长度的平方(L²)成比例。

循环神经网络(rnn): rnn只考虑当前输入和最后一个隐藏状态，按顺序更新隐藏状态。这种方法允许它们潜在地处理无限序列长度和恒定的内存需求。但是rnn的简单性是一个缺点，限制了它们记住长期依赖关系的能力。此外，rnn中的时间反向传播(BPTT)是内存密集型的，并且可能遭受梯度消失或爆炸的影响，尽管有LSTM等创新部分结解决了这个问题。

State Space Models(S4):这些模型已经显示出很好的特性。它们提供了一种平衡，比rnn更有效地捕获远程依赖关系，同时比transformer更高效地使用内存。

Mamba以其选择性状态空间的概念引入了传统状态空间模型的一个有趣的改进。这种方法稍微放松了标准状态空间模型的严格状态转换，使其更具适应性和灵活性（有点类似于lstm）。并且Mamba保留了状态空间模型的高效计算特性，使其能够在一次扫描中执行整个序列的前向传递-这一特性更让人想起Transformer。

在训练期间，Mamba的行为类似于Transformer，同时处理整个序列。而lstm必须一步一步地计算前向传递，即使所有输入都是已知的。在推理中，Mamba的行为更符合传统的循环模型，提供有效的序列处理。

先验状态空间模型(ssm)的一个关键限制是其刚性的、输入不变的结构。这些模型为整个序列使用一组固定参数(我们称它们为a和B)。这种结构甚至比lstm等模型更具限制性，在lstm中，信号的转换可能依赖于先前的隐藏状态和输入。

Mamba则一种范式转换，即如何计算向下一个隐藏状态的过渡？在Mamba的体系结构中，转换依赖于当前输入，这种方法在传统ssm的固定计算和循环神经网络的输入依赖动态性之间取得了平衡。

主要组成如下：

固定主干:从一个隐藏状态到下一个隐藏状态的转换仍然是一个固定的计算(由a矩阵定义)，允许跨序列的预计算。

输入相关转换:输入影响下一个隐藏状态(由B矩阵定义)的方式取决于当前输入，而不是之前的隐藏状态。与传统ssm相比，这种输入依赖性提供了更大的灵活性。

为了满足这种方法的计算需求，Mamba使用了一种硬件感知算法。该算法使用扫描操作而不是卷积来循环执行计算，这样在gpu上非常高效的。尽管输入依赖转换带来了算法复杂性，但这种效率对于保持高性能至关重要。

Mamba和选择性状态空间模型不是同义词。Mamba是一个使用选择性状态空间概念的实现。这种区别是至关重要的，因为它突出了Mamba的独特贡献:在保持计算效率的同时，使SSM框架更加灵活和响应输入。

S4架构

Structured state space sequence models

Mamba的架构主要基于S4，一种最新的状态空间模型（SSM，state space model）架构。

Selective State Space Models

Mamba的架构

SRAM和HBM

gpu包含两种主要类型的内存:HBM (High Bandwidth memory)和SRAM (Static Random-Access memory)。HBM虽然带宽很高，但与更快但更小的SRAM相比，它的访问时间相对较慢。Mamba则使用SRAM在矩阵乘法期间进行快速访问，这是其计算的关键。

计算中的主要瓶颈通常不是计算本身，而是数据在内存类型之间的移动。Mamba通过显著减少传输大量数据的需求来解决这个问题。它通过直接在SRAM中执行算法的关键部分(如离散化和递归计算)来实现，从而减少延迟。

还引入了一个融合选择扫描层，使其内存需求与使用flash attention的优化Transformer实现相当。这一层对于保持效率至关重要，尤其是在处理模型中依赖于输入的元素时。

概念

Reference

• 挑战Transformer的新架构Mamba解析以及Pytorch复现 (opens new window)
• Mamba正式被ICLR拒收！“年度最佳技术原理解读”却火了 (opens new window)
• https://jackcook.com/2024/02/23/mamba.html