Transformer和Mamba强强结合！最新混合架构全面开源，推理速度狂飙8倍

news2024/11/26 19:32:14

最近发现，将Mamba和Transformer模块混合使用，效果会比单独使用好很多，这是因为该方法结合了Mamba的长序列处理能力和Transformer的建模能力，可以显著提升计算效率和模型性能。

典型案例如大名鼎鼎的Jamba：Jamba利用Transformer架构的元素增强Mamba 结构化状态空间模型技术，提供了 256K 上下文窗口，吞吐量直接超了Transformer三倍。

除Jamba外，近日又有不少最新提出的Transformer结合Mamba的研究，效果都很赞，比如Mamba-2-Hybrid，推理速度比Transformer快8倍。

我从这些最新研究中挑选了8个高质量成果供同学们参考学习，全部都是已开源可复现，帮助各位寻找灵感、打磨论文。

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

An Empirical Study of Mamba-based Language Models

方法：本文主要研究了基于Mamba的选择性状态空间模型与Transformer模型的结合，提出了一种混合SSM-Transformer模型，即Mamba-2-Hybrid，通过将Mamba、自注意力和MLP层组合起来，在所有常见基准测试中都优于Transformer模型。

创新点：

提出了混合SSM-Transformer模型的概念，这种模型由Mamba-2、自注意力和MLP层组成。
提出了Mamba-2-Hybrid模型扩展到支持16K、32K和128K上下文长度的方法。在推理时，Mamba-2-Hybrid模型预测生成一个token的速度可能比Transformer快达8倍；在23个长上下文评估中，16K和32K模型与Transformer基准模型相当甚至超过其性能。
证明了Mamba-2-Hybrid模型在长上下文任务中的泛化能力。通过在Phonebook查找任务中的表现，发现将Mamba-2-Hybrid扩展到支持128K上下文可以完美执行任务，即使电话簿中包含超过150K个标记。

Dimba: Transformer-Mamba Diffusion Models

方法：Dimba 是一种全新的文本到图像扩散模型，它采用了一种结合了 Transformer 和 Mamba 元素的独特混合架构。具体来说，Dimba 在 Transformer 层和 Mamba 层之间交替顺序堆叠块，并通过交叉关注层整合条件信息，从而充分利用了两种架构范式的优势。

创新点：

引入了一个新的文本到图像扩散模型，采用了Transformer和Mamba元素的混合架构。Dimba模型通过在Transformer和Mamba层之间交替堆叠块，并通过交叉注意力层整合条件信息，充分发挥了两种架构范式的优势。
介绍了一种新的状态空间模型，即Mamba，它在处理序列数据方面具有高效性和模型灵活性的优势。Mamba在SSM中引入了时变参数，并提出了一种硬件感知的算法，实现了高效的训练和推理。

Integrating Mamba and Transformer for Long-Short Range Time Series Forecasting

方法：本文探索时间序列数据中Mamba和Transformer架构的融合潜力，提出了一种混合框架Mambaformer，通过内部整合Mamba和Transformer的优势实现长短范围的时间序列预测，实验证明Mambaformer家族在长短范围时间序列预测问题上优于Mamba和Transformer。

创新点：

作者提出了一种混合架构，该架构结合了Mamba和Transformer模型以进行时间序列预测。
核心的Mambaformer层将Mamba层和自注意力层交替排列，以结合Mamba和Transformer的优点。注意力层捕获时间序列数据中的短期依赖性，而Mamba层则增强了捕获长期依赖性的能力。这一层显著提高了Mambaformer系列在长短时范围时间序列预测中的性能。