为什么大语言模型难以处理长上下文？从 Transformer 到 Mamba

编者按： 大语言模型真的能像人类一样高效处理海量信息吗？我们今天为大家带来的这篇文章，作者揭示了大语言模型在长上下文处理中的技术挑战与未来发展路径。
文章重点聚焦于三个关键层面：首先，解析了 Transformer 模型注意力机制的计算成本问题，指出随着上下文长度增加，计算复杂度呈指数级增长；其次，探讨了 Mamba 等新兴架构在突破传统模型局限性方面的潜力；最后，强调需要跳出现有思维模式，寻找处理海量信息的创新方法。

作者 | Timothy B. Lee

编译 | 岳扬

OpenAI 在两年前推出 ChatGPT 时，其能够处理的上下文信息仅有 8,192 个 tokens¹。换言之，如果输入的文本超过大约 15 页，它就会“遗忘”最初的上下文内容。这一限制使得 ChatGPT 在处理任务时的规模和复杂度都受到了影响。

而现今的 LLMs 能力有了显著提升：

OpenAI 的 GPT-4o[1] 现在能够处理多达 128,000 个 tokens 的上下文。
Anthropic 的 Claude 3.5 Sonnet[2] 可以处理 200,000 个 tokens 的上下文。
Google 的 Gemini 1.5 Pro[3] 更是拥有 2 百万个 tokens 的上下文处理能力。

尽管如此，要想让 AI 系统达到人类水平的认知能力，我们还需要取得更多的进步。

许多人展望未来，认为 AI 将能够承担大部分甚至全部的人类工作。然而，人类在工作生涯会阅读和听到数以亿计的文字，并且还能通过视觉、听觉和嗅觉从周围环境中获取更多信息。要使 AI 达到人类智能水平，它们也需要具备处理如此大量信息的能力。

目前，处理大量信息的最流行 LLM 系统构建方法是“检索增强生成”（RAG）。这类系统会寻找与用户查询相关的文档，并将最相关的部分嵌入到 LLM 的上下文中。

尽管 RAG 系统在某些情况下能够有超越传统搜索引擎的表现，但目前这类系统仍存在诸多不足。它们只有在成功将最关键的文档嵌入 LLM 的上下文时，才能产出满意的结果。然而，用于检索这些文档的技术，通常是在向量数据库[4]中进行搜索 —— 并不够精细。如果用户提出的问题复杂或含糊不清，RAG 系统很可能会错误地检索文档，导致聊天机器人给出错误的回答。

此外，RAG 系统并未让 LLM 在处理大量文档时展现出更高级的推理能力：

例如，律师可能需要 AI 帮助审阅和总结数十万封电子邮件。
工程师可能需要 AI 分析数千小时的工厂监控视频。
医学研究者可能需要 AI 在数以万计的患者病历中识别趋势。

这些任务任何一个都可能需要超过 200 万个 tokens 的上下文处理能力。而且，我们希望 AI 系统在完成这些任务后，不是一切从头开始，而是能够像人类工作者一样，通过经验积累不断提升。计算机的超强记忆力和耐力一直是其重要优势，在 AI 时代，我们并不想放弃这些特性。但目前 LLMs 在吸收和解读大量信息的能力上，还远未能达到人类水平。

确实，LLMs 在训练过程中吸收的信息量远远超过了人类。最新的人工智能模型已经在数万亿个 tokens 上进行了训练，这远远超过了一个人一生中所能阅读或听到的信息量。然而，许多有价值的资料是保密的、具有时效性的，或者因为其他原因无法用于训练。

因此，我们希望 AI 模型在推理阶段能够阅读并记住远超 200 万个 tokens 的信息。但这并非易事。

基于 transformer 的 LLMs 的核心创新在于“注意力”机制，这是一种数学运算，使得模型能够“回顾”之前的 tokens。在 LLM 生成新 token 之前，它会执行一次注意力操作，将当前 token 与之前的所有 tokens 进行比较。这导致传统的 LLMs 在上下文增长时效率逐渐降低。

目前，许多人正在研究解决这一问题的方法，我将在本文后续部分讨论其中的一些方案。但在此之前，我需要解释一下我们是如何从一开始就形成了这样一个复杂的架构。

01 GPUs 让深度学习成为现实

个人电脑的核心 —— 中央处理单元（CPUs），曾是通过提高时钟频率来提升性能的。但进入 21 世纪初期，由于过热问题，芯片制造商大多放弃了这种提速方法。

芯片制造商转而开始研发能够同时处理多个指令的 CPU[5]。然而，它们的进步受到了需要指令按顺序执行的传统编程模式的限制。

为了充分发挥摩尔定律[6]的潜力，一种全新的架构应运而生，那就是 Nvidia 推出的 GPUs。

1999 年，Nvidia 开始销售 GPU，旨在加快 3D 游戏如《Quake III Arena》的渲染速度。这些作为 PC 扩展卡的 GPU，任务是迅速绘制构成游戏中墙壁、武器、怪物等物体的成千上万的三角形。

这种任务不需要顺序编程：屏幕上不同区域的三角形可以任意顺序绘制。因此，Nvidia 的首款 GPU[7] 并不是采用单个处理器逐个执行指令，而是拥有十几个专用核心 —— 类似于微型的CPU —— 它们并行作业，共同绘制场景。

随着摩尔定律的发展，Nvidia 制造的 GPU 计算核心数量从数十个增加到数百个，最终甚至达到数千个。人们逐渐意识到，GPU 强大的并行计算能力不仅可以用于视频游戏，还能应用于其他领域。

2012 年，多伦多大学的计算机科学家 Alex Krizhevsky、Ilya Sutskever 和 Geoffrey Hinton 利用两块 Nvidia GTX 580 GPUs[8] 训练了一个用于图像识别的神经网络。这两块 GPU 各自拥有 512 个核心，提供了巨大的计算能力，使他们能够训练出一个拥有 6000 万个参数的神经网络。他们在 ImageNet 图像分类竞赛[9]中取得了新的准确率纪录[10]（该竞赛的目标是将图像归类到 1000 个不同类别之一）。

不久之后，研究人员开始将这些技术应用于更多领域，自然语言处理便是其中之一。

02 Transformers 打破了自然语言理解的瓶颈

在 2010 年代初，循环神经网络（RNNs）是处理自然语言的主流架构。RNNs 采用逐词处理的方式。神经网络在处理完每个单词后，会更新其隐藏状态（hidden state），这是一组数字，代表了神经网络对句子当前理解的程度。

RNNs 在处理短句时表现尚可，但面对长句时就显得力不从心，更别提段落或更长的文本了。在分析长句时，RNN 有时会“遗忘”句子开头的关键词。 2014 年，计算机科学家 Dzmitry Bahdanau、KyungHyun Cho 和 Yoshua Bengio 发现[11]，通过引入一个注意力机制，允许网络“回顾”句子中的早期单词，可以提升循环神经网络的性能。

2017 年，谷歌发布了《Attention Is All You Need》[12]这篇论文，它被誉为机器学习史上最重要的论文之一。在 Bahdanau 及其团队的研究基础上，谷歌的研究人员摒弃了 RNN 及其隐藏状态的概念。他们采用的模型利用注意力机制来扫描先前的单词，以获取相关的上下文信息。

这种被谷歌命名为 transformer 的新架构，其重要性不言而喻，因为它消除了扩展语言模型的一个关键障碍。

以下动画展示了 RNNs 为何难以扩展：

在这个假想的 RNN2 中，神经网络试图预测句子中的下一个单词，预测结果展示在图表的顶部。这个神经网络由三层组成，每层用一个矩形表示。它的处理方式是线性的：必须先完成对第一个单词“How”的分析，然后将隐藏状态传递回底层，网络才能开始分析第二个单词“are”。

这种限制在机器学习算法在 CPU 上运行时还不是大问题。但当人们开始利用 GPU 的并行计算能力时，RNN 的线性架构就成为了瓶颈。

transformer 通过让神经网络能够同时“思考”输入中的所有单词，从而突破了这一限制：

如图所示，基于 transformer 的模型进行的计算量与前面图中的 RNN 模型相当。因此，在（单核）CPU 上，它的运行速度可能不会更快。但由于模型不需要在处理“are”、“you”或“doing”之前完成对“How”的分析，它可以同时处理这些单词。这意味着在拥有多个并行执行单元的 GPU 上，它的运行速度可以大幅提升。

速度提升有多大？速度的潜在提升与输入单词的数量成正比。以我的动画为例，transformer 模型处理四词输入的速度大约是 RNN 的四倍。而对于 LLMs，其输入可能包含数千个单词。因此，在强大的 GPU 支持下，基于 transformer 的模型速度可以比类似的 RNN 快出几个数量级。

你可能会问，为何不能同时用多个文档来训练RNN—— 即在文档层面而非单个单词层面利用 GPU 的并行处理能力。

这是因为训练过程的第二阶段——反向传播的限制。在这个过程中，训练软件会“逆向”工作，通过微积分调整模型的参数，以提高得出正确答案的概率。对于 RNN 来说，反向传播需要从输入的最后一个单词反向追溯到第一个单词。如下图中红色箭头所示：

反向传播需要保存前向传递中每一步的中间结果——也就是说，训练软件需要存储图表中每个矩形的输出。对于大模型，这些数据占用的空间极大，以至于无法同时并行训练大量实例。 3

简而言之，transformer 释放了 GPU 的全部处理能力，推动了语言模型规模的飞速增长。领先 LLMs 的参数量从 2018 年的数亿[13]增长到了 2020 年的数千亿[14]。由于传统的基于 RNN 的模型受到线性架构的限制，它们无法在 GPU 上高效训练，因此无法达到如此庞大的规模。

03 Transformers 模型存在扩展问题

我曾提到，在本文动画中，循环神经网络与 transformer 模型“大致完成了相等的工作量”。然而，两者的工作量并非完全一致。我们再来看看 transformer 模型的工作图：

注意到各层间那些交错的对角线箭头了吗？它们代表了注意力机制的运转。基于 transformer 的语言模型在创造新 token 前，会“审视”之前每一个已有的标记，以确定哪些最为相关。

在较小规模的上下文中，这些比较的成本微不足道。例如，对于仅有 10 个、100 个甚至 1000 个 tokens 的上下文，这些成本并不构成负担。但随着上下文长度的增加，注意力机制的计算成本也随之攀升。上下文越长，为了生成下一个 token，所需的注意力操作（以及相应的计算资源）就越多。

这导致了一个问题：注意力机制总的计算能力需求与 tokens 总数成二次方关系增长。举例来说，如果一个 10 个 tokens 的提示词需要 414,720 次注意力操作4，那么：

处理一个 100 个 tokens 的提示词，将需要 4560 万次注意力操作。
处理一个 1000 个 tokens 的提示词，将需要 46 亿次注意力操作。
处理一个 10000 个 tokens 的提示词，将需要 4600 亿次注意力操作。

这或许也解释了为何当上下文超过 128,000 个 tokens 时，谷歌会对 Gemini 1.5 Pro 的收费翻倍。因为生成第 128,001 个 token 时，需要与前面 128,000 个 tokens 进行比较，其成本远高于生成第一个、第十个或第一百个 token。

04 提升注意力的效率和可扩展性

研究者们投入了大量精力优化注意力机制。其中一条研究路径旨在最大化单个 GPU 的运算效率。

我们在前文了解到，现代 GPU 包含了成千上万的执行单元。但在 GPU 开始进行数学运算之前，它需要将数据从较慢的共享内存（即高带宽内存）转移到特定执行单元内更快的内存（即SRAM）。有时，GPU 在移动数据上耗费的时间甚至超过了执行计算的时间。

在一系列论文中15[17]，普林斯顿大学的计算机科学家 Tri Dao 及其合作者开发了 FlashAttention，这种计算注意力的方式能够最大限度地减少慢速内存操作的需求。Dao 等人的工作显著提升了现代 GPU 上 transformers 的表现。

另一条研究路径则着眼于如何在多个 GPU 上高效扩展注意力。其中一篇被广泛引用的论文介绍了环形注意力机制（ring attention）[18]，它通过将 input tokens 分成块，并将每个块分配给不同的 GPU 来工作。之所以称为环形注意力，是因为 GPU 被构想为一个环形结构，每个 GPU 将其数据传递给相邻的 GPU。

这让我想起了曾参加过的一堂交谊舞课，舞伴们围成一圈，女性保持不动，而男性则轮换舞伴。最终，每个男性都能与每位女性共舞。环形注意力的原理与之类似。"女性"代表查询（query）向量（描述每个 token 所“寻找”的内容），"男性"代表键（key）向量（描述每个 token 的特征）。键向量在一连串 GPU 中传递，依次与所有查询向量相乘。

总的来说，环形注意力机制通过在多个 GPU 间分配计算任务，使得大语言模型（LLM）能够处理更大的上下文窗口。然而，它并未降低单个注意力计算的成本。

05 RNN 能否卷土重来？

由于 RNN 拥有固定大小的隐藏状态（hidden state），因此它不会存在与 transformer 相同的扩展难题。无论是生成第一个、第一百个还是第一百万个 token，RNN 所需的计算资源都相差无几。这一点，相较于基于注意力机制的模型，RNN 具有显著优势。

尽管在 transformer 问世后，RNN 的地位有所下滑，但研究者们并未放弃，他们继续探索适合在现代 GPU 上训练的 RNN 新版本。

今年 4 月，谷歌推出了一款名为 Infini-attention[19] 的新模型。这个模型可谓是 transformer 与 RNN 的“混血儿”。Infini-attention 像传统 transformer 那样处理最近的 tokens，利用注意力机制记住它们并召回它们。

不过，Infini-attention 并未试图记住所有上下文中的 tokens。相反，它采用一种“压缩记忆（compressive memory）”来存储较旧的 tokens，这种方式与 RNN 的隐藏状态有几分相似。这种数据结构能够完美地存储和召回少量 tokens，但随着 tokens 数量的增加，召回率也会越来越低。

然而，机器学习领域的 YouTube 红人 Yannic Kilcher 对谷歌的这种做法并不感冒[20]。

“我非常愿意相信这个方法确实有效，也认同这是实现无限注意力的一种途径，但我还是持怀疑态度，”Kilcher表示。“它使用的是一种边走边存的压缩记忆方法，并没有真正学会如何存储，只是按照一种确定性的方式在存储，这意味着我们对存储的内容和方式几乎没有控制权。”

06 Mamba 会是未来吗？

在复兴循环神经网络（RNN）的众多尝试中，Mamba 架构无疑是最引人注目的。它是 2023 年 12 月发表的一篇论文[21]中公布的一种架构，其开发者是计算机科学家 Tri Dao（他也是我之前提到的 FlashAttention 的负责人）和 Albert Gu。

与传统的 RNN 一样，Mamba 并不依赖于注意力机制。它拥有一个充当“记忆”角色的隐藏状态。由于这个隐藏状态的大小是固定的，因此即使输入的提示词更长，也不会增加 Mamba 处理每个 token 的成本。

我在三月着手撰写这篇文章时，本打算深入剖析 Mamba 的架构。然而，到了五月，研究团队推出了 Mamba-2[22]，其架构较之初代 Mamba 有了显著的改变。坦白说，我一直在努力理解初代 Mamba 的原理，而对于 Mamba-2 的工作机制更是尚未完全弄清。

然而，我们需要明白的是，Mamba 有潜力将 transformer 模型的性能和传统 RNN 的效率结合起来。

在六月，Dao 和 Gu 与 Nvidia 的研究人员合作发表了一篇论文[23]，对拥有 80 亿参数的 Mamba 模型进行了评估。研究发现，Mamba 模型在多项任务中与同等规模的模型不相上下，但在“上下文学习”和“从上下文中提取信息”的能力上，Mamba 模型仍略逊一筹。

transformer 模型之所以擅长信息提取，是因为它们能够“记住”上下文中的每一个 token —— 这也是为什么随着上下文长度的增加，transformer 模型的效率会降低。而 Mamba 则试图将整个上下文压缩到一个固定大小的状态中，这意味着在处理长上下文时，它不得不舍弃一部分信息。

Nvidia 团队发现，通过采用一种混合架构，该架构将 24 个 Mamba 层与 4 个注意力层交错排列，他们获得了最佳性能。这种混合架构的表现优于单纯的 transformer 模型或单纯的 Mamba 模型。

模型需要一些注意力层来记住其早期上下文中的关键细节。但是，似乎只需要少量的注意力层就够了；其余的注意力层可以由成本更低的 Mamba 层替换，而对模型的整体性能影响很小。

在八月，一家名为 AI21 的以色列初创公司发布了其 Jamba 1.5 系列模型[24]。其中最大版本的参数数量达到了 3980 亿，使其在规模上与 Meta 的 Llama 405B 模型相当。Jamba 1.5 Large 模型的 Mamba 层数量是注意力层的七倍。因此，Jamba 1.5 Large 所需的内存远少于 Meta 和其他公司的同类模型。例如，AI21 估计 Llama 3.1 70B 需要 80 GB 的内存来跟踪 256,000 个上下文 token ，而 Jamba 1.5 Large 只需要 9 GB，这使得模型能够在性能较弱的硬件上运行。

Jamba 1.5 Large 模型的 MMLU 得分为 80，显著低于 Llama 3.1 70B 的 86 分。因此，按照这个标准，Mamba 并没有完全超越 transformer 模型。然而，这可能并不是一个完全公平的比较。像 Meta 这样的前沿实验室在训练数据和后训练基础设施上投入了大量资金，以在 MMLU 等基准测试中提高几个百分点的性能。同样的高强度优化可能会缩小 Jamba 与前沿模型之间的差距。

因此，虽然更长上下文窗口的好处显而易见，但达到这一目标的最优策略尚不明确。短期内， AI 公司可能会继续使用巧妙的效率和扩展技巧（如 FlashAttention 和 Ring Attention）来扩展标准的 LLMs。长期来看，我们可能会看到对 Mamba 以及其他无注意力架构的兴趣日益增长。或者也许有人会提出一种全新的架构，使 transformers 过时。

但我确信，仅仅依靠扩大基于 transformers 的前沿模型规模并不是一个完整的解决方案。如果我们想要能够处理数十亿个 tokens 的模型——许多人都有这样的需求，我们就需要跳出固有的思维模式，寻找新的方法。

1 有网络消息指出，ChatGPT 起初设定的上下文窗口为 4,096 个 tokens，但发布后不久的一次实验[25]显示，它能够记忆超过这个数量的信息。

2 在十年前，循环神经网络（RNN）通常会包含编码器和解码器两部分，而像 GPT-3 这样的现代大语言模型（LLM）则只有解码器。出于教学目的，我展示了一个与历史不符的、仅包含解码器的 RNN 模型，这样可以更容易地与 GPT-3 等现代 LLM 进行对比。同样的分析方法也适用于 2010 年代初的真实 RNN 模型，但那时的模型图会更加复杂。

3 GPU 能够将中间计算结果传输到它的大容量高带宽内存中。但是，由于高带宽内存（HBM）[26]的速度限制，这一操作并不会提升训练速度。

4 这是我针对拥有 1750 亿参数版本的 GPT-3 的一个初步估算，该模型包含 96 层，每层有 96 个注意力头。因此，实际上每对 tokens 之间需要进行 9,216 次注意力计算。

5 Jamba 模型是一种混合专家模型，这意味着对于任何一个 token，只有网络中的一部分（980 亿参数中的 3980 亿）会被激活和使用。

Thanks for reading!

Hope you have enjoyed and learned new things from this blog!

About the author

Timothy B. Lee

I write the newsletter Understanding AI and cohost the AI Summer podcast. Previously I was a reporter at Ars Technica, Vox, and the Washington Post. twitter.com/binarybits

END

本期互动内容 🍻

❓在您看来，长上下文处理能力对 AI 发展意味着什么？人类和 AI 在信息处理能力上最大的差距是什么？

🔗文中链接🔗

[1]https://platform.openai.com/docs/models/gp

[2]https://www.anthropic.com/news/claude-3-5-sonnet

[3]https://blog.google/technology/ai/google-gemini-update-flash-...

[4]https://en.wikipedia.org/wiki/Vector_database

[5]https://en.wikipedia.org/wiki/Multithreading_(computer_architecture)

[6]https://en.wikipedia.org/wiki/Moore%27s_law

[7]https://en.wikipedia.org/wiki/GeForce_256

[8]https://www.techpowerup.com/gpu-specs/geforce-gtx-580.c270