Dao, T. , & Gu, A. (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. International Conference on Machine Learning (ICML).✅
Gu, A. , & Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv preprint arXiv:2312.00752.✅
在人工智能和自然语言处理领域,Transformer模型一直占据主导地位。然而,近期出现的Mamba模型作为一种新的状态空间模型(SSM)架构,展现出了挑战Transformer地位的潜力。就在Mamba问世仅仅6个月之后,其原创团队再次带来了重大突破 – Mamba-2的正式发布。这一新版本不仅在性能上有了显著提升,更重要的是,它揭示了Transformer和状态空间模型之间深刻的理论联系,为序列建模领域带来了全新的视角。
Mamba-2的主要改进
相比于第一代Mamba模型,Mamba-2在以下几个方面实现了重大升级:
理论突破:Transformer与SSM的统一
Mamba-2最令人兴奋的发现是,Transformer中的注意力机制与状态空间模型(SSM)之间存在着密切的数学联系。这一发现不仅具有理论意义,还为未来模型设计提供了新的思路。
状态空间模型(SSM)视角
SSM定义了一个线性映射,可以表示为一个半可分离矩阵(Semiseparable Matrices)。这种矩阵具有特殊的低秩结构,与SSM中的状态变量直接对应。因此,矩阵乘法可以等价于SSM的线性时变系统。更进一步,带选择性的SSM实际上可以被视为一种广义的线性注意力机制。
注意力机制视角
研究团队提出了一种更抽象的方式来描述注意力机制的本质。他们发现,任何带有掩码的注意力机制都可以表示为4个张量的缩并(Contraction):
$Attention(Q, K, V, L. = \sum_{i,j} Q_i K_j V_j L_{ij}$✅
其中,Q. K、V分别对应注意力中的query、key、value,而L则对应掩码矩阵。基于这一联系,团队提出了”结构化掩码注意力”(Structured Masked Attention, SMA)的概念。当注意力的掩码矩阵是半可分离的,它就与SSM等价。✅
结构化状态空间二元性(SSD)
基于上述发现,作者进一步推导出了两种等价的计算形式,这就是论文核心思想——”状态空间二元性”(Structured State Space Duality, SSD)的由来。SSD提供了一种统一的视角,将Transformer和SSM这两大主流序列建模架构联系起来。
Mamba-2的技术创新
1. 更大的状态维度
Mamba-2支持将状态维度从16扩展到256,这大大增强了模型的表示能力。更大的状态空间使得模型能够捕捉更复杂、更长期的依赖关系。
2. 高效的计算方法
新方法采用了基于块分解的矩阵乘法,充分利用了GPU的存储层次结构,从而显著提升了训练速度。这种优化使得Mamba-2能够在相同的计算资源下处理更大规模的数据。
3. 架构改进
Mamba-2在架构设计上做了多项创新:
4. SSD层的性能提升
Mamba-2中的SSD层比Mamba-1中的关联扫描操作快得多。这一改进使得研究团队能够增加状态维度,从而提高模型质量,同时不会显著增加计算成本。
实验结果
Mamba-2在多项任务上展现出了优异的性能:
Mamba-2的潜在应用
Mamba-2的出现为多个领域带来了新的可能性:
结论与展望
Mamba-2的出现不仅带来了性能上的提升,更重要的是,它为我们理解和设计序列模型提供了全新的视角。通过揭示Transformer和SSM之间的深层联系,Mamba-2为未来的模型设计和优化开辟了新的道路。
然而,Mamba-2的研究仍处于早期阶段,还有许多值得探索的方向:
总的来说,Mamba-2的出现为序列建模领域注入了新的活力。它不仅是对现有技术的改进,更是对整个领域认知的重塑。随着研究的深入和应用的拓展,我们有理由相信,Mamba-2将在推动人工智能技术发展中发挥重要作用。
参考文献