对 SSM 感兴趣的研究者不妨读一下这篇博士论文。
在大模型领域,Transformer 凭一己之力撑起了整个江山。但随着模型规模的扩展和需要处置惩罚的序列不断变长,Transformer 的局限性也逐渐凸显,比如其自注意力机制的算计量会随着上下文长度的增加呈平方级增长。为了克服这些缺陷,研究者们开发出了很多注意力机制的高效变体,但收效甚微。
最近,一项名为「Mamba」的研究似乎打破了这一局面,它在语言建模方面可以媲美甚至击败 Transformer。这都要归功于作家提出的一种新架构 —— 选择性形态空间模型( selective state space model),该架构是 Mamba 论文作家 Albert Gu 此前主导研发的 S4 架构(Structured State Spaces for Sequence Modeling )的一个简单泛化。
在 Mamba 论文发布后,很多研究者都对 SSM(state space model)、S4 等相关研究产生了好奇。其中,有位研究者表示自己要在飞机上把这些论文都读一下。对此,Albert Gu 给出了更好的建议:他的博士论文其实把这些进展都梳理了一下,读起来可能更有条理。
在论文摘要中,作家写到,序列模型是深度进修模型的支柱,已在科学应用领域取得了广泛成功。然而,现有的方法需要针对不同的任务、模态和能力举行广泛的专业化;存在算计效率瓶颈;难以对更复杂的序列数据(如涉及长依靠关系时)举行建模。因此,继续开发对一般序列举行建模的原则性和实用性方法仍然具备根本性的重要意义。
论文链接:https://stacks.stanford.edu/file/druid:mb976vf9362/gu_dissertation-augmented.pdf
作家在论文中阐述了一种使用形态空间模型举行深度序列建模的新方法,这是一种灵活的方法,具备理论基础,算计效率高,并能在各种数据模态和应用中取得强大的结果。
首先,作家介绍了一类具备众多表征和属性的模型,它们概括了标准深度序列模型(如循环神经网络和卷积神经网络)的优势。然而,作家表明算计这些模型可能具备应战性,并开发了在当前硬件上运行非常快速的新型结构化形态空间,无论是在扩展到长序列时还是在自回归推理等其他设置中都是如此。最后,他们提出了一个用于对连续信号举行增量建模的新颖数学框架,该框架可与形态空间模型相结合,为其赋予原则性的形态表示,并提高其对长程依靠关系的建模能力。总之,这一类新方法为机器进修模型提供了有效而多用途的构建模块,特别是在大规模处置惩罚通用序列数据方面。
以下是论文各部分简介。
深度序列模型
针对序列数据的深度进修模型可被视为围绕循环、卷积或注意力等简单机制建立的序列到序列转换。
这些基元(primitive)可以被纳入标准的深度神经网络架构,形成主要的深度序列模型系列:循环神经网络(RNN)、卷积神经网络(CNN)和 Transformers,它们表达了强大的参数化变换,可以使用标准的深度进修技术(如梯度下降反向传播)举行进修。图 1.1 和定义 1.1 展示了本论文中使用的序列模型抽象,第 2.1 节将结合实例对其举行更正式的定义。
定义 1.1(非正式)。作家使用序列模型来指代在序列 y = f_θ(x) 上的参数化映射,其中输入和输出 x、y 是 R^D 中长度为 L 的特征向量序列,θ 是通过梯度下降进修的参数。
上述每个模型系列都为机器进修带来了巨大的成功:例如,RNN 为机器翻译带来了深度进修,CNN 是第一个神经音频生成模型,而 Transformers 则彻底改变了 NLP 的广阔领域。
不过,这些模型也有其序列机制所遗留的折衷课题。例如,RNN 对于序列数据来说是一个天然的有形态模型,每个时间步只需要恒定的算计 / 存储,但训练速度慢,而且存在优化困难(如梯度消失课题),这限制了它们处置惩罚长序列的能力。CNN 专注于局部上下文,编码 shift equivariance 等特性,并具备快速、可并行训练的特点,但其序列推理成本较高,且上下文长度受到固有限制。Transformers 因其处置惩罚长程依靠关系的能力和可并行性而获得巨大成功,但在序列长度上存在二次扩展课题。另一个最新的模型系列是神经微分方程(NDE),这是一种有理论基础的数学模型,理论上可以解决连续时间课题和长期依靠关系,但效率非常低。
这些课题显示了深度序列模型面临的三大应战。
应战一:通用能力
深度进修的一个广泛目标是开发可用于各种课题的通用构建模块。序列模型为解决许多此类课题提供了一个通用框架。它们可以应用于任何可被投射为序列的环境。然而,当前的模型通常仍需要大量的专业化能力,以解决特定任务和领域的课题,或针对特定的能力。各类模型的优势分析如下:
RNN:需要快速更新隐藏形态的有形态设置,例如在线处置惩罚任务和强化进修;
CNN:对音频、图像和视频等均匀采样的感知信号举行建模;
Transformers:对语言等领域中密集、复杂的交互举行建模;
NDE:处置惩罚非典型时间序列设置,如缺失或不规则采样数据。
反之,每个模型系列都可能在其不擅长的功能方面举步维艰。
应战二:算计效率
在实践中应用深度序列模型需要算计其定义的函数(即参数化序列到序列映射),这可以有多种形式。在训练时,任务一般可以用整个输入序列的损失函数来表述,算法的核心课题是如何高效地算计前向传递。在推理时(训练完成后部署模型),设置可能会发生变化;例如,在在线处置惩罚或自回归生成设置中,输入每次只显示一个时间步,模型必须能够高效地按顺序处置惩罚这些输入。
这两种情况对不同的模型系列都提出了应战。例如,RNN 本身是序列性的,很难在 GPU 和 TPU 等现代硬件加速器上举行训练,而并行性则能使其受益。另一方面,CNN 和 Transformers 则难以举行高效的自回归推理,因为它们不是有形态的;处置惩罚单个新输入的成本可能会与模型的整个上下文大小成比例关系。更奇特的模型可能会带来额外的功能,但通常会使其算计更加困难和缓慢(如需要调用昂贵的微分方程求解器)。
应战三:长程依靠
现实世界中的序列数据可能需要推理数以万计的时间步骤。除了处置惩罚长输入所需的算计课题外,解决这一课题还需要能够对此类长程依靠(LRD)中存在的复杂交互举行建模。具体来说,困难可能来自于无法捕捉数据中的交互,比如模型的上下文窗口有限;也可能来自于优化课题,比如在循环模型中通过长算计图举行反向传播时的梯度消失课题。
由于效率、表达能力或训练能力方面的限制,长程依靠是序列模型长期以来面临的应战。所有标准模型系列,如 NDE、RNN、CNN 和 Transformers,都包括许多旨在解决这些课题的专门变体。例如对抗梯度消失的正交和 Lipschitz RNN、增加上下文大小的空洞卷积,以及日益庞大的高效注意力变体系列,这些变体降低了对序列长度的二次依靠。然而,尽管这些解决方案都是针对长程依靠设计的,但在 Long Range Arena 等具备应战性的基准测试中,它们的表现仍然不佳。
形态空间序列模型
本论文介绍了基于线性形态空间模型(SSM)的新系列深度序列模型。作家将 SSM 定义为一个简单的序列模型,它通过一个隐式的潜在形态 x (t)∈R^N 映射一个 1 维函数或序列
SSM 是一种基础科学模型,广泛应用于控制论、算计神经科学、信号处置惩罚等领域。广义上,SSM 一词指的是对潜变量如何在形态空间中演化举行建模的任何模型。这些广义的 SSM 有许多种,可以改变 x 的形态空间(如连续、离散或混合空间)、y 的观测空间、过渡动态、附加噪声过程或系统的线性度。SSM 在历史上通常指隐马尔可夫模型(HMM)和线性动力系统(LDS)的变体,如分层狄利克雷过程(HDP-HMM)和 Switching Linear Dynamical 系统(SLDS)。
方程(1.1)的形态空间模型在形态空间和动力学上都是连续的,并且是完全线性和确定性的,但还没有被用作定义 1.1 意义上的深度序列模型。本论文探讨了形态空间序列模型的诸多优点,以及如何利用它们来解决一般序列建模难题,同时克服其自身的局限性。
通用序列模型
SSM 是一种简单而基本的模型,具备许多丰富的特性。它们与 NDE、RNN 和 CNN 等模型族密切相关,实际上可以以多种形式编写,以实现通常需要专门模型才能实现的各种功能(应战一)。
SSM 是连续的。SSM 本身是一个微分方程。因此,它可以执行连续时间模型的独特应用,如模拟连续过程、处置惩罚缺失数据,以及适应不同的采样率。
SSM 是循环的。可以使用标准技术将 SSM 离散化为线性 recurrence,并在推理过程中模拟为形态循环模型,每个时间步的内存和算计量保持不变。
SSM 是卷积系统。SSM 是线性时不变系统,可显式表示为连续卷积。此外,离散时间版本可以在使用离散卷积举行训练时并行化,从而实现高效训练。
因此,SSM 是一种通用序列模型,在并行和序列环境以及各种领域(如音频、视觉、时间序列)中都能高效运行。论文第 2 章介绍了 SSM 的背景,并阐述了形态空间序列模型的这些特性。
不过,SSM 的通用性也有代价。原始 SSM 仍然面临两个额外应战 —— 也许比其他模型更严重 —— 这阻碍了它们作为深度序列模型的使用。应战包括:(1)一般 SSM 比同等大小的 RNN 和 CNN 慢得多;(2)它们在记忆长依靠关系时会很吃力,例如继承了 RNN 的梯度消失课题。
作家通过 SSM 的新算法和理论来应对这些应战。
利用结构化 SSM 举行高效算计(S4)
遗憾的是,由于形态表示 x (t) ∈ R^N 对算计和内存的要求过高(应战二),通用的 SSM 在实践中无法用作深度序列模型。
对于 SSM 的形态维度 N 和序列长度 L,仅算计完整的潜在形态 x 就需要 O (N^2L) 次运算和 O (NL) 的空间 —— 与算计总体输出的 Ω(L + N) 下界相比。因此,对于合理大小的模型(例如 N ≈ 100),SSM 使用的内存要比同等大小的 RNN 或 CNN 多出几个数量级,因此作为通用序列建模解决方案,SSM 在算计上是不切实际的。
要克服这一算计瓶颈,就必须以一种适合高效算法的方式对形态矩阵 A 施加结构。作家介绍了具备各种形式结构矩阵 A 的结构化形态空间序列模型(S4)(或简称结构化形态空间)家族,以及能以任何表示形式(如循环或卷积)高效算计 S4 模型的新算法。
论文第 3 章介绍了这些高效 S4 模型的不同类型。第一种结构使用形态矩阵的对角参数化(diagonal parameterization),它非常简单、通用,足以表示几乎所有的 SSM。然后,作家通过允许低秩校正项对其举行推广,这对于捕捉后面介绍的一类特殊的 SSM 是必要的。通过结合众多技术思想,如生成函数、线性代数变换和结构矩阵乘法的结果,作家为这两种结构开发了时间复杂度为 和空间复杂度为 O (N + L) 的算法,这对于序列模型来说基本上是严密的。
使用 HIPPO 解决长程依靠关系
即使不考虑算计课题,基本的 SSM 在实验中仍然表现不佳,而且无法建模长程依靠关系(应战三)。直观地说,其中一种解释是线性一阶 ODE 求解为指数函数,因此可能会出现梯度随序列长度呈指数级缩放的课题。这也可以从它们作为线性 recurrence 的解释中看出,这涉及到反复对一个 recurrent 矩阵举行幂运算,这就是众所周知的 RNN 梯度消失/爆炸课题的起因。
在第 4 章中,作家从 SSM 后退一步,转而研究如何从第一性原理出发,用循环模型对 LRD 举行建模。他们开发了一个名为 HIPPO 的数学框架,它形式化并解决了一个名为在线函数逼近(或记忆)的课题。这种方法旨在通过保持对连续函数历史的压缩,以逐步记忆连续函数。尽管这些方法的动机完全独立,但它们都是 SSM 的具体形式。这些最终的方法被证明是 SSM 的特定形式 —— 尽管它们的动机是完全独立的。
论文第 5 章完善了这一框架,并将其与 SSM 抽象更严格地联系起来。它引入了一个正交 SSM 概念,广泛推广了 HIPPO,并推导出更多实例和理论结果,例如如何以原则性的方式初始化所有 SSM 参数。
HIPPO 概览
考虑一个输入函数 u (t)、一个固定的概率度量 ω(t),以及 N 个正交基函数(如多项式)组成的序列。在每个时间 t,u 在时间 t 之前的历史都可以投影到这个基上,从而得到一个系数向量 x (t)∈ R^N,这个向量代表了 u 的历史相对于所提供的度量 ω 的最佳近似值。函数 u (t)∈R 映射到系数 x (t)∈R^N 的映射被称为关于度量 ω 的高阶多项式投影算子 (HIPPO)。在很多情况下,在许多情况下,其形式为 x ′ (t) = Ax (t) + Bu (t),对于 (A, B) 有封闭形式的公式。
HIPPO 和 S4 的组合
HIPPO 提供了一个数学工具来构建具备重要属性的 SSM,而 S4 是关于算计表示的。第 6 章正式将两者联系起来,并说明它们可以结合起来,以获得两个世界的最佳效果。论文表明,HIPPO 生成的用于处置惩罚长程依靠关系的特殊矩阵实际上可以用第 3 章中开发的特定结构形式来编写。这就提供了结合 HIPPO 的 S4 的具体实例,从而产生了一个具备丰富功能、非常高效并擅长长程推理的通用序列模型。
应用、消融和扩展
通用序列建模功能
第 7 章对 S4 方法在各种领域和任务中的应用举行了全面的实证验证。当 S4 方法被纳入一个通用的简单深度神经网络时,它在许多基准测试中推进了 SOTA。
特别的亮点和功能包括:
通用序列建模。在不改变架构的情况下,S4 在语音分类方面超越了音频 CNN,在时间序列预测课题上优于专门的 Informer 模型,在序列 CIFAR 方面与 2-D ResNet 相媲美,准确率超过 90%。
长程依靠。在针对高效序列模型的 LRA 基准测试中,S4 的速度与所有基线一样快,同时比所有 Transformer 变体的平均准确率高出 25% 以上。S4 是第一个解决了 LRA Path-X 任务(长度为 16384)这一难题的模型,准确率达到 96%,而之前所有工作的随机猜测准确率仅为 50%。
采样分辨率变化。与专门的 NDE 方法一样,S4 无需再训练即可适应时间序列采样频率的变化。
大规模生成建模与快速自回归生成。在 CIFAR-10 密度估计方面,S4 与最好的自回归模型(每维 2.85 比特)不相上下。在 WikiText-103 语言建模方面,S4 大幅缩小了与 Transformers 的差距(在 0.5 困惑度范围内),在无注意力模型中实现了 SOTA。与 RNN 一样,在 CIFAR-10/WikiText-103 上,S4 利用其潜在形态生成像素 /token 的速度比标准自回归模型快 60 倍。
理论消融
作家对 S4 的处置惩罚讨论了训练 SSM 的许多理论细节,例如如何仔细初始化每个参数以及如何纳入 HIPPO 框架。他们对这些细节举行了全面的实证分析和消融研究,验证了他们的 SSM 理论的各个方面。例如,他们验证了 HIPPO 大大提高了 SSM 的建模能力,在标准序列模型基准上的性能比原始 SSM 实例提高了 15%。在算法上,他们的 S4 算法比传统的 SSM 算法提高了几个数量级(例如,速度提高了 30 倍,内存使用量减少到 1/400)。
应用:音频波形生成
作为一种具备多种特性的序列建模基元,S4 可以被整合到不同的神经网络架构中,并以多种方式使用。第 8 章介绍了 S4 在原始音频波形生成中的应用,由于音频波形的采样率较高,这是一个具备应战性的课题。这一章节介绍了围绕 S4 构建的 SaShiMi 多尺度架构,该架构在包括自回归和扩散在内的多种生成建模范式中,推动了无限制音频和语音生成技术的发展。该应用突显了 S4 的灵活功能,包括高效训练、快速自回归生成和用于连续信号建模的强大归纳偏置。
扩展:用于算计机视觉的多维信号
虽然作家主要关注一维序列,但某些形式的数据本身具备更高的维度,如图像(二维)和视频(三维)。序列模型的灵活性也适用于这些环境。第 9 章介绍了 S4ND,这是 S4 从一维到多维(N-D)信号的扩展。S4ND 继承了 S4 的特性,如直接对底层连续信号建模,并具备更好地处置惩罚输入分辨率变化等相关优势,是第一个在 ImageNet 等大型视觉任务中性能具备竞争力的连续模型。
更多细节请参考原论文。
最后,借机梳理介绍几篇 SSM 研究,供大家了解、进修。
论文一:Pretraining Without Attention
论文地址:https://arxiv.org/pdf/2212.10544.pdf
机器之心报道:预训练无需注意力,扩展到4096个token不成课题,与BERT相当
论文二:Mamba: Linear-Time Sequence Modeling with Selective State Spaces
论文地址:https://arxiv.org/ftp/arxiv/papers/2312/2312.00752…
项目地址:https://github.com/state-spaces/mamba
机器之心报道:五倍吞吐量,性能全面包围Transformer:新架构Mamba引爆AI圈
围绕 Mamba,已经有一些语言模型发布,包括 mamba-130m, mamba-370m, mamba-790m, mamba-1.4b, mamba-2.8b。
HuggingFace 地址:https://huggingface.co/state-spaces
也有人做出 Mamba-Chat:
Github 地址:https://github.com/havenhq/mamba-chat
论文三:苹果等机构的论文 Diffusion Models Without Attention
论文地址:https://arxiv.org/pdf/2311.18257.pdf
机器之心报道:丢掉注意力的扩散模型:Mamba带火的SSM被苹果、康奈尔盯上了
论文四:Mamba 作家 Albert Gu 的博士论文 MODELING SEQUENCES WITH STRUCTURED STATE SPACES
论文地址:https://stacks.stanford.edu/file/druid:mb976vf9362
论文五:Long Range Language Modeling via Gated State Spaces 认为 Transformer 和 SSM 完全可以互补。
论文地址:https://arxiv.org/abs/2206.13947
论文六:DeepMind 的论文 Block-State Transformer
论文地址:https://arxiv.org/pdf/2306.09539.pdf