赞
踩
Mamba: Linear-Time Sequence Modeling with Selective State Spaces
基础模型,它现在驱动着深度学习中大多数令人振奋的应用,几乎都是基于Transformer架构及其核心注意力模块。
许多subquadratic-time架构如线性注意力、门控卷积和循环模型以及structured state space models (SSMs)已经被开发,以解决Transformer在长序列上的计算效率低下问题,但它们在语言等重要模态上的表现不如注意力。
本文发现这些模型的一个关键弱点是它们不能执行基于内容的推理,对此,本文做了如下工作:
Mamba具有快速推理能力(比Transformers高5倍吞吐量)和序列长度的线性伸缩能力,在实际数据上的性能提高到百万序列长度。
作为一个通用的序列模型主干,Mamba在语言、音频和基因组学等多个模态上取得了最先进的性能。在语言建模方面,Mamba - 3B模型在预训练和下游评估方面都优于相同大小的Transformers,并且匹配了两倍大小的Transformers。
尽管当前的Foundatio model主要基于是Transformer,但它存在一些不足:
尽管有很多关于注意力变体的研究,但这些往往都会以牺牲精度为代价,且这些研究在跨域问题上并不有效。
SSMs可以理解为循环神经网络( RNNs )和卷积神经网络( CNNs )的组合,在序列长度上具有线性或近似线性的复杂度。它们在特定的数据模态上效果不错,但是它们对文本等离散、信息密集的数据建模效果较差。
提出了一类新的选择性状态空间模型,它改进了先前的工作,以实现Transformer的建模能力,同时在序列长度上是线性复杂度。具体来说:
模型的优点:
离散化: Structured state space sequence models (S4)将一个一维函数或者序列x(t)->y(t)通过一种隐性的隐状态h(t)进行映射,它最开始是针对连续数据的。而文本、基因序列都是离散的,因此需要将数据进行离散化,具体如下:
离散过程为:
关于以上公式的解释可以看这篇文章:通俗易懂解释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架构为:
本文认为序列建模的一个基本问题是将上下文压缩到一个较小的状态。Transformer没有压缩上下文,导致推理过程缓慢;递归模型具有有限的状态,实现了常数时间推理和线性时间训练,但是它们的有效性受限于这种状态对上下文的压缩程度。
本文重点关注合成任务的两个运行例子:
这些任务揭示了LTI模型的失败模式,具体如下:
序列模型的效率和有效性的tradeoff是由它们压缩状态的程度来表征的:
s
B
s_B
sB,
s
C
s_C
sC,
s
Δ
s_{\Delta}
sΔ都是线性层,以x为输入,如下图。
S6的参数有一个长度维度L,这意味着模型已经从时不变(输出跟输入的时间无关)变为时变。这失去了与卷积的等价性,影响了它的效率。
s
Δ
s_{\Delta}
sΔ和
τ
Δ
τ_{\Delta}
τΔ的选择与RNN门控机制有关.
隐藏状态维度较大的模型应该更高效,但速度较慢,因此本文希望在不付出速度和内存成本的情况下,最大化隐藏状态维度。
为了高效地运行,本文将SSM参数
(
A
,
B
,
C
,
Δ
)
(A,B,C,\Delta)
(A,B,C,Δ)直接从慢速HBM加载到快速SRAM中,在SRAM中进行离散化和递归,然后将最终输出的尺寸( B,L,D)写回HBM。
本文还避免保存反向传播所必需的中间状态。本文应用了经典的重计算技术来减少内存需求:当输入从HBM加载到SRAM时,中间状态不存储,而是在后向通道中重新计算。
选择性SSM是独立的序列变换,可以灵活地集成到神经网络中。
激活函数使用Silu或者Swish激活函数,同时使用了Layer Norm。
可变间距: 选择性允许过滤掉感兴趣的输入之间可能出现的不相关的噪声标记。例如’ 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的作用:
论文信息:
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。