【论文阅读笔记】Mamba模型代码理解

0.开源代码地址

官方实现：state-spaces/mamba (github.com)

最简化实现：johnma2006/mamba-minimal: Simple, minimal implementation of the Mamba SSM in one file of PyTorch. (github.com)

直接实现：alxndrTL/mamba.py: A simple and efficient Mamba implementation in PyTorch and MLX. (github.com)

官方代码做了大量优化，目录层级较多，对于理解模型含义较难，这里老师对上面第二最简化实现的代码进行了详细注释，该代码性能比官方实现差，但是对于理解模型原理比较直白。

这段代码的主要组成部分包括模型参数类ModelArgs、完整的Mamba模型类Mamba、残差块类ResidualBlock、单个Mamba块类MambaBlock、RMSNorm归一化类以及一些辅助函数。

1算法核心

的算法图，原始论文中给出的Mamba(SSSM:Selective state Space model of )的前身S4(SSM:Structured State Space Model):

S6模块

S6模块是Mamba架构中的一个复杂组件，负责通过一系列线性变换和离散化过程处理输入序列。它在捕获序列的时间动态方面起着关键作用，这是序列建模任务(如语言建模)的一个关键方面。这里包括张量运算和自定义离散化方法来处理序列数据的复杂需求。

离散化def discretization(self)中有两行代码提出来解释，论文中离散化采用零阶保持：

：对应代码中的self.dA​

:：对应代码中的self.dB

各个张量维度如下：