在深度学习领域,Transformer架构凭借其自注意力机制几乎统治了自然语言处理、计算机视觉等众多方向。然而,当序列长度持续增长时,Transformer的二次方复杂度(O(n²))成为难以回避的瓶颈——无论是处理长文档、基因序列还是长时间视频,计算和显存的消耗都会快速膨胀。近年来,一类基于状态空间模型(State Space Models, SSMs)的新型架构悄然崛起,而其中最具代表性的工作正是Mamba。
Mamba由Albert Gu和Tri Dao提出,以线性时间复杂度和媲美Transformer的性能迅速吸引了学界与工业界的关注。它的核心创新在于将结构化状态空间模型(Structured SSM)与选择性机制(Selective Mechanism)相结合,使得模型能够根据输入内容动态调整状态更新过程,从而在长序列中保留关键信息的同时有效丢弃噪声。
从技术原理来看,传统SSM本质上是线性时不变系统:给定输入序列u(t),通过隐藏状态h(t)和输出y(t)的微分方程进行变换。离散化后得到循环表示,计算复杂度为O(n),但难以处理需要选择性记忆的任务(如上下文中的重要实体)。Mamba引入了输入依赖的参数化——让状态转移矩阵A、输入矩阵B和输出矩阵C根据当前token动态生成,这样模型就能“选择性”地记住或遗忘信息,类似于门控机制在LSTM中的作用,但更高效。
更值得一提的是,Mamba摒弃了Transformer依赖的softmax注意力,也无需像许多SSM变体那样使用卷积加速。它通过硬件友好的并行扫描算法(Parallel Scan)实现训练时的快速计算,同时在推理时保持循环递推形式,使得显存占用和计算量均与序列长度成线性关系。在实验对比中,Mamba在Pile数据集上的困惑度(perplexity)与同等参数量的Transformer相当,但在处理16k长度的序列时吞吐量提升了5倍以上。
Mamba的出现并非要完全取代Transformer,而是提供了一种在特定场景下更优的选择。当任务涉及超长序列(如整个书籍、基因组分析、时间序列预测)时,Mamba架构的线性优势无可替代。同时,其简洁的数学形式和高效的实现也为理论研究提供了新视角——状态空间模型或许能弥合连续时间动力学与离散神经网络之间的鸿沟。未来,Mamba有望与Transformer混合使用,取长补短,推动下一代基础模型的发展。
——
一个热爱技术的程序员,喜欢分享前沿AI知识和开发经验。