长上下文技术

标签

AI/llm/长上下文

AI/llm/Attention

字数

3963 字

阅读时间

16 分钟

上下文窗口从 4K 涨到 128K、再到百万级，不是单纯「把 max_len 调大」。它受三重约束，对应三条技术路线：

约束	表现	路线
位置编码外推失效	超出训练长度的位置从没见过	位置插值 / 频率缩放（09-位置编码）
$O (n^{2})$ 的注意力	算力与显存随长度平方增长	稀疏 / 线性 / 状态空间模型
单卡装不下	1M token 的 KV Cache 数百 GB	跨设备切分（11-长序列训练与上下文并行）

还有第四重约束容易被漏掉：评估。 长上下文的「好」不能用短上下文基准衡量（14-预训练：数据、稳定性与评估里那张基准表全是短输入），需要专门的测试方法。

上下文长度的三重约束各对应一条技术路线，第四重约束在评估侧。

约束 1  位置编码外推失效
   超出训练长度的位置从没见过
   └─▶ 路线：位置插值 / 频率缩放（不动结构，只调位置编码）

约束 2  O(n²) 的注意力
   算力与显存随长度平方增长
   └─▶ 路线：稀疏 / 线性 / 状态空间模型（改结构）

约束 3  单卡装不下
   1M token 的 KV Cache 数百 GB
   └─▶ 路线：跨设备切分（改工程调度，不改变表达能力）

约束 4  评估（最容易漏掉的一条）
   短上下文基准衡量不了长上下文
   └─▶ 路线：NIAH / RULER 这类专门方法

约束 2 与约束 3 的两条路线是正交的：改结构 与 切设备 可以叠加使用。

$O (n^{2})$ 到底贵在哪

自注意力要算完整的 $n \times n$ 分数矩阵：

Attn (Q, K, V) = softmax (\frac{Q K^{⊤}}{\sqrt{d}}) V

项	复杂度	4K → 128K 的增长
计算（注意力分数）	$O (n^{2} d)$	1024×
显存（分数矩阵）	$O (n^{2})$	1024×
显存（KV Cache，推理）	$O (n \cdot d \cdot layers)$	32×（线性）

注意两行的差别：训练时那块 $n^{2}$ 的分数矩阵是真实的显存瓶颈，而推理时瓶颈换成了线性的 KV Cache —— 它线性增长，但长度够大时绝对值依然惊人。FlashAttention 解决的是前者的 IO 问题（08-FlashAttention），把训练侧的常数降下来、使 128K 在工程上可行；它不改变渐近复杂度，也不解决 KV Cache 的体量。

结构侧：三条降低复杂度的路线

稀疏注意力

观察到的事实： $n \times n$ 的注意力权重里，大部分位置接近零。稀疏注意力只算可能有显著权重的位置对。

方案	模式	用在
Sliding Window	每个 token 只看前后固定窗口 $w$	Mistral 7B
Dilated / Strided	窗口 + 固定间隔跳采样远处 token	LongNet 一类的扩展
Global + Local	少数 token 全局可见，其余走窗口	Longformer、BigBird
Multi-pattern Heads	不同头用不同窗口	多尺度感受野

Sliding Window 的复杂度是 $O (n \cdot w)$ 而不是 $O (n^{2})$ 。 Mistral 7B 用的窗口是 4096、共 32 层，该文给的论证是：信息可以逐层向外传播，有效感受野约为 $L \times w = 32 \times 4096 \approx 131072$ 。

「有效感受野」不等于可用上下文

那条乘法的含义是「第 $L$ 层理论上能触达 $\pm L w$ 的距离」，但它假定信息在每一层都能无损地往外传一步。实践中远距离信息的保真度会逐层衰减 —— Mistral 7B 的实际可用上下文仍是 8K 量级。引用 131072 这个数时要说清它是什么。

Longformer（2020）与 BigBird（2020）在窗口之外各加了「全局 token」与「随机连接」，BigBird 给出了稀疏注意力仍具通用逼近性的理论结果。

四种稀疏模式的可见范围不同。

Sliding Window     每个 token 只看前后固定窗口 ω
   ▓▓▓▓▓●▓▓▓▓▓      ω = 4096、L = 32 层时，第 L 层的理论感受野 ≈ L × ω ≈ 131072

Dilated / Strided  在窗口内按固定间隔跳采样远处的 token
   ▓ ▓ ▓●▓ ▓ ▓      同样的预算覆盖更大范围，代价是漏掉中间的点

Global + Local     少数 token 全局可见，其余走窗口
   █ ▓▓▓▓▓●▓▓▓▓▓ █  Longformer / BigBird 的「全局 token」放在关键位置

Multi-pattern Heads  不同头用不同窗口长度
   头 1   ▓●▓          短窗口，看局部
   头 h   ▓▓▓▓▓●▓▓▓▓▓  长窗口，看全局
                      └─ 多尺度感受野，成本按最长窗口算

Sliding Window 的复杂度是 O(n·ω)，而不是 O(n²)。
Mistral 7B 用的窗口是 4096、共 32 层 —— 但「有效感受野 ≈ L × ω」假定
信息在每一层都能无损地往外传一步，实践中远距离保真度会逐层衰减，
它实际可用的上下文仍是 8K 量级。

线性注意力

把 softmax 换成一个可分解的核函数 $ϕ$ ，然后改变乘法顺序：

Attn (Q, K, V) = ϕ (Q) (ϕ (K)^{⊤} V)

$ϕ (K)^{⊤} V$ 先算出来是一个 $d \times d$ 的矩阵，与序列长度无关。复杂度从 $O (n^{2} d)$ 降到 $O (n d^{2})$ —— 当 $d ≪ n$ 时就是线性的。代价是：

表达能力受限于 $d \times d$ 的容量 —— 「状态」大小固定，无论序列多长
精确检索能力弱 —— 需要从大量历史里精确取回某一个 token 的任务上明显掉点

Performer（2020）用随机特征（FAVOR+）构造无偏的 $ϕ$ 估计是这条线的代表工作。

线性注意力的关键只是换了一次乘法顺序。

标准注意力：先算 QK^T，得到 n × n 的矩阵，再乘 V
   (n × d) @ (d × n) ──▶ (n × n) ──▶ @ (n × d) ──▶ (n × d)
                          └─ 这一项随 n² 增长，就是瓶颈所在

线性注意力：把核函数 φ 提出来，先算 φ(K)^T V，得到 d × d 的矩阵
   φ(Q) (n × d)  @  [ φ(K)^T (d × n) @ V (n × d) ]  ──▶  (n × d)
                     └─ 括号里先算 ──▶ (d × d)，与序列长度无关

复杂度从 O(n²d) 降到 O(n d²) —— 当 d ≪ n 时就是线性的。

代价两条
① 表达能力受限于 d × d 的容量：「状态」大小固定，序列再长也不变
② 精确检索能力弱：需要从大量历史里精确取回某一个 token 的任务上明显掉点

状态空间模型与 Mamba

另一条线性路线完全不用注意力：用一个固定大小的隐状态压缩历史。

h_{t} = \bar{A} h_{t - 1} + \bar{B} x_{t} y_{t} = C h_{t}

这形式上和 RNN 一样，但参数化方式不同。演进的关键节点：

代际	年份	核心变化
S4	2021	用 HiPPO 矩阵初始化 $A$ ，长程依赖建模大幅改善；LRA 的 Path-X 任务 88%（Transformer 不到 60%）。但 $A, B, C$ 不随输入变化
Mamba	2023	让 $B$ 、 $C$ 、离散化步长 $Δ$ 由输入决定（选择性机制）： $Δ = softplus (Linear (x))$ 。模型因此能决定「记住这个 token」还是「忽略它」
Mamba-2	2024	提出 SSD（State Space Duality） —— 选择性 SSM 与一类结构化掩码注意力在数学上是同一族的两面。据此改用分块算法，训练比 Mamba-1 快 2–8×，状态维度从 16 提到 64–256

Mamba 的工程要害是那个「选择性」带来的副作用。 参数随输入变化后，S4 依赖的卷积模式（FFT）就不能用了（卷积核必须时不变），只能退回逐步递推。Mamba 的解法是硬件感知的并行前缀扫描加核融合：把离散化、递推、输出投影合进一个 CUDA kernel，中间状态留在 SRAM 不回写 HBM —— 这与 FlashAttention 是同一套思路（06-GEMM 性能优化里 cp.async 那一节讲的是同一个问题的另一面）。

推理侧的差别是结构性的：

	Transformer	Mamba
每 token 推理成本	$O (n)$ （要读整个 KV Cache）	$O (1)$ （状态大小固定）
需要 KV Cache	是	否
长序列显存	随长度线性增长	恒定

该文报告 Mamba-3B 的语言建模质量与同规模 Transformer 相当，推理吞吐约 5×。

混合架构：生产环境的现实选择

纯 SSM 有一处系统性短板：精确回忆。「第 47 段里的第三个数字是多少」这类任务，注意力可以直接索引任意历史位置，SSM 只能靠固定容量的状态重建。这也是 SSM 在 few-shot 场景下分数偏低（MMLU 5-shot 有约 10 个百分点的差距）的原因。

所以工程上走的是交替堆叠：

模型	结构
Jamba（AI21, 2024）	52B MoE，Transformer 层与 Mamba 层按固定比例交替（如 1:7），256K 上下文能装进单张 80 GB 卡，显存约为同上下文纯 Transformer 的 1/3
Zamba（Zyphra）	更省：只保留一层注意力并复用权重，插在多个 Mamba 块之间

Zamba 的取舍值得记住：它认为那些注意力层的职责主要是全局信息传播，而不是学习多样化的模式 —— 所以一层共享权重就够，参数和 KV Cache 都大幅减少。

混合架构的层间比例是超参数，不是常数

「9:1 优于纯注意力」这类结论来自特定规模与数据集的消融，不能当成通用配方。选型时该看的是目标负载里精确检索的占比 —— 检索需求重就往注意力侧偏。

系统侧：跨设备切分

单卡装不下序列长度时，把序列维切开（11-长序列训练与上下文并行）：Ring Attention 在环形拓扑上传递 K/V 块，单卡显存降到 $O (n / P)$ ；DeepSpeed-Ulysses 与 Megatron CP 走的是不同的切分维度。这一层是工程调度，不改变模型的表达能力 —— 与上面两条结构侧路线是正交的。

评估：长上下文怎么算「行」

短上下文基准在此完全失效：MMLU、GSM8K 的输入都在千 token 量级。长上下文需要专门的方法。

Needle in a Haystack

最常用的方法（Greg Kamradt, 2023）：在一段长无关文本（「干草堆」）的指定位置插入一条关键信息（「针」），在末尾提问。

系统扫描两个维度 —— 序列长度（4K 到 128K+）× 插入位置（开头 / 中间 / 末尾）—— 画成热力图，颜色是检索成功率。

发现的系统性现象是 Lost in the Middle：多数模型在序列开头和末尾的检索率高，中间位置明显掉。这与位置编码的衰减、注意力权重分布、以及训练数据里「关键信息靠前或靠后」的统计偏置都有关。

NIAH 扫两个维度，画出来的热力图有一个稳定的形状。

（示意：用高 / 中高 / 中 / 低代替具体成功率，只表示相对高低）

序列长度      开头     1/4     中间     3/4     末尾
   4K   ┤     高       高      中       高      高
  16K   ┤     高       高     中低      高      高
  64K   ┤     高      中高     低      中高     高
 128K   ┤    中高      中      低       中     中高
        └─────────────────────────────────────────────
              中间那一列始终最低 ──▶ Lost in the Middle

   成因与三条都有关：位置编码的衰减、注意力权重的分布、
   训练数据里「关键信息靠前或靠后」的统计偏置。

NIAH 的局限：它只考「能不能定位并取回一条信息」，不考「能不能在长文档上推理」。
一个模型可以在 NIAH 上全绿，在 RULER 的多跳追踪上崩掉。

其他基准

基准	考什么	长度
RULER	检索 + 多跳追踪 + 聚合，比 NIAH 更全	4K–128K
LongBench	中英文长文档理解	2K–32K
InfiniteBench	超长上下文理解	100K+
SCROLLS	多种长文档任务	数千到数万

NIAH 的局限要说清：它只考「能不能定位并取回一条信息」，不考「能不能在长文档上推理」。一个模型可以在 NIAH 上全绿、在 RULER 的多跳追踪上崩掉 —— RULER 的提出动机正是 NIAH 太容易饱和。看长上下文能力要同时看两者。

三条路线怎么选

长度量级	可行方案
≤ 128K	FlashAttention + RoPE 外推（YaRN 等）已经够用，不需要改结构
128K – 1M	全局注意力成本开始难扛；混合架构（注意力 + SSM）或稀疏注意力进入视野
> 1M	纯全局注意力基本不可行；需要 SSM / 线性注意力这类恒定状态的方案，配合跨设备切分

判断的核心量是「精确检索在负载里的占比」，不是长度本身 —— 长度只决定成本，检索需求才决定能走哪条路。

走哪条路要先看长度量级，再看负载里精确检索的占比。

长度量级
  ≤ 128K
     └─ FlashAttention + RoPE 外推（YaRN 等）已经够用，不需要改结构
        这一档要解决的是 IO 常数与位置编码，不是渐近复杂度

  128K – 1M
     └─ 全局注意力的成本开始难扛
        混合架构（注意力 + SSM）或稀疏注意力进入视野
        └─ 注意力层负责精确回忆，其余层用 O(1) 状态压成本

  > 1M
     └─ 纯全局注意力基本不可行
        需要 SSM / 线性注意力这类恒定状态的方案，并配合跨设备切分

精确检索在负载里的占比
  检索需求重   ──▶ 往注意力侧偏
                   （SSM 的短板正是精确回忆，MMLU 5-shot 有约 10 个百分点的差距）
  生成与摘要为主 ──▶ 可以更多依赖 SSM，成本优先

长度决定成本，检索需求决定能走哪条路 —— 这两件事要分开问。

参考

https://arxiv.org/abs/2312.00752
https://arxiv.org/abs/2405.21060
https://arxiv.org/abs/2111.00396
https://arxiv.org/abs/2006.16236
https://arxiv.org/abs/2009.14794
https://arxiv.org/abs/2004.05150 ｜ BigBird：https://arxiv.org/abs/2007.14062
https://arxiv.org/abs/2310.06825
https://arxiv.org/abs/2403.19887
https://arxiv.org/abs/2307.03172
https://arxiv.org/abs/2404.06654
ting.is-a.dev. LLM 原理 专栏第 08 篇.

该篇第 2 章的位置编码外推与本库 09-位置编码重叠、第 3.1/3.4 节的 FlashAttention 与 Ring Attention 与本库既有笔记重叠，本文只保留「 $O (n^{2})$ 的账 / 稀疏与线性注意力 / SSM / 混合架构 / 评估」五块。SSM 部分的代际与数字按 S4、Mamba、Mamba-2、Jamba 原论文核对；检索到的二手材料还提到 2026 年 3 月的 Mamba-3，未见于一手论文库，本文不采用

贡献者

文件历史

最后编辑于 5 分钟前查看完整历史

语法基础

gpu

cuda

train

infer

长上下文技术

$O (n^{2})$ 到底贵在哪

结构侧：三条降低复杂度的路线

稀疏注意力

线性注意力

状态空间模型与 Mamba

混合架构：生产环境的现实选择

系统侧：跨设备切分

评估：长上下文怎么算「行」

Needle in a Haystack

其他基准

三条路线怎么选

相关

参考

贡献者

文件历史

长上下文技术 ​

O(n2) 到底贵在哪 ​

结构侧：三条降低复杂度的路线 ​

稀疏注意力 ​

线性注意力 ​

状态空间模型与 Mamba ​

混合架构：生产环境的现实选择 ​

系统侧：跨设备切分 ​

评估：长上下文怎么算「行」 ​

Needle in a Haystack ​

其他基准 ​

三条路线怎么选 ​

相关 ​

参考 ​

贡献者 ​

文件历史 ​

长上下文技术

$O (n^{2})$ 到底贵在哪

结构侧：三条降低复杂度的路线

稀疏注意力

线性注意力

状态空间模型与 Mamba

混合架构：生产环境的现实选择

系统侧：跨设备切分

评估：长上下文怎么算「行」

Needle in a Haystack

其他基准

三条路线怎么选

相关

参考

贡献者

文件历史