目录
1. 引言:长上下文窗口的挑战
在Transformer架构中,注意力机制的核心是计算Query (Q), Key (K), 和 Value (V) 矩阵。其中,Q和K的点积 QKᵀ 会形成一个注意力分数矩阵,其维度为 N x N,其中 N 是上下文窗口的长度(序列长度)。
- 计算瓶颈: 当上下文窗口
N变得非常大时(例如从 1K 增长到 100万),N x N矩阵的计算量和内存占用会呈二次方增长,成为主要的性能瓶颈。 - 内存溢出: 巨大的注意力矩阵会导致显存(GPU内存)溢出。
- KV Cache: 推理过程中,KV Cache 的大小也随着
N线性增长,多层 Transformer 叠加会使其急剧膨胀,消耗大量内存。
为了解决这些问题,研究者们提出了多种优化方案,本节课主要介绍 Sparse Attention 和 Infini-Attention。
Screenshot-[00:48]
2. Sparse Attention (稀疏注意力)
Sparse Attention (稀疏注意力) 是一种由 Google AI 提出的算法,旨在通过牺牲一定的精度来换取处理超长上下文窗口的能力。其核心思想是,并非所有的 Token 都需要与所有其他 Token 计算注意力,我们只计算其中一部分“重要”的注意力分数。
2.1. 核心思想与动机
- 完整注意力 (Full Attention): 计算
N x N矩阵中的所有元素,计算量巨大。 - 稀疏注意力 (Sparse Attention): 只计算矩阵中的一小部分元素,从而大幅减少计算量。这是一种用精度换取效率的策略。
2.2. 三大核心组件
Sparse Attention 将注意力计算分解为三个部分,并将它们组合起来形成最终的稀疏注意力矩阵。
Screenshot-[01:15]
-
Random Tokens (随机注意力)
- 描述: 随机选择矩阵中的一部分位置进行计算。例如,可以设定一个比例(如15%),只计算这15%的
QKᵀ值。 - 目的: 捕捉一些可能被忽略的、非局部的长距离依赖关系。
- 描述: 随机选择矩阵中的一部分位置进行计算。例如,可以设定一个比例(如15%),只计算这15%的
-
Window Tokens (窗口注意力)
- 描述: 计算注意力矩阵中主对角线附近的元素。
- 目的: 保证每个 Token 都能关注到其邻近的 Token。因为在自然语言中,相邻的词语通常关系最密切。对于长文本,这个“邻近”窗口可以设置得较大(如一两百个词)。 Screenshot-[01:50]
-
Global Tokens (全局注意力)
- 描述: 让某些特殊的“全局”Token 可以关注到所有的其他 Token,同时所有其他 Token 也可以关注到这些“全局”Token。在矩阵上表现为计算某些特定的行和列。
- 目的: 捕捉对整个序列都至关重要的信息,类似于句子中的
[CLS]Token。 Screenshot-[02:52]
通过将这三种注意力计算出的分数合并,就得到了一个稀疏的注意力矩阵。 Screenshot-[03:20]
2.3. Big Bird 模型与可视化
Google 在其论文 Big Bird 中正式提出了这种结合了三种稀疏注意力组件的机制。
- Full Attention: 像一个完全连接的图,每个节点(Token)都与其他所有节点相连。 Screenshot-[05:07]
- Random Attention: 节点之间随机连接。 Screenshot-[05:47]
- Window Attention: 每个节点只与它相邻的节点连接,形成一个环。 Screenshot-[05:56]
- Global Attention: 存在一个或多个中心节点,所有其他节点都与它连接。 Screenshot-[06:14]
- Big Bird (Combined): 结合了以上三种连接方式,既保留了局部信息和全局信息,又通过随机连接捕捉了长距离依赖,同时计算量远小于 Full Attention。 Screenshot-[06:34]
2.4. 优缺点与权衡
- 优点: 大幅减少了计算量和内存占用,使得处理超长上下文成为可能。
- 缺点: 丢失了一部分非关键位置的注意力信息,可能会影响模型在某些任务上的精度。
- 权衡: 这是一种典型的效率与精度之间的权衡。在对精度要求不是极高,但上下文长度非常重要的场景(如长文档摘要)中非常有效。多头注意力机制(Multi-Head Attention)可以在一定程度上弥补单个头的精度损失,因为不同的头可能会捕捉到不同的稀疏关系。
3. Infini-Attention (无限注意力)
Infini-Attention 是 Google 在 Gemini 1.5 中发布的一项新技术,旨在实现百万级别的超长上下文窗口。
3.1. 背景与动机
在处理长文档(如一本小说)时,传统的固定窗口模型存在严重问题。例如,当模型处理到小说的第三章时,由于窗口大小限制,它可能已经完全忘记了第一章出现过的人物或情节,导致理解断裂。
要实现真正的长程记忆,模型需要能够回顾历史中所有的内容。然而,将所有历史信息都放入 KV Cache 会导致内存爆炸。Infini-Attention 正是为了解决这个问题而设计的。
Screenshot-[08:40]
3.2. 核心机制:压缩记忆 (Compressed Memory)
Infini-Attention 的核心思想包括两个部分:
- 压缩记忆 (Compressed Memory): 不再保留所有历史的 KV Cache,而是将过去的 KV Cache “压缩”成一个固定大小的记忆摘要。
- 线性访问 (Linear Access): 模型在处理新信息时,会同时关注当前窗口内的内容和这个被压缩的历史记忆。
Screenshot-[09:50]
工作流程:
- 模型按固定大小的段落(Segment)处理长文本。
- 每处理完一个段落,该段落的 KV Cache 就会被一种特殊方法压缩,并更新到全局的“压缩记忆”中。
- 在处理下一个段落时,模型不仅会计算当前段落内的注意力,还会用它的
Q去查询这个“压缩记忆”,从而将历史信息融入当前的计算中。 Screenshot-[10:07]
3.3. 技术溯源:Linear Attention (线性注意力)
Infini-Attention 的实现严重依赖于一篇由华人学者主导的论文,该论文提出了一种名为 Linear Attention 或 Efficient Attention 的机制。
Screenshot-[14:52]
3.4. 计算方式对比:标准注意力 vs. 线性注意力
Screenshot-[15:10]
-
标准注意力 (Standard Attention):
- 计算顺序:
Attention(Q, K, V) = Softmax((QKᵀ) / √d) V - 瓶颈:
QKᵀ产生一个N x N的大矩阵。
- 计算顺序:
-
线性注意力 (Linear Attention):
- 计算顺序:
Attention(Q, K, V) = Q (KᵀV) - 核心改变: 改变了矩阵乘法的结合顺序。先计算
KᵀV。 - 优势:
K的维度是N x Dk,V的维度是N x Dv。Kᵀ的维度是Dk x N。KᵀV的结果矩阵维度是Dk x Dv。Dk和Dv是头的维度(如 64 或 128),远小于序列长度N(如几十万)。- 这样就避免了生成
N x N的大矩阵,计算复杂度和内存占用从O(N²)降低到了O(N)。
- 计算顺序:
虽然在数学上,(QKᵀ)V 和 Q(KᵀV) 在不考虑 Softmax 的情况下是等价的,但论文指出,在加入 Softmax 后,两者的结果仍然高度近似,因此可以互相替代。
3.5. 概念解读:精确关注 vs. 整体总结
这两种计算顺序在概念上代表了不同的信息处理方式。 Screenshot-[18:51]
-
标准注意力 (
QKᵀ):- 过程: 将两个细长的矩阵(
Q和K)扩展成一个巨大的N x N方阵。 - 意义: 计算 每一对 Token 之间的精确关系。它回答的问题是:“对于当前的这个词,我应该关注上下文中的哪一个词?” 这是一种精细化的、点对点的关注。 Screenshot-[19:26]
- 过程: 将两个细长的矩阵(
-
线性注意力 (
KᵀV):- 过程: 将两个高维的矩阵(
Kᵀ和V)压缩成一个小的Dk x Dv矩阵。 - 意义: 计算 所有 Token 在每个特征维度上的加权和。它不再关注单个词之间的关系,而是对整个上下文进行 总结和概括,形成一个“记忆摘要”或“要点总结”。 Screenshot-[21:44]
- 过程: 将两个高维的矩阵(
Infini-Attention 正是利用了 KᵀV 这种“总结”特性来创建其“压缩记忆”。
3.6. 优势与局限性
- 优势:
- 无限上下文: 理论上,通过不断将新的
KᵀV结果累积到压缩记忆中,可以实现无限长度的上下文处理。 - 可微调实现: 无需从零开始预训练模型。可以在现有模型(如 Gemini 1.0)的基础上,通过全参数微调的方式,让模型学会使用这种压缩记忆,从而快速升级到支持超长上下文(如 Gemini 1.5)。
- 无限上下文: 理论上,通过不断将新的
- 局限性:
- 信息丢失: “总结”必然会丢失细节。模型能够记住过去发生过某件事(如第一章出现过某个人),但可能无法精确回忆起原文的每一个字。
- “大海捞针”测试: 虽然 Google 的发布会展示了模型在100万 Token 中找到特定信息的“大海捞针”能力,但这很可能是针对性优化的结果,在泛化场景下性能仍有待检验。
4. 实现超长上下文窗口的技术总结
截至 2024 年 5 月,业界实现超长上下文窗口主要依赖以下技术的组合: Screenshot-[29:53]
-
预训练与微调策略:
- 预训练: 使用较小的窗口(如 4K)进行,以节省成本。
- 微调: 使用更大的窗口(如 64K+)进行微调,并配合
ALiBi或RoPE的变体等位置编码方法来扩展模型的上下文理解能力。
-
稀疏注意力 (Sparse Attention):
- 通过只计算部分注意力分数来降低计算量。
-
底层计算优化 (Flash Attention, MQA/GQA):
Flash Attention: 从 GPU 底层优化内存读写,提高计算效率。MQA/GQA: 减少 Key 和 Value 头的数量,降低 KV Cache 大小。
-
无限注意力 (Infini-Attention):
- 通过压缩历史信息为摘要,实现理论上无限的上下文长度,是当前实现超长上下文的最重要技术之一。国内的
Kimi可能也采用了类似原理。
- 通过压缩历史信息为摘要,实现理论上无限的上下文长度,是当前实现超长上下文的最重要技术之一。国内的
-
MLP 路由 (MLP Routing):
- 类似于
Mixture of Experts (MoE),通过一个“路由器”将不同的任务或请求分配给不同大小的 MLP(前馈网络)层。简单任务用小网络,复杂任务用大网络,从而优化整体的推理速度和资源消耗。
- 类似于
-
硬件堆砌 (More GPUs):
- 最直接但成本最高的方法:使用更多、更大、更新的 GPU。
-
Temp-LoRA (临时 LoRA):
- 一种新兴技术,在推理过程中为当前上下文动态训练一个临时的
LoRA适配器。这样,模型可以将上下文信息“学习”到临时的权重中,而非仅仅存储在 KV Cache 里。推理结束后,该临时适配器被丢弃,不影响原始模型。
- 一种新兴技术,在推理过程中为当前上下文动态训练一个临时的
AI 总结
该视频深入探讨了两种为解决大语言模型长上下文窗口瓶颈而设计的先进注意力机制:Sparse Attention 和 Infini-Attention。
首先,视频介绍了 Sparse Attention (稀疏注意力),它通过牺牲部分计算精度来换取处理超长序列的能力。其核心思想是,并非所有词元间的注意力都同等重要,因此它只选择性地计算三种关键的注意力:随机注意力 (Random)、窗口注意力 (Window) 和 全局注意力 (Global)。这种方法源于 Google 的 Big Bird 模型,它能显著降低 N x N 注意力矩阵的计算和存储开销。
接着,视频重点讲解了 Infini-Attention (无限注意力),这是 Google Gemini 1.5 实现百万级上下文窗口的核心技术。它通过一个创新的“压缩记忆 (Compressed Memory)”机制,将历史信息的 KV Cache 压缩成一个固定大小的“摘要”,从而避免了 KV Cache 随序列长度无限增长的问题。其技术基础是 Linear Attention (线性注意力),它通过改变矩阵乘法的顺序(从 (QKᵀ)V 变为 Q(KᵀV)),将计算瓶颈从 O(N²) 降至 O(N)。这种方法在概念上将注意力从“词元间的精确关系”转变为对“上下文整体的总结”。
最后,视频总结了当前实现超长上下文窗口的多种主流技术,包括位置编码优化、Flash Attention、MQA/GQA、MLP 路由以及新兴的 Temp-LoRA 等,为理解大模型长文本处理能力的技术全景提供了清晰的路线图。