Efficient AI Lecture 15: Long-Context LLM

长上下文并不仅仅是一个更大的提示窗口。

关键问题是：哪些过去的 token 仍值得被记忆？

RoPE 插值通过将更长的序列压缩到模型预训练期间见过的角度范围内来拉伸位置编码。

大海捞针测试展示了一种常见的失败模式：模型通常能从开头和结尾很好地检索信息，但会遗漏埋藏在中间的事实。

StreamingLLM揭示了一个重要的系统洞察：早期 token 通常充当注意力汇聚点。即使它们在语义上不重要，它们也会吸收多余的 softmax 注意力。

与其保留整个 KV 缓存，StreamingLLM 选择保留：

• 少数初始的汇聚 token
• 一个滚动窗口内的近期 token
• 并驱逐中间的 token

这使得在有限内存下实现稳定的流式推理成为可能。

我的笔记： https://ickma2311.github.io/ML/HW-SW-codesign/efficient-ai-lecture-15-long-context-llm.html…

长上下文 LLM – ∇ ickma.dev

来源：https://ickma2311.github.io/ML/HW-SW-codesign/efficient-ai-lecture-15-long-context-llm.html 长上下文 LLM 不仅仅是“更大的提示窗口”。

它们需要在多个层面进行改进：

• 位置编码必须在预训练范围之外也能正常工作
• 注意力机制必须避免内存和计算量的二次增长
• 评估必须测试模型是否真的使用了长上下文的中间部分
• 推理系统必须决定哪些 KV 缓存条目值得占用内存
• 替代序列模型必须压缩历史信息而不丢失重要细节

核心矛盾很简单：长上下文只有在模型既能存储又能检索正确信息时才有用。

上下文扩展

RoPE: 旋转位置编码

旋转位置编码（RoPE）通过旋转查询和键向量来注入位置信息。

嵌入维度\(d\)被分成\(d/2\)个二维对。每个对被视为一个二维坐标，每个坐标对根据 token 位置进行旋转。

基本频率为：

\[ \Theta = \left\{ \theta_i = 10000^{-2(i-1)/d}, \quad i\in\{1,\dots,d/2\} \right\}。 \]

从复数的角度看待每个二维坐标对，RoPE 可以写成：

\[ \operatorname{RoPE}(x,m) = x e^{im\theta}。 \]

一个有用的性质是，位置为\(m\)的查询与位置为\(n\)的键之间的点积取决于相对距离\(m-n\)：

\[ \left\langle \operatorname{RoPE}(q_j,m), \operatorname{RoPE}(k_j,n) \right\rangle = q_j k_j e^{i(m-n)\theta}。 \]

等价地，位置交互是\(m-n\)的函数，而不仅仅是两个绝对位置分开的函数。

这就是 RoPE 对自回归 Transformer 有用的原因：注意力分数可以编码相对距离，同时仍然通过变换后的查询和键来实现。

通过位置插值扩展上下文

挑战在于 LLM 是在固定的上下文长度下训练的。例如：

• 原始 LLaMA：大约 2k 上下文
• Llama 2：大约 4k 上下文
• GPT-4 系列模型：更长的原生上下文窗口

如果一个在位置\([0,L)\)上训练的模型被迫评估远超出该范围的位置，它可能会因外推而失败。位置模式超出了预训练时见过的分布。

RoPE 使得一种简单的上下文扩展策略成为可能：通过缩放频率来插值位置。

假设一个模型训练时的上下文长度为\(L=2048\)，我们想要评估两倍长的上下文。如果不进行插值：

\[ m\in[0,2L), \qquad \theta_i’=\theta_i。 \]

现在模型看到了训练范围之外的角度位置。

通过位置插值：

\[ m\in[0,2L), \qquad \theta_i’=\frac{\theta_i}{2}。 \]

将频率减半可将两倍数量的位置压缩到模型已经见过的角度范围内。更一般地，将上下文扩展因子\(s\)可使用：

\[ \theta_i’=\frac{\theta_i}{s}。 \]

其理念是优先在已知位置范围内进行插值，而不是外推到范围之外。

LongLoRA

标准注意力机制随序列长度呈二次方缩放：

\[ O(N^2)。 \]

随着上下文长度的增长，注意力成为主要的训练瓶颈。

LongLoRA 通过两个思路解决这个问题：

• 在长上下文适应过程中使用稀疏注意力
• 在普通 LoRA 之外微调少量低代价参数

移位稀疏注意力

LongLoRA 使用移位稀疏注意力，记作\(S^2\)-Attention。

注意力头被分成两组：

• 一组在不移位的情况下关注局部块内的内容
• 另一组将块边界移动半个块

例如，如果块大小为 4：

• 未移位的头将 token 分组为\([0,1,2,3]\),\([4,5,6,7]\), …
• 移位的头将 token 分组为\([2,3,4,5]\),\([6,7,8,9]\), …

每个头仍然执行廉价的局部注意力，但移位的分组在相邻块之间创建了重叠。跨层时，信息可以通过序列传播，而每一层避免了完整的密集注意力。

增强的 LoRA

LongLoRA 还发现，仅对注意力权重进行普通 LoRA 通常不足以实现激进的上下文扩展。

它为训练开放了一些额外的组件：

• 自注意力前的输入归一化
• 前馈网络前的注意力后归一化
• 输入嵌入

这些参数与全微调相比是廉价的。归一化层只占总参数的一小部分，嵌入层仍然远小于所有 Transformer 块。但调整它们有助于缩小 LoRA 式微调与全微调之间的差距。

评估长上下文 LLM

大海捞针

大海捞针测试询问模型是否能够检索到隐藏在长上下文中的特定事实。

这揭示了一种反复出现的失败模式：当相关事实出现在提示的开头或结尾附近时，模型通常表现良好，但当事实埋藏在中间时，表现会差得多。

由此产生的曲线通常是 U 形的：

• 开头附近准确率高
• 中间准确率低
• 结尾附近准确率高

实际教训是，接受长输入并不意味着能稳健地使用整个输入。埋藏在长文档中间的关键事实可能会被遗漏。

对于提示设计，重要的指令和事实放在上下文的开头或结尾附近更安全。

LongBench

合成检索测试很有用，但还不够。

LongBench 在更广泛的真实任务上评估长上下文模型：

• 单文档问答
• 多文档问答
• 摘要
• 少样本学习
• 编程任务
• 合成任务

它包含英文和中文输入，上下文长度可达 13000 个 token 以上。根据任务不同，指标包括 F1 和 ROUGE。

两个教训值得注意：

• 位置编码优化可以直接改善长上下文理解
• 检索或压缩的变通方法有帮助，但原生长上下文架构通常表现更好

高效注意力机制

KV 缓存优化

在自回归生成过程中，每个新 token 都需要访问之前的键和值。KV 缓存避免了重新计算它们，但其内存随序列长度和层数线性增长。

对于长上下文，KV 缓存内存成为部署瓶颈。

主要问题变成了：

哪些过去的 token 必须保持可用，哪些可以被压缩、驱逐或近似？

以下方法以不同的方式回答了这个问题。

StreamingLLM

StreamingLLM 从一个经验观察开始：早期 token 会获得异常大的注意力质量，即使它们在语义上并不重要。

这被称为注意力汇聚现象。

注意力汇聚

在自回归 Transformer 中，注意力分布的每一行在 softmax 之后必须总和为 1。

当一个 token 不需要从更早的 token 获取太多信息时，模型仍然需要将剩余的概率质量放在某处。由于第一个 token 对所有后来的 token 可见，它成为了剩余注意力的便捷放置点。

结果是，在第一个 token 处出现一条垂直的注意力条带，跨越许多后续位置和层。

关键点是，这些汇聚 token 不一定是重要的。它们的作用是结构性的：它们稳定注意力分布。

利用注意力汇聚实现无限流式推理

StreamingLLM 使用一种混合缓存：

• 保留前几个注意力汇聚 token 固定不变
• 保留一个滚动窗口内的近期 token
• 驱逐中间的 token

例如，使用 4 个汇聚 token 和 4 个近期 token 预算：

步骤汇聚 token被驱逐的 token滚动缓存生成 token 7\([0,1,2,3]\)\([4]\)\([5,6,7]\)生成 token 8\([0,1,2,3]\)\([4,5]\)\([6,7,8]\)生成 token 9\([0,1,2,3]\)\([4,5,6]\)\([7,8,9]\)这提供了有界的内存，同时保留了模型期望的注意力汇聚行为。

复杂度变为：

\[ O(TL), \]

其中\(T\)是生成的 token 总数，\(L\)是固定的活动缓存长度。

位置编码分配

对于 StreamingLLM，位置应遵循缓存顺序，而不是原始的绝对文本位置。

如果中间 token 被驱逐，但剩余的 token 保留其旧的绝对位置，RoPE 会看到人为的间隙。例如，模型可能从位置 3 跳到位置 10000，即使缓存现在包含相邻的活动 token。

更安全的规则是：

根据紧凑的缓存布局分配位置。

这保留了模型在预训练期间看到的相对距离模式。

使用专用注意力汇聚进行预训练

普通模型通常需要多个第一个 token 来稳定流式推理。

一个自然的问题是，是否可以用一个专用的汇聚 token 来预训练模型。

方法很简单：在每个预训练样本的开头添加一个额外的可学习 token。这个 token 从一开始就被训练为注意力汇聚。

报告的结果是，这可以将永久汇聚开销减少到一个 token，而不损害预训练收敛。

模型0 + 10241 + 10232 + 10224 + 1020Vanilla27.8718.4918.0518.05Zero Sink2921419.9018.2718.01Learnable Sink123518.0118.0118.02可学习汇聚模型仅用一个汇聚 token 就达到了稳定的困惑度。

DuoAttention: 检索头与流式头

DuoAttention 始于另一个观察：

并非所有注意力头都需要全局上下文。

有些头充当检索头。它们需要访问长距离上下文，并且对 KV 缓存压缩敏感。

其他头充当流式头。它们主要使用局部近期上下文和注意力汇聚，因此可以使用较小的固定缓存运行。

预填充瓶颈

在预填充阶段，模型处理整个输入提示并构建初始 KV 缓存。

对于长上下文，预填充期间的全注意力代价高昂，因为每个 token 都关注整个历史：

\[ O(N^2)。 \]

DuoAttention 通过仅对需要的头保留全注意力来减少这种开销。

头类型

检索头：

• 捕获长距离依赖
• 需要完整的 KV 缓存历史
• 如果被积极压缩会损失精度

流式头：

• 关注近期局部上下文
• 使用注意力汇聚加上滚动缓存
• 能容忍固定大小的稀疏 KV 缓存

基于门的头选择

DuoAttention 为每个头分配一个可训练的门控值\(\alpha_{i,j}\)，其中\(i\)是层，\(j\)是头。

混合注意力输出为：

\[ \operatorname{attn}{i,j} = \alpha{i,j}\operatorname{full\attn} + (1-\alpha{i,j})\operatorname{streaming\_attn}。 \]

全注意力为：

\[ \operatorname{full\attn} = \operatorname{softmax}(QK^T\odot M{\mathrm{causal}})V。 \]

流式注意力为：

\[ \operatorname{streaming\attn} = \operatorname{softmax}(QK^T\odot M{\mathrm{streaming}})V。 \]

\(\alpha\)高于阈值的头被视为检索头。低于阈值的头使用轻量级流式缓存。

训练目标

训练目标在保持全注意力模型和鼓励稀疏头选择之间取得平衡：

\[ \mathcal{L} = \mathcal{L}{\mathrm{distill}} + \lambda \mathcal{L}{\mathrm{reg}}。 \]

蒸馏损失匹配全注意力模型与混合注意力模型之间的隐藏状态：

\[ \mathcal{L}{\mathrm{distill}} = \frac{1}{N} \sum{i=1}^{N} \sum_{j=T-l+1}^{T} \left( H_{\mathrm{full}}^{(i)}[j] - H_{\mathrm{mixed}}^{(i)}[j] \right)^2。 \]

正则化损失推动更多头进入流式模式：

\[ \mathcal{L}{\mathrm{reg}} = \sum{i=1}^{L} \sum_{j=1}^{H} |\alpha_{i,j}|。 \]

在一个报告的约 1.84M token 的长上下文设置中：

指标全注意力DuoAttention影响GPU 内存78.4 GB49.3 GB大约节省 37% 内存预填充时间4570 s2348 s大约快 2 倍更广泛的要点是，并非每个头都同样需要全上下文内存。

SpAtten

SpAtten 通过从 KV 缓存中修剪不重要的 token 来降低注意力成本。

它通过累积注意力来估计 token 的重要性。对于每个历史 token，它对其随时间获得的注意力概率求和。累积重要性低的 token 被视为冗余并可被移除。

这是查询无关的：一旦一个 token 被判定为不重要，即使后续查询可能需要它，它也可能被丢弃。

Quest: 查询感知稀疏性

Quest 使稀疏注意力具有查询感知能力。

它不是仅仅根据年龄或历史重要性永久驱逐旧的 token，而是根据当前查询决定检索什么。

缓存被分成页面。每个页面存储轻量级的键摘要，例如键的最小值和最大值。

然后 Quest 运行一个两阶段过程。

首先，它估计哪些页面可能对当前查询重要。这是一个对页面摘要进行廉价评分的步骤。

其次，它只加载前\(K\)个选中的页面，并对这些页面计算全注意力。

这提供了稀疏注意力的内存优势，同时保留了在旧 token 变得相关时恢复它们的途径。

超越 Transformer

状态空间模型

状态空间模型（SSM）通过隐藏状态处理序列。

在时间\(t\)：

• \(x_t\)是输入
• \(h_t\)是隐藏状态
• \(y_t\)是输出

在推理过程中，循环公式为：

\[ \mathbf{h}t = \overline{\mathbf{A}}\mathbf{h}{t-1} + \overline{\mathbf{B}}\mathbf{x}_t, \]

\[ \mathbf{y}_t = \mathbf{C}\mathbf{h}_t。 \]

这提供了每个 token 的常数时间状态更新：

\[ O(1) \]

相对于序列长度。

对于训练，循环可以重写为全局卷积：

\[ \overline{\mathbf{K}} = ( \mathbf{C}\overline{\mathbf{B}}, \mathbf{C}\overline{\mathbf{A}}\overline{\mathbf{B}}, \dots, \mathbf{C}\overline{\mathbf{A}}^k\overline{\mathbf{B}}, \dots ), \]

\[ \mathbf{y} = \mathbf{x} * \overline{\mathbf{K}}。 \]

这种卷积形式可以在训练期间并行化，而循环形式在推理期间高效。

Mamba

经典 SSM 有一个局限性：矩阵\(\mathbf{A}\),\(\mathbf{B}\), 和\(\mathbf{C}\)是静态的。

它们不根据输入内容改变。这使得有选择地记住重要 token 而忽略填充 token 变得困难。

Mamba 通过使关键参