阅读量:10
Mamba 面世以来，社区反应热烈，获得了广泛关注。
遗憾的是，Mamba 的论文竟然被 ICLR 拒稿，这让许多研究者颇感意外。
仅仅六个月之后，原作者团队带来更加改进和强大的 Mamba 2 正式发布啦！！
Mamba-2，其状态空间扩大了8倍，训练速度提升了50%！
——Transformer的新挑战者，最新架构Mamba刚刚推出了其第二代版本！
论文地址： Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality
GitHub 地址： https://github.com/state-spaces/mamba
更令人惊讶的是，团队的研究发现Transformer和状态空间模型（SSM）实际上是亲近关系。
这两大主流的序列建模架构，终于在此实现了统一。
OpenAI发布GPT-4o一夜创历史，超越所有AI
https://www.zhihu.com/pin/1773645611381747712
没体验过OpenAI最新版GPT-4o？快戳最详细升级教程，几分钟搞定：
手把手升级ChatGPT-4o Turbo详细步骤教程https://www.zhihu.com/pin/1768399982598909952
01 状态空间模型SSM
没错，这篇论文揭示了一个重大的发现：Transformer中的注意力机制与状态空间模型（SSM）之间存在着非常紧密的数学联系。
团队通过提出一个名为结构化状态空间二元性（Structured State Space Duality, SSD）的理论框架，将这两大模型家族统一了起来。
值得一提的是，Mamba一代论文在年初被ICLR拒稿，曾引发学术界广泛热议，许多学者因此纷纷表态。
此次推出的二代论文在理论和实验上都有了更丰富的内容，并成功入选ICML 2024。
论文的作者依然是Albert Gu和Tri Dao两位。
他们透露，论文题目中的“Transformers are SSMs”是向四年前那篇经典的线性注意力论文“Transformers are RNNs”致敬。
Transformer的核心组件是注意力机制，而SSM模型的核心则是一个线性时变系统。
乍看之下，两者似乎毫无关联，但论文指出：它们都可以表示成可半分离矩阵（Semiseparable Matrices）的变换。
让我们先从SSM的角度来看。
SSM本身定义了一个线性映射，恰好可对应于一个半可分离矩阵。
半可分离矩阵具有特殊的低秩结构，而这种结构恰好与SSM模型中的状态变量相对应。
因此，矩阵乘法就相当于SSM的线性时变系统。带选择性的SSM实际上是一种广义线性注意力机制。
从注意力的角度来看，又会有什么发现呢？
团队试图通过更抽象的方式来描述注意力机制的本质，毕竟“Softmax自注意力”只是其中的一种形式。
更一般地说，任何带有掩码的注意力机制，都可以表示为4个张量的缩并（Contraction）。
其中，QKV分别对应注意力中的query、key、value，L则对应掩码矩阵。 借助这一联系，团队在线性注意力的基础上提出了结构化掩码注意力（Structured Masked Attention，SMA）。
当注意力的掩码矩阵是半可分离的，就与SSM等价。
基于这一发现，作者进一步推导出了两种等价的计算形式，这就是论文核心思想——“状态空间二元性”（Structured State Space Duality，SSD）的由来。
02 状态空间二元性SSD
基于SSD思想的全新算法，Mamba-2 支持更大的状态维度，从16扩展到了256，从而能学习到更强的表示能力。
新方法采用了基于块分解的矩阵乘法，利用了GPU的存储层次结构，从而提升了训练速度。
在架构设计上，Mamba-2简化了块的设计，同时汲取了注意力机制的启发，做出了一些改动。
借鉴多头注意力机制，创建了多输入SSM。
通过与注意力机制之间的联系，SSD能够轻松将Transformer架构多年来积累的优化方法引入SSM。
例如，引入张量并行和序列并行，以扩展到更大的模型和更长的序列。
同时，还可以引入可变序列长度，以实现更快的微调和推理。
Mamba-2的SSD层比Mamba-1中的关联扫描快了很多，这使团队能够增加状态维度并提高模型质量。
在实验中，Mamba-2在3B参数规模上训练了300B tokens，超越了相同规模的Mamba-1和Transformer。
尤其在需要更大状态容量的任务上，如硬关联召回任务（MQAR），Mamba-2相较于Mamba-1有了显著的改进。
03 小结
man, what can i say!
推荐阅读：
如何免费使用GPT-4o？如何升级GPT4.0 Turbo？（内附详细步骤教程）
最新版手把手升级GPT-4o、GPT-4 Turbo详细步骤教程！！ 【2024年6月】
