我们还需要Transformer中的注意力吗？
输入“/”快速插入内容
我们还需要Transformer中的注意力吗？
我们还需要Transformer中的注意力吗？
2024年7月2日创建
作者：机器之心
🏖️
状态空间模型正在兴起，注意力是否已到尽头？
最近几周，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 中隐藏状态所需能力的瓶颈。
•
整合检索记忆等功能的能力，尤其是对于长文档。这更侧重于整合复杂的信息源，而不是已经存在的长文本扩展。