14.7 长上下文技术:从理论到工程实践

从 GPT-3 的 2K 上下文窗口到 Gemini 1.5 的 100 万词元,大语言模型的上下文处理能力在短短几年内增长了三个数量级。这种增长不仅仅是“数字变大”——它涉及注意力机制的平方复杂度瓶颈(2.5 节)、KV 缓存的线性显存增长(10.2 节)、位置编码的外推极限(4.3 节)等多个层面的深层挑战。

本节系统梳理长上下文技术的工程实现、评测方法和有效利用的核心问题,将书中分散讨论的长上下文相关内容串联成一个完整的技术图景。

14.7.1 百万级上下文的工程挑战

将上下文窗口从 4K 扩展到 1M 并非简单的参数调整,而是需要在计算、存储和通信三个维度上同时解决瓶颈。

计算瓶颈。 标准自注意力的计算量与序列长度的平方成正比。处理 128K 词元的注意力计算量是 4K 的 1024 倍。即使有 Flash Attention(10.3 节)将 IO 复杂度优化到 $O(n^2 d / M)$,纯计算量本身仍然是 $O(n^2 d)$。当序列长度达到百万级时,单个注意力层的计算就需要数百 TFLOPS——单张 GPU 在合理时间内无法完成

显存瓶颈。 KV 缓存的大小与序列长度线性增长。以 Llama 2-70B 为例,缓存 4K 词元需要约 5 GB(FP16),缓存 128K 词元则需要约 160 GB——已超过单张 A100(80 GB)的显存。百万级上下文的 KV 缓存需要超过 1 TB 的显存,必须分布到多张 GPU 上。

通信瓶颈。 当序列被分片到多个设备上时,注意力计算需要跨设备访问 KV 缓存,引入了大量的设备间通信。如何在并行计算和通信之间找到平衡,是分布式长上下文系统设计的核心难题。

14.7.2 分布式注意力:Ring Attention 与序列并行

解决百万级上下文挑战的关键技术路线是序列维度的分布式计算——将一条超长序列切分到多个设备上并行处理。

Ring Attention

Ring Attention(Liu 等人,2023 年)是分布式长上下文注意力的代表性工作。其核心思想是将 Flash Attention 的分块计算策略从单设备扩展到多设备。

Ring Attention 的工作流程如下:

  1. 序列分片:将长度为 $N$ 的序列均匀分配到 $P$ 个设备上,每个设备持有 $N/P$ 个词元的 Q、K、V

  2. 环形传递:所有设备组成一个逻辑环。在每一步中,每个设备用本地的 Q 与当前持有的 K、V 计算局部注意力(使用 Flash Attention 的分块算法),同时将 K、V 块沿环传递给下一个设备

  3. 在线聚合:经过 $P$ 步后,每个设备的 Q 已经与所有位置的 K、V 交互过。利用 Flash Attention 的在线 Softmax 技巧,各步的局部结果可以正确地增量聚合为最终结果

Ring Attention 的精妙之处在于计算与通信的重叠——当设备在计算当前 KV 块的注意力时,下一块 KV 正在通过网络传输。只要单块的计算时间大于传输时间(即计算密集,而非通信密集),通信开销就可以被完全隐藏。

在理想条件下,Ring Attention 的理论上下文长度可以随设备数量线性扩展——$P$ 个设备可以处理 $P$ 倍于单设备极限的序列长度。Google 的研究团队已使用该技术在 TPU 集群上实现了超过百万词元的上下文处理。

序列并行

序列并行(Sequence Parallelism)是另一种在序列维度上分布计算的策略,与 Ring Attention 互补。

Megatron-SP(NVIDIA)在张量并行的基础上,将非注意力操作(如 LayerNorm、Dropout)也沿序列维度分片,避免了这些操作在单设备上需要完整序列长度显存的问题。

DeepSpeed Ulysses(Microsoft,2023 年)采用了一种更直接的方法。它将输入序列沿序列维度在参与计算的 GPU 之间进行均匀分片。每个 GPU 先在本地对其分片进行 Q、K、V 投影,然后通过高效的 All-to-All 集合通信完成全局数据重组——通信后每个 GPU 持有所有词元在部分注意力头上的完整 Q、K、V,从而可以独立计算各自分配到的注意力头(相当于在注意力头维度上并行)。这种设计将通信量从 Ring Attention 的 $O(N)$(与序列长度成正比)降低到了 $O(N/P)$,在大规模集群上通信效率更高。

与 Flash Attention 的协同

分布式注意力技术与 Flash Attention 并非替代关系,而是不同层次的互补优化

  • Flash Attention 解决的是单设备内的内存效率问题——通过分块计算和 IO 优化,在 SRAM 层面避免了 $O(n^2)$ 的 HBM 读写

  • Ring Attention / 序列并行解决的是跨设备的计算分布问题——将超出单设备能力的长序列分片到多个设备

实践中,两者通常组合使用:每个设备内部使用 Flash Attention 高效处理本地的注意力分块,设备之间通过 Ring Attention 或 All-to-All 通信完成全局注意力的协调。

14.7.3 长上下文训练策略

拥有了处理长序列的工程能力后,下一个问题是:如何训练模型真正理解和利用长上下文?

渐进式长度扩展

直接在百万级长度上预训练不仅代价高昂(计算量与长度的平方成正比),而且数据获取困难——互联网上真正有意义的超长连贯文本相对稀少。因此,主流做法是渐进式训练

  1. 短序列预训练:先用标准长度(如 4K-8K)完成大规模预训练,让模型学到扎实的语言理解能力

  2. 长序列微调:在预训练完成后,用较小的学习率和长序列数据继续训练,逐步扩展上下文窗口

Llama 3.1 就是一个典型案例——预训练使用 8K 上下文,然后通过长序列微调逐步扩展到 128K。这种策略的效率比从头用长序列训练高出数个数量级。

位置编码的外推与适配

渐进式训练的关键使能技术是位置编码的外推能力。如 4.3 节 所述,RoPE 配合一系列外推技术(位置内插、NTK 感知缩放、YaRN)可以将训练长度高效地扩展到数倍甚至数十倍。

YaRN 为例,它结合了频率缩放和注意力温度调整:对不同频率的旋转分量施加不同的缩放系数——低频分量(编码远距离关系)需要更激进的缩放,高频分量(编码近距离关系)则保持相对稳定。这种差异化处理使得模型在短距离上保持原有精度的同时,获得了远距离的建模能力。

长文本数据的构造

长上下文训练的另一个挑战是高质量长文本数据的稀缺性。常用的数据构造策略包括:

  • 文档拼接:将来自同一领域或主题的多个短文档拼接为长序列,用特殊分隔符标记文档边界

  • 书籍与学术论文:自然存在的长连贯文本来源,但领域偏向严重

  • 代码仓库:将完整仓库的多个文件组织成超长上下文,文件间的引用关系提供了自然的长距离依赖

  • 合成数据:利用模型自身生成的长文本作为训练数据的补充

数据质量对长上下文能力的影响巨大——如果训练数据中的长距离信息多为冗余重复,模型可能“学到”忽略远端上下文的习惯。

14.7.4 长上下文评测方法

支持长上下文是一回事,真正有效利用长上下文是另一回事。评测方法的设计对于区分这两者至关重要。

大海捞针测试

Needle-in-a-Haystack(NIAH) 测试是最直观的长上下文评测方法。其原理很简单:

  1. 在一段很长的无关文本(“干草堆”)中,随机位置插入一条关键信息(“针”)

  2. 在上下文末尾提问,要求模型精确检索该信息

  3. 通过改变“针”的位置(从头部到尾部)和“干草堆”的长度,构建一张热力图——展示模型在不同位置和长度组合下的检索准确率

NIAH 测试的优势在于直观可理解,但它只评估了最基础的精确检索能力——相当于测试一个人能否在一本书中找到某句话,而不是理解这本书的内容。实际应用中需要的长上下文能力远不止于此。

RULER 基准

RULER(Hsieh 等人,2024 年)基准针对 NIAH 的局限性,设计了四类更全面的评测任务:

  • 检索(Retrieval):不仅是单一关键词,还包括多键检索和多值检索——在上下文中找到分散在不同位置的多条相关信息

  • 追踪(Tracing):跟踪上下文中实体之间的变量传递链、因果链等逻辑关系

  • 聚合(Aggregation):对分散在上下文中的信息进行汇总、计数、排序等统计操作

  • 问答(Question Answering):基于长上下文的综合理解和推理

RULER 的实验揭示了一个重要发现:几乎所有现有模型的有效上下文长度都远短于其声称的最大上下文长度。例如,一个标称支持 128K 上下文的模型,在 RULER 的聚合任务上可能在 32K 之后性能就开始显著下降。

其他评测基准

LongBench(Bai 等人,2023 年)是一个多任务长上下文基准,覆盖单文档问答、多文档问答、摘要、少样本学习、代码补全等 21 个任务,为中英文模型提供了全面的长上下文评测。

InfiniteBench(Zhang 等人,2024 年)专门关注 100K+ 词元的超长上下文场景,包含需要理解完整小说情节、分析极长代码等任务,挑战的是现有模型的极限。

S-NIAH(Shifted Needle-in-a-Haystack)在原始 NIAH 的基础上增加了干扰项——不仅插入一根“针”,还在其他位置插入若干“假针”(语义相似但答案不同的信息),测试模型区分和精确定位的能力。

评测的反思

所有评测方法都面临一个根本问题:通过测试不等于有效利用。 模型可能在人工构造的测试中表现良好,但在真实应用场景(如分析长文档、理解代码库)中却无法充分利用长上下文。这一差距源于真实任务的信息密度、关联复杂度和推理深度远超当前评测方法的覆盖范围。

14.7.5 有效利用长上下文的挑战

拥有一个大上下文窗口与真正有效利用这个窗口之间,存在巨大的鸿沟。

注意力分布的不均匀性

Lost in the Middle(Liu 等人,2024 年)揭示了一个引人注目的现象:当相关信息放在长上下文的中间部分时,模型的表现显著下降。注意力权重倾向于集中在上下文的开头(首因效应)和末尾(近因效应),中间位置的信息更容易被“忽视”。

这一现象的根源可能在于训练数据的分布——大多数训练样本中,开头包含任务描述或系统提示,末尾包含最新的对话轮次,这两个位置的信息密度天然较高。模型可能隐式地学到了“重要信息在两端”的启发式规则。

上下文越大是否越好

直觉上,更大的上下文窗口意味着更多信息,应该总是更好。但实际情况更为微妙:

成本与延迟。 上下文越长,计算成本越高(平方增长的注意力计算),推理延迟越大(首次生成的预填充时间与上下文长度正相关)。在实际应用中,将一份 10 页的文档放入 1M 上下文窗口与放入 32K 上下文窗口相比,前者的计算成本可能高出几十倍,但信息利用效率却未必更高。

注意力稀释。 当上下文中无关信息过多时,注意力机制可能被“稀释”——模型需要在更大的信息空间中寻找相关内容,准确性反而下降。研究表明,在某些任务上,有选择地提供最相关的上下文(如通过 RAG 检索)比一次性输入所有可用信息效果更好。

推理深度的限制。 即使模型能“看到”长上下文中的所有信息,它在多跳推理(需要综合多个分散位置的信息并进行推理链条构建)上的能力仍然有限。上下文长度的增加并不能自动提升推理深度。

上下文工程

面对这些挑战,上下文工程(Context Engineering)——即如何组织和呈现上下文信息以最大化模型的利用效率——正在成为一个重要的实践领域:

  • 关键信息前置:将最重要的信息放在上下文开头,利用“首因效应”

  • 结构化组织:使用明确的分隔符、标题和元数据标注上下文中不同部分的角色和重要性

  • 渐进式细化:先提供概要信息,再在需要时提供详细内容

  • 冗余消除:删除重复和无关信息,提高信息密度

RAG 与长上下文的关系

检索增强生成(RAG)和长上下文在本质上服务于同一目标——让模型访问更多的外部信息。两者之间不是替代关系,而是互补的:

  • RAG 的优势:精准检索、成本可控、信息可追溯、支持实时更新

  • 长上下文的优势:保持完整的信息连贯性、支持全局推理、不受检索质量限制

  • 融合策略:先用 RAG 检索最相关的文档片段,再将它们放入长上下文中供模型综合分析。这种“检索 + 长上下文”的组合方案在实践中往往优于单独使用任何一种方法

长上下文技术的终极意义不在于模型能“记住”多少信息,而在于它能否真正理解和推理这些信息之间的关系。随着模型架构(如混合注意力、SSM 等)和训练方法的持续演进,这一目标正在逐步接近现实。

最后更新于