Mamba：旨在取代Transformer的架构

2023年12月,研究者Albert Gu与Tri Dao提出了基于选择性状态空间模型(selective state space models,简称SSM)的Mamba架构。与Transformer中每个token都要与所有其他token进行比较(计算复杂度为平方级,O(N²))不同,Mamba使用一个大小固定的隐藏状态,并以线性复杂度对其进行序列化更新。其关键创新在于一种选择性机制:状态更新参数(A、B、C)会针对每个token重新计算,使模型能够自主决定保留哪些信息、丢弃哪些信息,这与注意力机制的作用类似,但不会带来平方级的计算成本。由于计算过程具有结合性,Mamba的训练可以并行化,从而解决了旧式循环神经网络(recurrent neural network,简称RNN)训练缓慢的固有问题。 不过,Mamba也存在缺陷:固定大小的状态会压缩信息,导致其在精确细节提取和复制能力方面不如显式存储所有token的Transformer。因此,Mamba并未取代Transformer,而是被用于混合模型中——大致每七层Mamba层就搭配一层注意力层。这种组合既保留了Mamba廉价的长程记忆能力,又兼顾了Transformer精确检索的优势。因此,如今许多现代模型实际上可能都在隐式地使用Mamba。