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
计算： 模型有两种计算模式：
- 线性递归：用于高效的自回归推理(其中输入在某一时刻被看作一个时间步)
- 全局卷积：用于高效的可并行化练习(其中整个输入序列被提前看到)
为什么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
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
选择性SSM是独立的序列变换，可以灵活地集成到神经网络中。
激活函数使用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使用扫描取代卷积，加快了练习速度
论文信息：
- 论文里面的 Δ \Delta Δ相当于RNN中的门控进制
- Mamba是输入依赖的，里面的A，B，C都是通过输入数据得到的，因此实现了选择性。
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。 |
