欢迎关注我的公众号 [极智视界]，获取我的更多技术分享
大家好，我是极智视界，本文分享一下 解读Mamba对LLM基础架构的冲击 作者一定是科密吧。
邀您加入我的知识星球「极智视界」，星球内有超多好玩的项目实战源码下载，链接：t.zsxq.com/0aiNxERDq
很多时候大家都会在讨论，Transformer 到底是不是 AI 基础架构的最终解，直面这个话题，Mamba 在给出否定答案的同时直接给出了另一种解法。
Mamba，一个冲击 Transformer 的新的高性能 LLM 架构。"五倍吞吐量，性能全面包围 Transformer"，在 AI 领域 "Transformer 逐渐呈现大一统" 的情况下，Mamba 无疑是那个出来 "意图破局" 的架构，而这一切，就要源于下面这篇工作，选择状态空间下的线性时间建模。从定位上来说，Mamba 不算算法层面的，所以自然也算不上什么新算法，它的定位是对标 Transformer 的基础架构。
Mamba，不禁让人想到那场比赛后的 "Mamba Out"，怀念那个男人，作者一定是科密吧，所以你如果也是科密，这篇工作就好好学习一下吧。对于看了 10+ 年球的我来说，看到 Mamba 多少是有亿点点感性了。
回到 Mamba 的学术工作，最初的 Mamba 肯定只是面向 LLM，或者说只是考虑先征服它最擅长的自然语言处理领域，而后不久其实又有 "Mamba 入局 CV" 的工作 Vision Mamba，基于双向状态空间建模的高效视觉表征学习，做的就是将 Mamba 引入 ViT，也就是下面这个华科的工作。当然，其实这就跟当初的 Transformer 的发展很像了，区别只是当时大家对于大模型、对于 Transformer 还没有现在这么敏感，所以从 Transformer 到 ViT 也是隔了好几年，在 Transformer 验证了从 NLP 到 CV 甚至是到多模态迁移的可行性后，从 Mamba 到 Vision Mamba 真的就很快了，特别是在图像编码打 patch、对比学习打标、整体架构等基本都是可以借鉴的情况下。
Transformer 虽然是强，但有个弊端很致命，那就是 "Transformer 在长序列输入的计算效率问题"，这表现在随着序列的增长，Transformer 计算效率和算力负载呈现指数级上升。而 Mamba 就是为了解决这个问题而来，Mamba 相比于 Transformer，最大的优势是五倍于 Transformer 的推理效率，但由于架构天然不具备 Transformer 并行计算的优势，所以训练效率其实不如 Transformer，但 Mamba 本身也清楚这点，所以作者也是专门结合硬件设计了一种并行计算的方法，原话是这么说的：
Even though this change prevents the use of efficient convolutions, we design a hardware-aware parallel algorithm in recurrent mode.
具体来看 Mamba 的工作。
下面是 Mamba 的概览，结构化的 SSMs (State Space Sequence Models) 通过更高维的隐状态 h (例如 N = 4) 独立地将输入 x 的每个通道 (例如 D = 5) 映射到输出 y。先前的 SSMs 通过巧妙的交替计算路径，避免了这种大的状态 (DN，批次大小 B 和序列长度 L) 的实现，这种路径需要时间不变性，也即 (∆, A, B, C) 参数是恒定的。Mamba 的选择机制增加了输入依赖动态，这也需要一个合适的硬件感知算法，最开始的 Mamba 工作其实也只是针对 GPU 内存层次结构级别实现了扩展。
上面说到了选择机制，Mamba 对于 SSMs 改进的一个很重要的点就是在选择机制。将选择机制融入模型的一种做法是让影响序列相互作用的参数与输入有关，例如嵌入到 RNN 的循环动态或者是 CNN 的卷积核。下面的伪代码表达了 Mamba 中使用的选择机制，主要的区别只是对输入的几个参数 ∆, B, C 进行了函数化处理以及对整个张量形状进行了相关变化，区别已用红色标出。
进一步讲解选择机制，左边是复制任务的标准版本，主要特点就是输入和输出元素之间具有恒定间隔并且可以通过线性递归和全局卷积等时间恒定的模型轻松解决掉。右边顶部是选择性复制任务，主要特点是具有随机的输入间隔，并且在实现上需要能够根据输入内容选择性记忆或忽略输入的具有时间变化的模型。右边底部是 Induction Heads 任务，是一个关联性回忆的例子，它要求基于上下文来检索答案，这是 LLMs 的关键能力。
下面是 Mamba 的整体模型架构。SSMs 是独立的序列转换模块，可以融入到端到端的神经网络架构中。其实有时候也将 SSM 结构称为 SSNN，这就如同 CNN 之于线性卷积层。在 Mamba 的 block 设计中，结合了大多数 SSMs 架构基础 H3 block 以及现代神经网络架构中无处不在的 MLP block。在 Mamba 中，并没有将这两种 block 进行交错，而是简单均匀地重复于 Mamba block 中。与 H3 block 相比，Mamba 用激活函数替换了第一个乘法门 (multiplicative gate)，而与 MLP block 相比，Mamba 在其主分支上增加了一个 SSM。对于激活函数，Mamba 主要是使用 SiLU / Swish。
再来看看 Mamba 的性能表现。下面是 Zero-shot 能力测试，可以看到 Mamba 的优势还是比较明显的。
下面是效率测试，在 A100 80GB 上对比了 Mamba 和 Transformer 对于不同 Sequence length 和 Batch size 的推理耗时和吞吐量，Mamba 依旧优势明显。而下面右边的图，就是本篇内容开头那个 "五倍吞吐量，性能全面包围 Transformer" 的由来。
好了，以上分享了 解读Mamba对LLM基础架构的冲击，希望我的分享能对你的学习有一点帮助。
【公众号传送】
