仅用于站内搜索，没有排版格式，具体信息请跳转上方微信公众号内链接
点击上方“小白学视觉”，选择加\“星标\“或“置顶”
重磅干货，第一时间送达
今天我们来详细研究这篇论文“Mamba:具有选择性状态空间的线性时间序列建模”
Mamba一直在人工智能界掀起波澜，被吹捧为Transformer的潜在竞争对手。到底是什么让Mamba在拥挤的序列建中脱颖而出?
在介绍之前先简要回顾一下现有的模型
Transformer:以其注意力机制而闻名，其中序列的任何部分都可以动态地与任何其他部分相互作用，特别是具有因果注意力机制的的Transformer，擅长处理序列中的单个元素。但是它们带来了显著的计算和内存成本，与序列长度的平方(L²)成比例。
循环神经网络(rnn):rnn只考虑当前输入和最后一个隐藏状态，按顺序更新隐藏状态。这种方法允许它们潜在地处理无限序列长度和恒定的内存需求。但是rnn的简单性是一个缺点，限制了它们记住长期依赖关系的能力。此外，rnn中的时间反向传播(BPTT)是内存密集型的，并且可能遭受梯度消失或爆炸的影响，尽管有LSTM等创新部分结解决了这个问题。
StateSpaceModels(S4):这些模型已经显示出很好的特性。它们提供了一种平衡，比rnn更有效地捕获远程依赖关系，同时比transformer更高效地使用内存。
选择性状态空间:Mamba建立在状态空间模型的概念之上，但引入了一个新的变化。它利用选择性状态空间，支持跨长序列更高效和有效地捕获相关信息。
线性时间复杂度:与Transformer不同，Mamba在序列长度方面以线性时间运行。这个属性使得它特别适合涉及非常长的序列的任务，而传统模型在这方面会遇到困难。
Mamba以其选择性状态空间的概念引入了传统状态空间模型的一个有趣的改进。这种方法稍微放松了标准状态空间模型的严格状态转换，使其更具适应性和灵活性（有点类似于lstm）。并且Mamba保留了状态空间模型的高效计算特性，使其能够在一次扫描中执行整个序列的前向传递-这一特性更让人想起Transformer。
在训练期间，Mamba的行为类似于Transformer，同时处理整个序列。而lstm必须一步一步地计算前向传递，即使所有输入都是已知的。在推理中，Mamba的行为更符合传统的循环模型，提供有效的序列处理。
先验状态空间模型(ssm)的一个关键限制是其刚性的、输入不变的结构。这些模型为整个序列使用一组固定参数(我们称它们为a和B)。这种结构甚至比lstm等模型更具限制性，在lstm中，信号的转换可能依赖于先前的隐藏状态和输入。
Mamba则一种范式转换，即如何计算向下一个隐藏状态的过渡？在Mamba的体系结构中，转换依赖于当前输入，这种方法在传统ssm的固定计算和循环神经网络的输入依赖动态性之间取得了平衡。
主要组成如下：
固定主干:从一个隐藏状态到下一个隐藏状态的转换仍然是一个固定的计算(由a矩阵定义)，允许跨序列的预计算。
输入相关转换:输入影响下一个隐藏状态(由B矩阵定义)的方式取决于当前输入，而不是之前的隐藏状态。与传统ssm相比，这种输入依赖性提供了更大的灵活性。
为了满足这种方法的计算需求，Mamba使用了一种硬件感知算法。该算法使用扫描操作而不是卷积来循环执行计算，这样在gpu上非常高效的。尽管输入依赖转换带来了算法复杂性，但这种效率对于保持高性能至关重要。
Mamba和选择性状态空间模型不是同义词。Mamba是一个使用选择性状态空间概念的实现。这种区别是至关重要的，因为它突出了Mamba的独特贡献:在保持计算效率的同时，使SSM框架更加灵活和响应输入。
计算中的主要瓶颈通常不是计算本身，而是数据在内存类型之间的移动。Mamba通过显著减少传输大量数据的需求来解决这个问题。它通过直接在SRAM中执行算法的关键部分(如离散化和递归计算)来实现，从而减少延迟。
还引入了一个融合选择扫描层，使其内存需求与使用flashattention的优化Transformer实现相当。这一层对于保持效率至关重要，尤其是在处理模型中依赖于输入的元素时。
Mamba代表了序列建模的重大进步，特别是在其高效使用GPU内存和计算策略方面。它具有高效率处理长序列的能力，使其成为各种应用的有前途的模型，我们下面来使用Pytorch代码来对其进复现。
导入基本库
设置标志和超参数
定义超参数和初始化
这里的超参数，如模型维度(d_model)、状态大小、序列长度和批大小。
S6模块是Mamba架构中的一个复杂组件，负责通过一系列线性变换和离散化过程处理输入序列。它在捕获序列的时间动态方面起着关键作用，这是序列建模任务(如语言建模)的一个关键方面。这里包括张量运算和自定义离散化方法来处理序列数据的复杂需求。
这个S6的模块，可以处理离散化过程和正向传播。
MambaBlock类是一个定制的神经网络模块，被设计为Mamba模型的关键构建块。它封装了几个层和操作来处理输入数据。
包括线性投影、卷积、激活函数、自定义S6模块和残差连接。该块是Mamba模型的基本组件，负责通过一系列转换处理输入序列，以捕获数据中的相关模式和特征。这些不同层和操作的组合允许MambaBlock有效地处理复杂的序列建模任务。
MambaBlock是Mamba核心功能
Mamba模型
包括一系列MambaBlock模块。每个块都顺序处理输入数据，一个块的输出作为下一个块的输入。这种顺序处理允许模型捕获输入数据中的复杂模式和关系，使其对涉及顺序建模的任务有效。多个块的堆叠是深度学习架构中的常见设计，因为它使模型能够学习数据的分层表示。
RMSNorm是一个自定义规范化层，这一层用于规范神经网络的激活，这可以帮助稳定和加快训练。
这一层的用法：
上面就是模型的全部基本代码，下面就可以进行数据准备和训练
我们自定义一个Enwiki8Dataset
pad_sequences_3d用于将一批序列填充到统一的长度，确保批中的每个序列具有相同数量的元素(或时间步长)。这在许多机器学习任务中尤其重要，因为输入数据必须具有一致的形状。
训练过程还是传统的pytorch过程：
评估函数也是一样：
最后，calculate_perplexity用于评估语言模型(如Mamba)的性能。
load_enwiki8_dataset函数用于下载和提取enwiki8数据集，该数据集通常用于对语言模型进行基准测试。
encode_dataset函数设计用于标记和编码数据集，为神经网络模型(如Mamba)处理数据集做准备。
下面就可以进行训练了
以上就是训练的完整代码
我们介绍了Mamba的概念和架构，并且从头开始构建Mamba复现，这样可以将理论转化为实践。通过这种动手的方法，可以看到Mamba序列建模方法和效率。如果你想直接使用，可以看论文提供的代码
论文地址：
https ://arxiv. org/abs/2312. 00752
论文提供的源代码：
https ://github. com/state-spaces/mamba
下载1：OpenCV-Contrib扩展模块中文版教程
在「小白学视觉」公众号后台回复：扩展模块中文教程，即可下载全网第一份OpenCV扩展模块教程中文版，涵盖扩展模块安装、SFM算法、立体视觉、目标跟踪、生物视觉、超分辨率处理等二十多章内容。
下载2：Python视觉实战项目52讲
在「小白学视觉」公众号后台回复：Python视觉实战项目，即可下载包括图像分割、口罩检测、车道线检测、车辆计数、添加眼线、车牌识别、字符识别、情绪检测、文本内容提取、面部识别等31个视觉实战项目，助力快速学校计算机视觉。
下载3：OpenCV实战项目20讲
在「小白学视觉」公众号后台回复：OpenCV实战项目20讲，即可下载含有20个基于OpenCV实现20个实战项目，实现OpenCV学习进阶。
交流群
欢迎加入公众号读者群一起和同行交流，目前有SLAM、三维视觉、传感器、自动驾驶、计算摄影、检测、分割、识别、医学影像、GAN、算法竞赛等微信群（以后会逐渐细分），请扫描下面微信号加群，备注：”昵称+学校/公司+研究方向“，例如：”张三+上海交大+视觉SLAM“。请按照格式备注，否则不予通过。添加成功后会根据研究方向邀请进入相关微信群。请勿在群内发送广告，否则会请出群，谢谢理解~
