【论文笔记】Mamba:挑战Transformer地位的新架构

这篇具有很好参考价值的文章主要介绍了【论文笔记】Mamba:挑战Transformer地位的新架构。希望对大家有所帮助。如果存在错误或未考虑完全的地方,请大家不吝赐教,您也可以点击"举报违法"按钮提交疑问。

Mamba

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

摘要

背景

基础模型,它现在驱动着深度学习中大多数令人振奋的应用,几乎都是基于Transformer架构及其核心注意力模块。

存在的问题

许多subquadratic-time架构如线性注意力、门控卷积和循环模型以及structured state space models (SSMs)已经被开发,以解决Transformer在长序列上的计算效率低下问题,但它们在语言等重要模态上的表现不如注意力。

本文的做法

本文发现这些模型的一个关键弱点是它们不能执行基于内容的推理,对此,本文做了如下工作:

  • 将SSM参数作为输入的函数,以离散的模态解决了它们的弱点,允许模型根据当前token沿序列长度维度选择性地传播或遗忘信息。
  • 尽管这种变化阻碍了高效卷积的使用,但在循环模式下设计了一种硬件感知的并行算法。将这些选择性SSM集成到一个简化的端到端神经网络架构中,它不需要注意力,甚至不需要MLP模块。

Mamba具有快速推理能力(比Transformers高5倍吞吐量)和序列长度的线性伸缩能力,在实际数据上的性能提高到百万序列长度。

实验结果

作为一个通用的序列模型主干,Mamba在语言、音频和基因组学等多个模态上取得了最先进的性能。在语言建模方面,Mamba - 3B模型在预训练和下游评估方面都优于相同大小的Transformers,并且匹配了两倍大小的Transformers。

文章内容

Transformer的缺点

尽管当前的Foundatio model主要基于是Transformer,但它存在一些不足:

  • 无法对有限窗口之外的任何东西进行建模
  • 窗口长度的二次复杂度

尽管有很多关于注意力变体的研究,但这些往往都会以牺牲精度为代价,且这些研究在跨域问题上并不有效。

Structured state space sequence models (SSMs)介绍

SSMs可以理解为循环神经网络( RNNs )和卷积神经网络( CNNs )的组合,在序列长度上具有线性或近似线性的复杂度。它们在特定的数据模态上效果不错,但是它们对文本等离散、信息密集的数据建模效果较差

本文的工作

提出了一类新的选择性状态空间模型,它改进了先前的工作,以实现Transformer的建模能力,同时在序列长度上是线性复杂度。具体来说:

  • Selection Mechanism :之前的模型根据输入数据进行高效选择的能力不强,对此,本文基于输入对SSM参数进行参数化,设计了一种简单的选择机制。这使得模型可以过滤掉不相关的信息,无限期地记住相关信息。
  • Hardware-aware Algorithm:所有的先验SSMs模型必须是时间和输入不变的,才能具有计算效率,为了克服这个问题,文章通过扫描而不是卷积来循环计算模型,但不实现expanded state,以避免在GPU内存层次结构的不同级别之间进行IO访问。
  • Architecture:将之前的SSM架构的设计与Transformers的MLP模块合并为一个模块来简化先验深度序列模型架构,从而得到一个包含选择性状态空间的简单同构架构设计( Mamba )。

模型的优点

  • High quality:选择性在诸如语言和基因组学等密集模态上带来了强大的性能。
  • Fast training and inference:在训练过程中,计算和记忆按序列长度线性缩放,并且在推理过程中自回归地展开模型只需要每一步恒定的时间,因为它不需要先前元素的缓存。
  • Long context:质量和效率一起提高了在序列长度为1M的真实数据上的性能。

模型介绍

State Space Models

离散化: Structured state space sequence models (S4)将一个一维函数或者序列x(t)->y(t)通过一种隐性的隐状态h(t)进行映射,它最开始是针对连续数据的。而文本、基因序列都是离散的,因此需要将数据进行离散化,具体如下:
mamba: linear-time sequence modeling with selective state spaces,论文阅读,transformer,计算机视觉,深度学习,人工智能
离散过程为:
mamba: linear-time sequence modeling with selective state spaces,论文阅读,transformer,计算机视觉,深度学习,人工智能

关于以上公式的解释可以看这篇文章:通俗易懂解释Mamba

计算: 模型有两种计算模式:

  • 线性递归:用于高效的自回归推理(其中输入在某一时刻被看作一个时间步)
  • 全局卷积:用于高效的可并行化训练(其中整个输入序列被提前看到)

为什么S4有这两种模式,可以看这篇文章:通俗易懂解释Mamba

线性时间不变性 ( A , B , C , Δ ) (A,B,C,\Delta) (A,B,C,Δ) ( A ‾ , B ‾ ) (\overline{A},\overline{B}) (A,B)在所有时间步都是固定的,它与递归和卷积有着深刻的联系。

本文认为LTI模型在建模某些类型的数据时具有根本的局限性,因此本文将会消除LTI约束,同时克服效率瓶颈。

结构和维度: 结构化SSM的命名原因是由于高效地计算它们还需要在A矩阵上施加结构,最常用的结构形式是对角线。

A的维度为 N × N N \times N N×N,B的维度为 N × 1 N \times 1 N×1,C的维度为 N × 1 N \times 1 N×1,其中N为数的个数。输入的一个x的维度为 B × L × D B \times L \times D B×L×D,其中D为通道数,L为序列长度,B为batch数,在序列长度上计算需要O ( BLDN )的时间和内存。

一般的State Space Models:状态空间模型简单地表示了任何具有潜在状态的循环过程的概念。

它在不同的学科中被用来指称许多不同的概念,包括马尔可夫决策过程,动态因果建模,隐马尔可夫模型和线性动力系统。

SSM体系结构: SSMs是独立的序列变换,可以结合到端到端的神经网络架构中。一些知名的SSM架构为:

  • Linear attention:一个包含递归的自注意力的近似,可以看作是一个退化的线性SSM。
  • H3:可以被看作是由两个门控连接夹在一个SSM中的体系结构。还在主SSM层之前插入了一个标准的局部卷积,是一个shift-SSM。
  • Hyena:使用与H3相同的架构,但将S4层替换为MLP参数化的全局卷积
  • RetNet:在架构中添加一个额外的门,并使用更简单的SSM,允许一个可选择的可并行化计算路径,使用多头注意力( MHA )的变体代替卷积。
  • RWKV:基于另一种线性注意力近似( attention-free Transformer )设计的用于语言建模的RNN。其主要的" WKV "机制涉及LTI重现,可视为两个SSM的比率。
Selective State Space Models
本文的灵感来源

本文认为序列建模的一个基本问题是将上下文压缩到一个较小的状态。Transformer没有压缩上下文,导致推理过程缓慢;递归模型具有有限的状态,实现了常数时间推理和线性时间训练,但是它们的有效性受限于这种状态对上下文的压缩程度。

本文重点关注合成任务的两个运行例子:

  • Selective Copying task:通过改变token的位置来记忆,从而实现修改Copying task。它要求内容感知推理能够记忆相关的token并过滤掉不相关的token
  • Induction Heads task:假设可以解释LLMs的大部分上下文学习能力。它需要上下文感知的推理来知道何时在适当的上下文中产生正确的输出。

这些任务揭示了LTI模型的失败模式,具体如下:

  • recurrent view:constant dynamics不能让它们从上下文中选择正确的信息,也不能以输入依赖的方式影响沿序列传递的隐藏状态。
  • convolutional view:由于缺乏内容意识,在选择性复制任务中存在困难。更具体地说,输入到输出之间的间距是变化的,不能用静态卷积核来建模

序列模型的效率和有效性的tradeoff是由它们压缩状态的程度来表征的:

  • 高效的模型必须有一个小的状态
  • 有效的模型必须有一个包含来自上下文的所有必要信息的状态
Improving SSMs with Selection

mamba: linear-time sequence modeling with selective state spaces,论文阅读,transformer,计算机视觉,深度学习,人工智能
s B s_B sB s C s_C sC s Δ s_{\Delta} sΔ都是线性层,以x为输入,如下图。
mamba: linear-time sequence modeling with selective state spaces,论文阅读,transformer,计算机视觉,深度学习,人工智能
S6的参数有一个长度维度L,这意味着模型已经从时不变(输出跟输入的时间无关)变为时变。这失去了与卷积的等价性,影响了它的效率。
s Δ s_{\Delta} sΔ τ Δ τ_{\Delta} τΔ的选择与RNN门控机制有关.

Efficient Implementation of Selective SSMs

隐藏状态维度较大的模型应该更高效,但速度较慢,因此本文希望在不付出速度和内存成本的情况下,最大化隐藏状态维度

为了高效地运行,本文将SSM参数 ( A , B , C , Δ ) (A,B,C,\Delta) (A,B,C,Δ)直接从慢速HBM加载到快速SRAM中,在SRAM中进行离散化和递归,然后将最终输出的尺寸( B,L,D)写回HBM。
本文还避免保存反向传播所必需的中间状态。本文应用了经典的重计算技术来减少内存需求:当输入从HBM加载到SRAM时,中间状态不存储,而是在后向通道中重新计算。

A Simplified SSM Architecture

选择性SSM是独立的序列变换,可以灵活地集成到神经网络中
mamba: linear-time sequence modeling with selective state spaces,论文阅读,transformer,计算机视觉,深度学习,人工智能
激活函数使用Silu或者Swish激活函数,同时使用了Layer Norm。

Selection Mechanisms的解释

可变间距: 选择性允许过滤掉感兴趣的输入之间可能出现的不相关的噪声标记。例如’ um '等语言填充词的存在。

过滤上下文: 许多序列模型并没有随着时间的延长而得到改善,按道理应该上下文越多性能越好。一种解释是,许多序列模型在必要时不能有效地忽略不相关的上下文;一个直观的例子是全局卷积(和一般的LTI模型)。选择性模型可以在任意时刻简单地重置它们的状态以去除无关的历史,因此它们的性能原则上随上下文长度单调地提高。

边界重置:在多个独立序列拼接在一起的场景中,Transformer可以通过实例化特定的注意力掩码来保持它们之间的分离,而LTI模型则会在序列之间引入信息。选择性SSM可以在边界处重置状态。

Δ \Delta Δ的解释:一般而言, Δ \Delta Δ控制着关注或忽略多少当前输入xt的平衡。它可以泛化到RNN的门控机制,一个大的 Δ \Delta Δ重置状态h并聚焦在当前输入x上,而一个小的 Δ \Delta Δ保持状态并忽略当前输入。

A的解释:A参数也可能是选择性的,但它最终只通过其和 Δ \Delta Δ的交互作用对模型产生影响。

B和C的解释:在SSM中,修改B和C使其具有选择性,可以更细粒度地控制输入xt进入状态ht或状态进入输出yt。这可以解释为允许模型分别基于内容(输入)和上下文(隐藏状态)来调节递归动态。

结论

Mamba的作用:

  • 可以对token或者特征进行有效选择,过滤无用的信息,保留感兴趣的信息
  • Mamba在模型的效率和性能之间得到了权衡,因为它能够很好地过滤无用的信息,最大化隐藏状态维度
  • Mamba使用扫描代替卷积,加快了训练速度

论文信息:文章来源地址https://www.toymoban.com/news/detail-839241.html

  • 论文里面的 Δ \Delta Δ相当于RNN中的门控进制
  • Mamba是输入依赖的,里面的A,B,C都是通过输入数据得到的,因此实现了选择性。

到了这里,关于【论文笔记】Mamba:挑战Transformer地位的新架构的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处: 如若内容造成侵权/违法违规/事实不符,请点击违法举报进行投诉反馈,一经查实,立即删除!

领支付宝红包 赞助服务器费用

相关文章

  • 【论文阅读笔记】Mamba模型代码理解

    官方实现:state-spaces/mamba (github.com) 最简化实现:johnma2006/mamba-minimal: Simple, minimal implementation of the Mamba SSM in one file of PyTorch. (github.com) 直接实现:alxndrTL/mamba.py: A simple and efficient Mamba implementation in PyTorch and MLX. (github.com) 官方代码做了大量优化,目录层级较多,对于理解模型含

    2024年04月13日
    浏览(67)
  • 【论文笔记】Mamba: Linear-Time Sequence Modeling with Selective State Spaces

    原文链接:https://arxiv.org/abs/2312.00752 基石模型(FM)的主干网络通常是序列模型,处理任意的输入序列。但现代FM主要基于Transformer这一序列模型,及其核心的注意力。但是,自注意力仅能在上下文窗口中密集地传递信息,而无法建模窗口外部的数据;此外,其尺度与窗口长度

    2024年04月26日
    浏览(42)
  • 占有统治地位的Transformer究竟是什么

    讲个有趣的小故事 我高二那年从乙班考入了甲班,对于那时的我 偏科英语最高只有108+班级平均英语成绩125+暴躁难为人女英语老师,使我上英语课时战战兢兢。英语老师很时尚,喜欢搞花里胡哨的词语让我们放松,也很尊重我虽然暴躁但维护着我的面子不让我出丑。 当时正

    2024年01月18日
    浏览(31)
  • AIGC在全球市场中的地位:机会与挑战

    博主 默语带您 Go to New World. ✍ 个人主页—— 默语 的博客👦🏻 《java 面试题大全》 🍩惟余辈才疏学浅,临摹之作或有不妥之处,还请读者海涵指正。☕🍭 《MYSQL从入门到精通》数据库是开发者必会基础之一~ 🪁 吾期望此文有资助于尔,即使粗浅难及深广,亦备添少许微薄

    2024年02月09日
    浏览(45)
  • AI革命新篇章:法国天才团队挑战ChatGPT霸主地位

    Mistral AI: Guillaume Lample, Arthur Mensch et Timothée Lacroix. ChatGPT 的霸主地位已被三位来自法国的天才所颠覆!如上图这三个人,其中一位曾在 DeepMind 工作,另外两位来自 Meta ,他们联手为 AI 领域带来了革命性的变革 我们可以想象一下,一家仅成立 8 个月的公司,如今已经在这短暂

    2024年02月20日
    浏览(41)
  • VIT 论文精读 | transformer架构引入CV的开创性工作

    目录 目录 1. 背景 2. 方法 2.1 怎么把2D图像变成1D序列输入到transformer中 像素? 先提取特征图? 打成多个patch 2.2 transformer和卷积网络比较 2.3 结构 2.4 cls Vs 全局平均池化 2.5 位置编码 3 实验 VIT是基于transformer的在图像分类大放异彩的变体,transformer是VIT的亲爹,可以和resnet相媲

    2024年02月11日
    浏览(38)
  • Jamba: A Hybrid Transformer-Mamba Language Model

    相关链接:arXiv : hybrid architecture 、 Transformer 、 Mamba 、 mixture-of-experts (MoE) 、 language model 我们介绍了Jamba,一种新的基于新颖混合Transformer-Mamba混合专家(MoE)架构的基础大型语言模型。具体来说,Jamba交错了Transformer和Mamba层的块,享受两个模型家族的好处。在这些层

    2024年04月09日
    浏览(66)
  • Transformer 论文学习笔记

    重新学习了一下,整理了一下笔记 论文 :《Attention Is All You Need》 代码 :http://nlp.seas.harvard.edu/annotated-transformer/ 地址 :https://arxiv.org/abs/1706.03762v5 翻译 :Transformer论文翻译 特点 : 提出一种不使用 RNN、CNN,仅使用注意力机制的新模型 Transformer; 只关注句内各 token 之间的关

    2024年02月14日
    浏览(39)
  • 【论文+代码】1706.Transformer简易学习笔记

    Transformer 论文: 1706.attention is all you need! 唐宇迪解读transformer:transformer2021年前,从NLP活到CV的过程 综述:2110.Transformers in Vision: A Survey 代码讲解1: Transformer 模型详解及代码实现 - 进击的程序猿 - 知乎 代码讲解2:: Transformer代码解读(Pytorch) - 神洛的文章 - 知乎 输入:词向量(

    2024年02月09日
    浏览(43)
  • [论文笔记] Transformer-XL

            这篇论文提出的 Transformer-XL 主要是针对 Transformer 在解决  长依赖问题中受到固定长度上下文的限制 ,如 Bert 采用的 Transformer 最大上下文为 512(其中是因为计算资源的限制,不是因为位置编码,因为使用的是绝对位置编码正余弦编码)。         Transformer-X

    2024年03月12日
    浏览(67)

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

博客赞助

微信扫一扫打赏

请作者喝杯咖啡吧~博客赞助

支付宝扫一扫领取红包,优惠每天领

二维码1

领取红包

二维码2

领红包