状态空间模型正在兴起，注意力是否已到尽头？

最近几周，AI 社区有一个热门话题：用无注意力架构来实现语言建模。简要来说，就是机器学习社区有一个长期研究方向终于取得了实质性的进展，催生出 Mamba 两个强大的新模型：Mamba 和 StripedHyena。它们在很多方面都能比肩人们熟知的强大模型，如 Llama 2 和 Mistral 7B。这个研究方向就是无注意力架构，现在也正有越来越多的研究者和开发者开始更严肃地看待它。

近日，机器学习科学家 Nathan Lambert 发布了一篇题为《状态空间 LLM：我们需要注意力吗？》的文章，详细介绍了 2023 年无注意力模型的发展情况。他还表示：2024 年你将会有不同的语言模型架构可选。需要说明，这篇文章包含不少数学内容，但深度理解它们是值得的。鉴于这篇文章较长，所以这里先列出分节目录，以方便读者索引：

• 引言：我们为什么可能并不想使用注意力以及什么是循环神经网络。
• Mamba 模型：这种新的状态空间模型能为未来多种类别的语言模型提供功能和硬件加速。
• StripedHyena 模型：这种来自 Together AI 的新 7B 模型组合了 RNN 和 Transformer 方向的最近研究成果，表现很出色。
• Monarch Mixers 研究：这篇新论文给出了一个范例，展示了该研究的运作方式，以及为什么没有注意力或 MLP 也能成功。
• Zoology 研究：这是一个用于高效 LLM 架构研究的代码库，另外还有基于这些研究得到的模型 Based。

此外还有更多链接、阅读材料和资源。

如果你对这些内容感兴趣，可以阅读作者 Nathan Lambert 对这一领域的两位领先研究者的采访，参阅机器之心报道《谁能撼动 Transformer 统治地位？Mamba 作者谈 LLM 未来架构》。

注意力 vs 循环和状态空间模型（SSM）

本文的核心是理解不同的计算方式会怎样为模型带来不同的能力。本文关注的主题是语言，但其中的思想也适用于其它许多模态（事实上，这些新架构最早取得成功的模态是音频）。当模型的内部不同时，就会出现不同的归纳方式、训练会有新的扩展律、不同的推理时间成本、新的表达能力水平（即模型可学习的任务的复杂度）等等。架构会改变有关模型的表达方式的一切，即便数据一样也是如此。

一如既往，不同的架构选择都有各自的优劣之处。现如今最流行的 Transformer 架构的核心组件注意力有非常出色的性能和可用性，原因有很多。本文不会把这些原因都列出来；简单来说，注意力有利于处理语言任务时有一个自然的归纳偏差的模型、可以轻松在 GPU 和 TPU 上扩展训练的模型、可以高效处理大批量输入的模型（例如存储键 - 值矩阵）等等。

究其核心，注意力中有从过去每个 token 到当前 token 的映射。正是这种密集架构，让模型有能力表征许多不同的内容并关注长上下文样本。

而循环神经网络（RNN）将时间整合进模型的方式却大不相同，这是本文要讨论的主要竞争方法。这些模型会在每次遇到新的输入数据时更新一个内部状态变量（以下记为 x）原理上讲，这种内部状态可以捕获任意系统的相关长期行为，而无需与数据之间有直接的计算链接。这能让模型在计算长序列时的效率非常高，但直到最近，人们都还没能证明其在性能上能媲美基于注意力的模型。下图比较了注意力和 RNN 的计算图谱：

在讨论这些模型时，会遇到很多奇特的术语。而研究社区想做的是创造一种具有 RNN 那样的时间依赖能力，同时又能维持注意力或卷积等架构的高效训练能力的模型。为此，最近出现了许多围绕状态空间模型（SSM）的研究成果，其遵照的是状态的连续时间或离散时间演变：x'(t) = Ax (t) + Bu (t)， y (t) = Cx (t) + Du (t)。使用巧妙的线性代数或微分方程，根据它是连续时间或离散时间，控制这个状态演变的矩阵可以表示成一个一维卷积。卷积架构的效率很高，所以这是个好预兆，但除此之外，本文不会出现艰深的数学。

下面展示了其方程，来自 Mamba 论文（https://arxiv.org/abs/2312.00752 ）。除非你想成为这方面的专家，否则你只需要知道这是在连续时间中构建的（1a 和 1b），往往会被离散化（2a 和 2b），并会得到一个核 K（3a 和 3b）。从技术上讲，这是一个一维卷积。

Mamba 的 SSM 方程

作者表示，尽管可以预期这不会在 2024 年改变一切，但却有可能在 2-4 年内带来天翻地覆的改变。不同的任务将会使用不同的 LLM 架构。作者还预计 ChatGPT 这样的系统将会使用多种类型的语言模型来执行日常任务。正如本文将描述的那样，基于这种 RNN 结构构建的模型（因为一些技术原因进行了许多修改）在长上下文任务的潜在准确度和推理成本方面有明显的规模扩展优势。

如果你对语言建模和机器学习的艰深数学感兴趣，那么十二月必定是个好月份。很多理性的人都知道注意力多半会被替代，只是疑惑会被怎样替代以及何时会发生。考虑到对特定于注意力的基础设施的投资之巨，作者预计这种方面短期内还无法达到 GPT-N 或 Gemini 那样的地位。如果它成功了且注意力被放弃，那谷歌就会面临一个大麻烦，因为 TPU 不见得也能用于这些新技术（就像 TPU 已经不能很好地处理 MoE 一样）。尽管如此，SSM 及相关技术依然面临诸多挑战，而且很多东西都还没有得到概念验证，例如：

• 高效利用 GPU 的能力，这是有效扩展所需的。
• 轻松微调模型并维持大多数性能的能力。
• 执行上下文学习以及系统 prompt 等功能的能力。
• 事实上，大型 Transformer 模型的大多数参数和计算依然还是前向网络（FFN），这也是 SSM 要么使用，要么不修改的部分。
• RNN 中隐藏状态所需能力的瓶颈。
• 整合检索记忆等功能的能力，尤其是对于长文档。这更侧重于整合复杂的信息源，而不是已经存在的长文本扩展。

列表不尽于此。当你在阅读此文章时，请记住这种新的无注意力技术很有潜力，但这并不意味着它现在就能与当前最先进的技术竞争，即便它确实有几个非常亮眼的结果。这只是意味着我们还处于很早期。

最近发布的模型

Mamba 以及 RNN 的高效计算

12 月 4 日，Albert Gu 和 Tri Dao 公布了 Mamba，一个力图让无注意力 LLM 的性能比肩 Transformer 的新模型，同时还能解决长上下文场景中计算受限的问题。Mamba 有三大关键特性：

1.（数据）选择机制：「我们设计了一种简单的选择机制，可以根据输入对 SSM 的参数进行参数化。」

2. 硬件感知型算法：一个将卷积切换成对特征的扫描的开关，可让模型在已有硬件上的运行效率更高。

3. 架构：将之前 SSM 的循环与 transformer 的前向模块风格结合起来。

这三方面都涉及大量数学知识。归根结底，它们是设计用于在不导致计算低效的前提下提升 SSM 表达能力的技术。

数据选择机制能够将循环空间 B 和 C 的处理矩阵表示为输入文本 x 的函数（这也被称为移除矩阵的线性时不变性（LTI））。这能以通用性为代价提升表达能力，因为输入序列会根据领域的不同而有非常不同的行为。这些矩阵可以学习哪些输入 token 是最重要的，因此这个机制名为「选择」。

Mamba 论文给出的算法

硬件感知型组件关注的重点是如何将隐藏状态 h 保存在内存中效率最高的部分。SSM 更新中使用的核心参数（线性化的 A、B 和 C 矩阵）会被保存在一种名为 SRAM 的缓存中，这样一来，到处搬运权重就不会造成较大的计算瓶颈。下图展示了所使用的内存类型：

最后，Mamba 论文还包含一个新的模型模块，其设计灵感来自 SSM 和 Transformer 模型。本文作者并不清楚如此设计的原因，但考虑到激活和 MLP 对当前机器学习的重要性，以及当前最佳的 LLM 基本都涉及到它们，因此这种做法应当是合理的。

这一项目可被看作是基于 SSM 社区大量研究成果的集大成之作（StripedHyena 确非如此）。使用定制版的 CUDA 核，GPU 可以说是火力全开。

专门设计了架构的 CUDA 核能将推理速度提得飞快，如下所示：

最后，与 Pythia 套件相比，考虑到模型大小，其平均评估性能较低。需要指出，Pythia 在参数方面的模型效率上已经不是当前最佳了，但其做对比的是 Mistral 等模型（还在 y 轴上更上面），所以这其实是褒奖。此外，该模型不一定稳健和灵活，参见下图，但它也值得了解。

正如之前的采访提到的那样，Tri Dao 表示架构只会以更好的拟合度将扩展律曲线向上移，而数据驱动型 LLM 依然是创造最佳模型的最大要素。作者认为这会进一步限制模型架构，以便在可用计算和相关任务上获得更好的性能。这是很棒的东西。Mamba 的模型和代码地址：https://github.com/state-spaces/mamba

此外，GitHub 上还有 Mamba 的最小实现版本：https://github.com/johnma2006/mamba-minimal

真实世界性能：StripedHyena-7B

尽管之前的两个项目旨在推进 LLM 的架构发展，但 StripedHyena（SH）却在新 LLM 架构方面显得最惹人眼球，但其目标是将许多先进方法和架构（包括注意力）组合起来推进 LLM 的性能表现。

12 月 8 日，Together AI 发布了第一个模型：StripedHyena-7B。StripedHyena（SH） 让人震惊。这种新的语言模型与很多常用的语言模型表现相当。根据其博客描述（https://www.together.ai/blog/stripedhyena-7b ），可知该模型采用了一种名为 grafting（嫁接）的技术。本质上讲，就像是 Together AI 取用了不同预训练模型的模块，将它们连接到一起，并继续训练模型使其性能稳定。这篇博客中写到：

「我们嫁接了 Transformer 和 Hyena 的架构组件，并在 RedPajama 数据集的一个混合数据集上进行了训练，并用更长的上下文数据进行了增强。

Hyena 这个名称来自论文《Hyena Hierarchy: Towards Larger Convolutional Language Models》。

在 OpenLLM 排行榜任务上，StripedHyena 轻松击败了 Llama 2 和 Yi 7B！

和许多无注意力架构一样，该模型的一大卖点是长上下文性能。这篇论文使用 ZeroScrolls 基准表明 StripedHyena 在该任务上的平均 f1 分数比 Mistral 7b v0.1 高 3 分（尽管其并未在每个子类上都获胜）。尽管胜过 Mistral 的 StripedHyena 的得分也不过只有 27.5，但考虑到 GPT-4 的平均分也只有 41.7，因此这个成绩还算是不错了。

除了性能之外，这项研究的还有一大部分是围绕 LLM 的不同概念的计算效率。其博客详细说明了不同架构的模型的 Chinchilla 风格的扩展律。下图左边比较了 Llama-2 与 StripedHyena，右边则是根据预算的最优注意力比例：

与 Mamba 一样，发布的 StripedHyena 包含有关推理提升的大量细节。首先是端到端的完整速度：

其次，随着上下文长度增长，其总的内存使用量呈现次二次的扩展趋势。

鉴于该模型更注重真实世界性能，作者将其与 Mistral 7b 进行了比较。作者表示尽管他更喜欢 Mistral 的答案，但 StripedHyena 的答案是正确的，如果它早一点发布，说不定还能当一段时间的最佳模型。这表明这些新架构并没有落后太多。

不过其也有些局限。研究团队没有分享基础模型使用的数据，只是说「RedPajama 数据集的一个混合数据集，并用更长的上下文数据进行了增强」。

在他们博客的最后，他们明确表示我们应该关注 Together 在这个方向的下一步动作：

• 有更长上下文的更大模型
• 多模态支持
• 更进一步的性能优化
• 将 StripedHyena 整合进检索流程，以充分利用更长的上下文。

近期研究

这一节将介绍可能会有 Mamba 和 StripedHyena 那样影响力的新论文，但它们可能还需要进一步的研究才能达到这一步。

Monarch Mixer：没有注意力或多层感知器的模型

• 论文题目：Monarch Mixer:A Simple Sub-Quadratic GEMM-Based Architecture
• 论文：https://arxiv.org/abs/2310.12109
• 代码：https://github.com/HazyResearch/m2
• 博客：https://hazyresearch.stanford.edu/blog/2023-07-25-m2-bert

这篇论文研究的不仅仅是去除 Transformer 中的注意力，还去除了占用了大部分参数的 MLP。这类研究将在未来 6-12 个月内出现在类 Mamba 的模型中。

GEMM 是一种广义的矩阵乘法算法，是许多流行架构的基础运算，包括 FFN、RNN、LSTM 和 GRU。GEMM 能在 GPU 上非常高效地执行。可以预见，由于矩阵乘法是现代机器学习的核心，但以新的方式设置它们并不总是能成功地扩展！

我们先过一遍其抽象描述，理解下这是什么意思：

「在序列长度和模型维度方面，机器学习模型在不断得到扩展，以支持更长的上下文并取得更优的性能。但是，Transformer 等现有架构会随这些轴呈二次扩展。我们就会问：是否存在能随序列长度和模型维度呈次二次扩展的高性能架构？」

尽管现在的长上下文长度已经很长了，但其推理效率并不高。作者认为模型大小在这里不是一个核心因素，或者换一种说法：寻找一种扩展律遵循比指数幂律更小规律的架构。下面继续：

「我们提出了 Monarch Mixer (M2)，一种沿序列长度和模型维度都是次二次扩展的新架构。」Monarch 矩阵是一类简单的表达结构化矩阵，可以捕获很多线性变换，能在 GPU 上实现很高的计算效率，并次二次地扩展。」

下图展示了 Monarch 矩阵的图片，它会首先按序列长度有效地混合输入，然后再使用模型维度，而不是两者同时。

「为了验证概念，我们探索了 M2 在三个领域的性能：非因果 BERT 式的语言建模、ViT 式的图像分类、因果 GPT 式的语言建模。」

本文关注的是 GPT 式的语言建模，因为这种技术势头正盛，但无注意力架构很可能助益许多领域（就像扩散模型那样）。对于最近才涉足机器学习领域的人来说，Transformer（BERT）的双向编码器表征是值得了解的 —— 它是第一个产生大量微调版本并带来诸多好处的 Transformer 模型。继续谈 BERT 式模型和 ViT（视觉 Transformer）的性能：

「对于非因果式 BERT 式建模，M2 在下游的 GLUE 质量上可以比肩 BERT-base 和 BERT-large，同时参数量少 27%，在 4k 序列长度上也能实现高 9.1 倍的吞吐量。在 ImageNet 上，M2 只用一半的参数量就让准确度超过了 ViT-b 1%。」

现在回到 GPT 式模型：

「因果 GPT 式模型带来了一个技术挑战：通过掩码实现强制因果性会带来二次的计算瓶颈。为了减轻这个瓶颈的影响，我们基于多元多项式评估和插值对 Monarch 矩阵进行了全新的理论分析，这让我们可以将 M2 参数化成因果模型，同时保留次二次特性。使用这种参数化，在 The PILE 上，M2 能在预训练困惑度指标上比肩 360M 参数的 GPT 式 Transformer—— 这首次表明也许不使用注意力或 MLP 也能达到比肩 Transformer 的质量。」

这一段的信息量很大。本质上讲，当使用 transformer 执行推理时，需要将注意力矩阵遮掩成上三角矩阵，以便每个生成的 token 仅会看过去的 token。这在该模型的解码器部分，如果你观察 BERT 这样的编码器，你会看到一个完全激活的注意力矩阵。M2 论文中的数学缓解了这个二次瓶颈（粗略地讲，直接为 N 个生成的 token 关注 M 个上下文 token）。这很了不起，但其中的数学难以理解。但你可以在那个采访中看理解的人如何解释。

那篇论文中的模型有 360M 参数，所以这在很多 GPT2 模型的范围内，还有很长的路要走。

模型 Zoology 和 Based 模型

在这波模型发布热潮中还有 Zoology：一个在合成任务上理解和测试语言模型架构的软件库。地址：https://github.com/HazyResearch/zoology

Hazy Research 的 Christopher Re 研究组分享了两篇相关的博客文章。

第一篇文章研究了不同的架构如何管理语言建模中的关联召回（associate recall）问题，即模型检索和组合来自不同数据源或概念的信号的能力。地址：https://hazyresearch.stanford.edu/blog/2023-12-11-zoology1-analysis

简而言之，这是基于注意力的模型表现很好的任务。其中给出了一个惹眼的例子：

「我们发现一个 7000 万参数的注意力模型的表现优于一个 14 亿参数的门控卷积模型。尽管关联召回听上去像是只有内行人才懂的任务，但其在机器学习领域已经有很长的历史，并且之前已有研究表明解决这些任务与上下文学习等吸引人的功能有关。」

第二篇文章更是引起了人们关注，因为他们基于自己的发现发布了一个新模型架构：Based 地址：https://hazyresearch.stanford.edu/blog/2023-12-11-zoology2-based

这篇文章中写到：

我们在这三个维度上展示了这些性质，结果表明 Based 能够提供：

简单直观的理解：Based 的基础是一种观点，即简单的卷积和注意力就善于建模不同类型的序列。我们没有引入新的复杂性来克服它们各自的缺点，而是以一种直观的方式将各自常用的版本（短一维卷积，「尖峰」线性注意力）组合起来，从而兼取了两者之长。

高质量建模：尽管模型很简单，但评估中发现，在语言建模困惑度方面， Based 在多个规模等级上都优于完整版的 Llama-2 式的 Transformer（旋转嵌入、SwiGLU MLP 等）和现代的状态空间模型（Mamba、Hyena）。

高效的高吞吐量推理：对于使用纯 PyTorch 实现的模型，Based 的推理吞吐量比相竞争的 Transformer（一个参数量相当的 Mistral 并使用了滑动窗口注意力和 FlashAttention 2）高 4.5 倍。为了让 LLM 执行批处理任务，高吞吐量是非常关键的。

总而言之，不同的架构都有各自不同的优势和短板。

原文链接：https://www.interconnects.ai/p/llms-beyond-attention

文章来自于微信公众号 “机器之心”，作者 “Nathan Lambert”