Fast training and inference:在练习过程中,计算和影象按序列长度线性缩放,并且在推理过程中自回归地展开模型只须要每一步恒定的时间,因为它不须要先前元素的缓存。
Long context:质量和效率一起提高了在序列长度为1M的真实数据上的性能。
模型介绍
State Space Models
离散化: 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架构为:
s B s_B sB, s C s_C sC, s Δ s_{\Delta} sΔ都是线性层,以x为输入,如下图。
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