Transformer的霸主地位，要被掀翻了？

深度学习进入新时代，变形金刚的霸主地位即将被颠覆？

2017年6月12日，Heng 空诞生，直接改变了NLP，统治了自然语言领域多年。最后，新架构将打破垄断。

Transformer虽然强大，但是有一个致命的bug:核心注意力层无法扩展到长期上下文。

今天，来自CMU和普林斯顿的研究人员释放了曼巴。这种SSM架构在语言建模上堪比变形金刚，并且可以线性扩展，同时拥有五倍的推理吞吐量！

地址:https://arxiv.org/abs/2312.00752.

第一篇论文Albert Gu说，二次注意对于信息密集型的模型是必不可少的，现在已经不需要了！

论文一出，直接炸了AI界。

Nvidia首席科学家Jim Fan表示，他一直期待有人能推翻Transformer，并感谢Albert Gu和Tri Dao多年来试图取代Transformer序列架构。

“你做的研究太酷了，一会儿就冒出来一个，你就不能停一会儿！”

“湖人球迷表示很满意曼巴这个名字！”

作者还解释了为什么将这种架构命名为Mamba:

性能碾压变压器？

Mamba源于Albert Gu之前关于“结构化状态空模型”的工作，可以看作是一个强大的循环算子。这支持序列长度的线性缩放和快速自回归解码。

地址:https://arxiv.org/abs/2111.00396.

但是，以前的递归模型的缺点是其固定大小的状态难以压缩上下文。

Mamba的主要贡献是引入了“选择性SSM”，这是对S4的简单概括，可以选择性地注意或忽略输入。

这个小小的改变——只是让一些参数成为输入的函数——就能让它立刻解决之前模型极其困难的任务。

例如，它可以无限推断重要的“联想记忆”任务的解决方案！(训练长度256，测试长度1米)

关键是这种改变涉及到一个严重的效率权衡，S4的原设计有其特定的原因。

在其他模式的真实数据上，比如DNA和音频，Mamba的训练前表现超过了之前的专业基线(比如鬣狗和生鱼片)。

值得注意的是，无论是在合成、DNA还是音频数据方面，Mamba的性能都在随着序列长度达到1M+而不断提升！

另一位作者Tri Dao介绍了如何使用硬件感知设计来应对这一挑战，以及Mamba在语言方面的强大性能。

他说，正如阿尔伯特所说，SSM的特点是它的固定大小的递归状态。如果你想达到更好的表现，你需要这个状态更大，更有表现力。

遗憾的是，由于大状态太慢，实际中无法使用递归进行计算。

过去，基于S4的SSM已经通过作出结构假设(即线性时间不变性)解决了这个问题，从而可以在不实现大状态的情况下进行等效的“卷积模式”计算。

但这一次，CMU和普林斯顿的研究人员的方法是选择性SSM，只能循环计算。

为了解决这一计算瓶颈，他们使用了其他高效的硬件感知算法(如FlashAttention)所使用的技术。

需要注意的是，对于Mamba(以及一般的SSM)来说，这种方法只能在SRAM中实现扩展状态，不能在主存中实现。

此外，scan的实现比基本的PyTorch/JAX快30倍，当序列长度变长时，比二次FlashAttention快几个数量级。

而且Mamba的LM推理速度比Transformer快5倍，因为是固定大小的循环状态(没有KV缓存)。

从经验上来说，两位作者最重要的成果是在语言建模方面，这也是SSM之前瞄准的领域(比如曼巴的前身H3)。

地址:https://arxiv.org/abs/2212.14052.

但是，这个时候我的作品还是不如《变形金刚》。而且他说，那时候还没有哪个型号能真正和微调变压器抗衡。

然而，惊喜突然来了！

当按照龙猫的缩放定律进行预训练时，Mamba的表现突然优于一个非常强大的现代“Transformer++”模型(接近于Llama模型)！

电子针灸降压仪有效果吗

经过在300B token上的训练，Mamba的性能比同类开源机型要好得多。

最后，作者得出结论，硬件感知思想可以打开一个新的架构设计。

展望未来，这种新架构能否利用围绕变形金刚构建的硬件/库？它将如何改变其他领域(基因组学、音频和视频)的序列扩展？

为此，作者还发布了一系列模型权重(最大参数可达2.8B，在300B token上训练)和快速推理代码。

项目地址:https://github.com/state-spaces/mamba.

打败《变形金刚》的架构是如何诞生的？

目前的基本模型几乎都是基于变压器架构和核心注意模块。

为了解决Transformer在处理长序列时计算效率低下的问题，人们开发了许多具有二次时间复杂度的体系结构，如线性注意、门控卷积和循环模型以及结构化状态空模型(SSM)。

然而，在处理语言时，这些框架的表现不如传统的注意力模型。

研究人员发现，这些模型的主要弱点是它们难以进行基于内容的推理，因此他们进行了几项改进:

首先，通过使SSM参数成为输入数据的函数，可以解决这种模型在处理离散数据类型时的缺点。

这使得模型能够根据当前令牌在序列长度的维度上选择性地传播或忘记信息。

其次，虽然这种调整使得模型无法使用高效的卷积，但研究人员设计了一种适应硬件的并行算法，并以循环模式实现。

研究人员将这种选择性SSM集成到一个简化的端到端神经网络架构中，该架构不需要注意机制，甚至不需要MLP(多层感知器)模块，这就是研究人员提出的Mamba。

Mamba在快速推理方面非常优秀(比变形金刚快5倍)，其性能随着序列长度的增加而线性增加，尤其是在处理长度为一百万的序列时。

作为一个通用的序列处理模型，Mamba在语言、音频、基因组学等多个领域都取得了最先进的性能。

在语言建模方面，曼巴-3B模型在预训练和后续评测中的性能可以达到变形金刚模型的两倍。

通过实证研究，研究人员验证了Mamba作为基本模型(FM)核心框架的巨大潜力。

这种潜力不仅体现在预培训的质量上，还体现在具体任务的执行上，涵盖多种模式和环境:

在重要的合成任务中，如复制、归纳等，Mamba不仅能轻松求解，还能推断出无限长(>；100万代币)。

无论是训练前质量还是下游指标，Mamba在建模音频波形和DNA序列方面都优于之前的SOTA模型，如生鱼片、鬣狗和变形金刚(例如，在具有挑战性的语音生成数据集上，FID降低了一半以上)。

在这两种情况下，它的性能随着上下文长度的增加而增加，直到一百万长度的序列。

Mamba是第一个线性时间序列模型，在预训练复杂度和下游任务评估上都可以达到Transformer级别的性能。

在将模型规模扩大到10亿个参数后，研究人员证明了Mamba的性能超过了大量的基线模型，如美洲驼。

与同体积的变压器相比，曼巴语言模型的生成吞吐量是5倍，曼巴-3B的质量相当于两倍体积的变压器(与皮媞亚-3B相比，常识推理的平均值高出4分，甚至超过皮媞亚-7B)。

选择性状态空模型

研究人员利用合成任务的直觉来刺激他们的选择机制，然后解释如何将这种机制纳入state 空模型。由此产生的时变SSM不能使用卷积，这对如何有效地计算它们提出了技术挑战。

研究人员通过在现代硬件上使用内存层次的硬件感知算法克服了这个问题。然后，研究人员描述了一个简单的SSM架构，没有注意机制，甚至没有MLP模块。最后，研究人员讨论了选择机制的一些附加属性。

动机:选择作为一种压缩手段

研究人员认为，序列建模的一个基本问题是将上下文压缩成更小的状态。他们从这个角度看待流行序列模型的权衡。

比如，注意力在某些方面非常有效，但在另一些方面效率很低，因为它根本没有压缩上下文。从这一点可以看出，自回归推理需要显式存储整个上下文(即KV cache)，这直接导致了变压器的线性时间推理和二次时间训练的缓慢。

另一方面，循环模型是有效的，因为它的状态是有限的，这意味着推理时间是恒定的，训练时间将是线性的。

然而，注意力的有效性受到这种状态压缩上下文的程度的限制。

为了理解这一原理，研究人员专注于两个合成任务的例子(下图2)。

选择性复制任务通过改变要记忆的标记的位置来修改普通复制任务。它需要内容感知推理来记住相关标签(颜色)，过滤掉不相关的标签(白色)。

诱导头任务是一个众所周知的机制，以前的研究认为它可以解释大多数LLM的情景学习能力。它需要上下文感知推理来知道何时在适当的上下文中产生正确的输出(黑色)。

这些任务揭示了LTI模式的失败模式。从循环的角度来看，它们的恒定动力学(如(2)中的(a，b)跃迁)并不能使它们从上下文中选择正确的信息，或者影响沿着输入相关序列传递的隐藏态模式。

从卷积的角度来看，众所周知，全局卷积只需要时间感知就可以解决普通的复制任务，但由于缺乏内容感知，它们在选择性复制任务中存在困难(如上图)。

更具体地说，输入和输出之间的间隔是可变的，不能用静态卷积核来建模。

总之，序贯模型的效率和有效性之间的权衡是通过它们对状态的压缩程度来表征的:一个高效模型的状态必须很小，而一个好模型的状态必须包含上下文中所有必要的信息。

相反，研究者构建的序列模型的基本原理是选择性:要么关注，要么过滤输入序列状态的上下文感知能力。

特别地，选择机制控制信息如何沿着序列维度传播或交互。

通过选择改进的SSM将选择机制纳入模型的一种方法是使影响序列相互作用的参数(如RNN的循环动力学或CNN的卷积核)依赖于输入。

算法1和2说明了研究人员使用的主要选择机制。

主要区别在于简单地把输入参数δ，B，C做成函数，整个张量形状的相关变化。

应当注意，这些参数现在具有长度维度，这意味着模型已经从时不变变为时变。

这就失去了与卷积的等价性，影响了它的效率。

简化的SSM建筑

与结构化SSM一样，选择性SSM是一种独立的序列变换，可以“灵活地并入神经网络”。

H3建筑是最著名的SSM建筑的基础，通常由线性注意力启发的积木和MLP(多层感知器)积木组成。研究人员通过将这两组同构堆栈合并为一个组件来简化这种架构(如下所示)。

这种治疗的原因是受门控注意力单元(GAU)的启发。该架构涉及通过可控的扩展因子来扩展模型维度。对于每个区块，大多数参数(3ED^2)位于线性投影中，而内部SSM贡献较小。SSM参数的数量要少得多。

研究人员重复这个区块，与标准化和剩余连接交织，形成Mamba架构。

在实验中，X总是被设置为E=2，并且两堆块被用于匹配变换器交错MHA和MLP块的122个参数。滴滴大规模裁员怎么办

研究人员使用了路斯/Swish激活功能，其动机是使门控MLP成为一种流行的“SwiGLU”变体。最后，研究人员还使用了一个可选的归一化层，这是受RetNet在类似位置使用归一化层的启发。

选择机制是一个更广泛的概念，可以以不同的方式应用，如更传统的RNN或CNN，不同的参数(如算法2中的A)，或使用不同的变换。

王者荣耀s9皮肤

经验评估

合成任务:选择性复制

复制任务是一种经典的综合任务，用于测试序列模型，尤其是循环模型的记忆能力。

LTI·SSM(线性递归和全局卷积)可以通过只关注时间而不是对数据进行推理来轻松解决这个任务。例如，构建一个长度完全正确的卷积核(图2)。

在这方面，选择性复制任务可以通过随机改变令牌的间距来防止这种捷径。

表1表明，H3和曼巴等门控架构只能部分提高性能，而选择机制(即把S4提高到S6)可以很容易地解决这个问题，尤其是在与更强大的架构结合时。

合成任务:感应头

归纳头从机械可解释性的角度来看是一个相对简单的任务，但它却能出人意料地预测LLMs的上下文学习能力。

这个任务需要模型回忆和复制:例如，模型以前在一个序列中遇到过短语“哈利·波特”，所以当“哈利”再次出现在同一序列中时，模型应该能够通过回顾历史信息来预测“波特”。

表2显示，Mamba模型，或者更准确地说，它的选择性SSM层，可以完美地完成任务，因为它可以选择性地记住相关的令牌，而忽略中间的其他令牌。

而且可以完美地推广到长度为一百万的序列，也就是训练时遇到的长度的4000倍。相比之下，其他方法的泛化能力不能超过2倍。

语言建模

研究人员将Mamba与标准的Transformer架构(GPT-3架构)和最先进的transformer (Transformer R++)进行了比较。

后者基于PaLM和LLaMa架构，其特点包括旋转嵌入、SwiGLU MLP、用RMSNorm代替LayerNorm、消除线性偏差和采用更高的学习率。

图4显示，Mamba是第一个在大约1.25亿到大约13亿的参数范围内具有与最强变压器架构(变压器R++)相当的性能的疏忽模型。

表3显示了Mamba在一系列下游零炮评估任务中的表现。

其中，Mamba在训练中使用了与皮媞亚和RWKV相同的tokenizer、数据集和训练长度(300B token)。

应该注意，曼巴和皮媞亚训练的上下文长度是2048，而RWKV是1024。

DNA建模

随着大语言模型的成功，人们开始探索基本模型范式在基因组学中的应用。

DNA由具有特定词汇的离散符号序列组成，并且需要长程依赖性来建模，因此将其比作语言。

Mamba一直作为预训练微调的FM骨干进行研究，研究背景与最近的DNA长序列模型研究相同。

在预训练方面，研究者基本遵循标准的因果语言建模(下一个令牌预测)设置。

数据集方面，基本沿用了鬣狗DNA的设定，使用HG38数据集进行预训练，该数据集由单个人类基因组组成，在训练分割中包含约45亿token(DNA碱基对)。

结果如图5(左)所示，Mamba的预训练混淆度随着模型规模的增大而稳步增加，Mamba的扩展能力优于HvenaDNA和Transformer++。

例如，当最大模型大小为≈40M参数时，曲线显示Mamba可以使用3到4倍的参数，与Transformer++和HvenaDNA模型相当。

另外，图5(右)显示，Mamba可以使用更长的上下文，甚至是长达1M的极长序列，其训练前的混乱程度会随着上下文的增加而增加。

另一方面，鬣狗DNA模型会随着序列长度的增加而退化。

从卷积的角度来看，一个很长的卷积核是聚合了一个长序列上的所有信息。

图6是类人猿DNA的分类，显示了使用具有相同上下文长度的预训练模型对长度为2 ^ 10到2 ^ 20的序列进行微调的准确性。

音频建模和生成

在音频波形处理领域，刺身架构是主要对比。该模型包括:

1.一个U-Net主干，通过两个阶段的池化操作，其中模型的维度D在每个阶段翻倍，池化因子为p，

2.S4和MLP模块在每个阶段交替使用。

对于长语境下的自回归预训练，研究人员使用标准钢琴音乐数据集——YouTube Mix对其进行评估。数据集包含4小时的钢琴独奏音乐，采样率为16000Hz。

图7示出了将训练序列的长度从8192(2 ^ 13)增加到≈1000000(2 ^ 20)同时保持计算量不变的效果。

曼巴和生鱼片(S4+MLP)基线模型都随着上下文长度的增加而稳步提高。其中，Mamba在整个过程中更胜一筹，序列越长优势越明显。

在自回归语音生成方面，使用基准语音生成数据集SC09进行评估。它由一个时长为1秒的语音段组成，采样频率为16000 Hz，包括数字“0”到“9”，特征多变。

表4给出了Mamba-UNet和一系列基准模型的自动评测结果，包括WaveNet、SampleRNN、WaveGAN、DiffWave和SaShiMi。

可以看出，小规模的Mamba模型在性能上已经超过了那些基于GAN和扩散的最先进技术的较大模型。而同参数尺度的Mamba模型在保真度上遥遥领先。

表5采用了小规模的Mamba模型，并探讨了不同架构在外部和中心阶段的综合效果。

研究发现，无论是外部块还是中心块，Mamba模型的性能都优于S4+MLP架构，而中心块的性能为Mamba >；S4+MLP & gt；MHA+MLP .

速度和内存基准测试

图8显示了扫描操作的速度(状态扩展N = 16)和Mamba端到端推理吞吐量的基准测试。

结果表明，当序列长度超过2k时，高效的SSM扫描比最好的注意机制——闪光注意-2更快。而且相比PyTorch标准的扫描实现，速度提升高达20到40倍。

由于没有键值(KV)缓存，Mamba可以支持更大的批量，使得推理吞吐量比同等大小的Transformer高4到5倍。

例如，一个未经训练的Mamba(Mamba-6.9B)，拥有69亿个参数，可以超越只有13亿个参数的Transformer模型，规模小5倍。

像大多数深度系列模型一样，使用的内存量与激活张量成正比。表15显示，Mamba的内存需求相当于优化转换器的内存需求。

125M机型在单个A100 80GB GPU上训练时的内存需求。

在论文的最后，作者说选择态空模型在构建不同领域的基本模型中的广泛应用是如此令人兴奋。

种种实验结果表明，Mamba很可能成为通用序列模型的主流框架，甚至有与Transformer一争高下的潜力。

参考资料:

https://arxiv.org/abs/2312.00752

https://twitter.com/tri_dao/status/1731728602230890895

https://twitter.com/_albertgu/status/1731727672286294400

未经允许不得转载：科技让生活更美好 » Transformer的霸主地位，要被掀翻了？