Mamba 2的发布是否可以撼动Transformer模型的AI大一统的江湖地位

人工智能研究所 2024-09-02 17:01:02 阅读 78

Transformer模型

Transformer模型是深度学习领域的一种神经网络架构,特别适用于自然语言处理(NLP)任务。它由Vaswani等人在2017年的论文《Attention is All You Need》中提出。Transformer模型的关键创新在于其使用注意力机制,而不是传统的递归神经网络(RNN)或卷积神经网络(CNN)来处理序列数据。正是由于Transformer模型强大的注意力机制,让模型进行并行运算成功降低了训练时间,加速了模型的训练。

动图封面

1.架构概述

Transformer模型由两个主要部分组成:编码器(Encoder)和解码器(Decoder)。

编码器:由多个相同的编码器层(通常是6层)堆叠而成。每个编码器层包含两个子层:多头自注意力机制(Multi-Head Self-Attention)和前馈神经网络(Feed Forward Neural Network),每个子层后都有一个Add & Norm层进行残差连接和规范化。解码器:由多个相同的解码器层(通常也是6层)堆叠而成。每个解码器层包含三个子层:Masked多头自注意力机制、编码器-解码器多头注意力机制(Encoder-Decoder Attention)和前馈神经网络。和编码器一样,每个子层后都有一个Add & Norm层。

2.注意力机制

注意力机制是Transformer的核心组件,它允许模型在处理一个元素时关注输入序列中的其他元素。Transformer使用的是Scaled Dot-Product Attention

Q(Queries):查询矩阵K(Keys):键矩阵V(Values):值矩阵d_k:键向量的维度,用于缩放点积的结果

动图封面

以上过程展示了QKV三矩阵的来历,其中输入的 I 矩阵为已知矩阵,而Wq Wk Wv为未知矩阵,而注意力机制的过程便是需要优化其中的未知参数了。

动图封面

注意力机制计算的过程也是极其简单,根据注意力机制的计算公式,其Q矩阵乘以K矩阵的转置矩阵,然后乘以一个缩放系数后,再乘以v矩阵,就得到了注意力机制的计算结果。

正是注意力机制的强大算法,不仅NLP领域各种任务中可以采用 transformer 模型,而在计算机视觉任务中也成功采用了 transformer 模型,且效果显著,正是这样的模型统一,让人工智能领域一度认为 transformer 模型将会承担模型大一统的任务。

Mamba模型的兴起

随着模型越来越大,训练数据越来越庞大,其输入序列也越来越长。而我们知道 transformer 模型随着输入序列长度,其模型复杂度成平方的增加。这就无意导致了大量的内存开销,虽然 transformer 模型的优点在于可以并行运算,但是较大的内存开销无法进行模型有效的训练,如何能够找到一个既能降低内存开销,又能并行运算的模型成为了大家研究的对象。

Mamba 是一种新的状态空间模型架构,在语言建模等信息密集数据上显示出良好的性能,而以前的二次模型在 Transformers 方面存在不足。Mamba基于结构化状态空间模型的,并使用FlashAttention进行高效的硬件设计和实现。

Mamba 享有快速推理(吞吐量比 Transformer 高 5 倍)和序列长度线性缩放,并且其性能在高达百万长度序列的实际数据上得到了很多的提高。 作为通用序列模型骨干,Mamba 在语言、音频和基因组学等多种任务上实现了最先进的性能。 在语言建模方面, Mamba-3B 模型在预训练和下游任务方面都优于相同大小的 Transformer模型,并且与两倍大小的 Transformer 相匹配。

而本次 2 位作者在Mamba 的基础上,更新了Mamba模型,发布了Mamba 2 。

而Mamba 2 模型的论文就叫transformer are SSMs,transformer 模型本身就是一个状态空间模型。作者用了大量的篇幅介绍了 transformer 模型与状态空间的关系,既然两者有着密切的联系,那么 transformer 模型就可以结合状态空间模型的优点,并行运算的同时降低内存开销。

上图是不是很熟悉,这个跟Retnet模型解决的大模型不可能三角很像,都是为了解决大模型的内存开销与计算复杂度问题。而根据其Mamba 2 的注意力机制的计算公式跟标准的注意力公式就是删除了 soft max 计算,取而代之的是乘以了一个 L 矩阵。而Retnet模型是一个衰减系数矩阵D.

而L 矩阵同样是一个下三角矩阵,一来屏蔽未来的输入信息,另外一个便是帮忙输入数据计算单词与单词之间的相似性。后期我们将在动画详解 transformer 专栏中详细介绍Retnet与Mamba模型,并用动画的方式来介绍相关的技术难点

Mamba 采用状态空间模型作为其核心架构,而非 Transformer 的自注意力机制。状态空间模型通过递归方式更新隐藏状态,能够更高效地处理长序列数据。Mamba 的架构允许在训练过程中进行高度并行化,从而显著缩短训练时间。与 Transformer 的平方复杂度不同,Mamba 在序列长度上的复杂度为线性,这意味着它在处理长序列时效率更高。虽然 Mamba 旨在提高效率,但初步研究表明,它在各种序列建模任务上都能保持与 Transformer 相当的性能。

Mamba的模型框架不仅可以进行高度并行化,而可以通过状态空间的递归方式更新模型状态,这就可以让模型可以进行递归方式的训练,降低内存开销。

Mamba2 的发布证明了 transformer 模型就是一个SSMs模型,完全可以采用递归的方式降低内存开销,这个跟 retnet模型的分块并行与块与块进行递归操作类似,都成功降低了内存开销。当然Mamba2模型是否在各个人工智能领域任务上都适合使用,需要各大模型的不断适配。而 transformer 模型已经被很多大模型成功使用,是否可以进行替代,需要模型与时间的验证。

https://arxiv.org/pdf/2405.21060

更多transformer,VIT,swin tranformer

参考头条号:人工智能研究所

v号:启示AI科技

 动画详解transformer  在线教程



声明

本文内容仅代表作者观点,或转载于其他网站,本站不以此文作为商业用途
如有涉及侵权,请联系本站进行删除
转载本站原创文章,请注明来源及作者。