一举颠覆Transformer！最新Mamba结合方案刷新多个SOTA，单张GPU即可处理140k

news2024/9/29 17:26:36

还记得前段时间爆火的Jamba吗？

Jamba是世界上第一个生产级的Mamba大模型，它将基于结构化状态空间模型 (SSM) 的 Mamba 模型与 transformer 架构相结合，取两种架构之长，达到模型质量和效率兼得的效果。

在吞吐量和效率等关键衡量指标上，Jamba处理128k长上下文时吞吐量是 Mixtral 8x7B的3倍；在成本上，Jamba一共支持256k上下文，单张A100 GPU即可处理140k。

这种十分炸裂的效果得益于其作者对Mamba和Transformer两种架构的创新性结合。受此启发，为了让Mamba也可以在其他方面达到两全其美的效果，研究者们开始探索Mamba与其他技术的结合，以期解决单一模型或方法难以克服的挑战。

目前已出现不少非常值得学习的研究成果，我从中挑选了12种Mamba结合方案，都是2024最新，可借鉴的方法和创新点我做了简单介绍，已开源的代码也都整理了，方便同学们学习。

论文原文以及开源代码需要的同学看文末

结合MoE

MoE-Mamba: Efficient Selective State Space Models with Mixture of Experts

方法：本文提出了将SSM与Mamba相结合的研究方向，以便将SSM的潜力扩展到更大规模，并与现有的最先进的语言模型竞争。 MoE-Mamba通过Mamba与Mixture of Experts层的结合，实现了SSM和MoE的效率提升，并在2.35倍的训练步骤中达到了与Mamba相同的性能。

创新点：

将Mixture of Experts与State Space Models相结合，开辟了一个新的研究方向。这条道路将使得更大规模的语言模型更有效地扩展。
作者发现了两种表现相似但基于不同架构的模型之间奇怪的度量不一致情况。作者假设这一差异可能暗示了Mamba和其他SSM的潜在失效模式。

结合多模态

Fusion-Mamba for Cross-modality Object Detection

方法：本文提出了一种名为Fusion-Mamba的方法，旨在在隐藏状态空间中融合特征，这可能为跨模态特征融合开辟了一种新的范例。受到Mamba的启发，作者采用具有线性复杂度的Mamba构建了隐藏状态空间，并通过门控机制进一步改进，实现更深入和复杂的融合。

创新点：

Fusion-Mamba方法：作者引入了一种名为Fusion-Mamba的新方法，该方法首次利用Mamba进行多模态特征融合。在Fusion-Mamba中，作者设计了两个模块：State Space Channel Swapping (SSCS)模块用于浅层特征融合，Dual State Space Fusion (DSSF)模块用于在隐藏状态空间中进行深层特征融合。
2D选择性扫描（SS2D）机制：作者引入了一种名为2D选择性扫描机制，用于解决二维视觉数据和一维语言序列之间的不兼容性。SS2D机制通过将图像补丁沿四个不同方向进行扩展，生成四个独立的序列，并利用这些序列建立全局感受野。

结合SAR

Simba:Mamba augmented U-ShiftGCN for Skeletal Action Recognition in Videos

方法：论文提出了一种新的Skeleton Action Recognition (SAR)框架，将选择性状态空间模型Mamba与图数据相结合。通过在具有Shift-GCN骨干的新型编码器-解码器架构中利用Mamba，作者解决了SAR任务中有效建模长序列的挑战。与缺乏结构先验并且性能低于GCN方法的纯Transformer不同，该方法利用Mamba的力量来增强时间建模，同时保留空间信息。

创新点：

首次将选择性状态空间模型Mamba集成到骨骼动作识别（SAR）和图数据领域中。通过在新颖的编码器-解码器架构中与Shift-GCN骨干网络结合使用Mamba，解决了SAR任务中有效建模长序列的挑战。
模型Simba在三个基准SAR数据集（NTU RGB+D、NTU RGB+D120和Northwestern-UCLA）上取得了最先进的性能。

结合PM扩散

P-Mamba: Marrying Perona Malik Diffusion with Mamba for Efficient Pediatric Echocardiographic Left Ventricular Segmentation

方法：论文提出一种名为P-Mamba的模型，用于高效的儿科心脏超声检查左心室分割，该模型通过使用DWT-based PMD分支和Vision Mamba分支相结合的方法，在准确性和效率方面优于现有模型。

创新点：

P-Mamba模型：在儿科心脏超声图像中，引入了一种名为P-Mamba的创新架构，用于儿科心脏左心室的分割。该模型采用了Vision Mamba层来提高计算和内存效率，并能够捕捉全局依赖关系。同时，在基于DWT的PMD编码器分支中，引入了基于DWT的Perona-Malik扩散（PMD）块来抑制噪声，同时保留左心室的局部形态特征。
数据集：该研究使用了从Lucile Packard Children's Hospital Stanford（2014-2021）收集的1,958名儿科患者的4,467个心脏超声图像数据集，其中包括7,643个灰度2D视频剪辑和17,600个标记图像。