注意力复杂度分析：O(n²)瓶颈与突破路径

可视化演示

展示自注意力O(n²)计算量随序列长度爆炸增长、显存瓶颈分析与稀疏化线性化低秩近似三种突破路径，切换模式掌握复杂度优化核心。

复杂度本质

回到我们的图书馆。管理员最初设计的“全馆名片交换”确实优雅：每本书都能精准知道其他任何一本的位置和内容概要。但管理员很快发现了规律——藏书量 n 从 100 涨到 200 时，名片交换总量不是翻倍，而是从 4950 次涨到 19900 次，整整翻了 4 倍！这就是数学家口中的“平方级增长”：总量约等于 n²/2。在自注意力机制里，这个“交换名片”的过程发生在每个词元（Token）之间，词元总数 n 就是藏书量。

自注意力之所以强大，是因为它让任意两个位置直接建立依赖，捕获长距离关系。但也恰恰是这份“事无巨细的全连接”，让它的时间和空间开销都随 n 的平方增长。具体来说，自注意力的核心计算可以写作：

时间上，矩阵乘法和 softmax 都绕不开这 n×n 的表格；空间上，这张表格本身就得存在内存里。因此我们说，自注意力的时间复杂度和空间复杂度同量级，都是 O(n²)。当 n 从 512 涨到 4096，这张表格的规模直接膨胀 64 倍——再厚的钱包也经不起这样花。理解了“全连接”这个根，后面的所有优化都是在回答同一个问题：能不能不跟所有人交换名片？

时间复杂度推导

图书馆管理员的日常工作可以拆成三步：先派发问卷（每本书向全馆询问“你和我有多相关”），再筛选重点书目（把相关度不高的礼貌性忽略），最后把重点书的内容摘抄回来综合。自注意力的一次计算，也恰好是这三个流程的矩阵流水线。

第一步是 QK^T 矩阵乘法。Q 的每一行是一个词元的查询向量，K 的每一行是键向量。一个 n×d 的矩阵乘以一个 d×n 的矩阵，结果是一个 n×n 的得分矩阵。这个乘法的操作数是 2n²d 次浮点运算。当序列长度 n 增长时，这一项随 n 的平方膨胀。第二步是 Softmax 归一化，它对得分矩阵的每一行 $n$ 个元素计算指数和归一，$n$ 行合计操作数是 O(n²)。第三步是 注意力加权求和，把归一化后的 n×n 权重矩阵与 V 矩阵（n×d）相乘，同样是 O(n²·d)。三项相加：

这解释了为什么 GPT 系列早期只敢把上下文窗口设在 2048 或 4096——一旦序列翻倍，时间成本直接翻四倍。试想：图书馆从 1000 本书涨到 2000 本，管理员的名片工作量不是翻倍而是四倍，他当然要限制“馆藏规模”。这也是为什么你在实际训练大模型时，会明显感觉到：序列长度加一点，训练时间就“坐火箭”。理解了这三步，你就拿到了亲手估算任意 Transformer 前向耗时的基础：先数 n，再乘 d，其他都是小头。

空间复杂度分析

图书馆里除了交换名片要花时间，还有第二个噩梦：登记簿。管理员需要一张巨大的表格，记录每两本书之间的亲疏关系得分——行列都是书，一页纸根本画不下。更麻烦的是，softmax 需要先算出整张表的完整得分，才能逐行归一化，所以这张 n×n 的“得分矩阵”必须完整驻留在内存里，不能边算边扔。这就是自注意力的空间复杂度 O(n²) 的由来。

我们来算一笔账。假设序列长度 n = 4096，模型用 FP16（每个数占 2 字节）存储得分矩阵，那么仅这一张表就占用 $4096^2 \times 2 = 32$ MB 显存。听起来不多？但注意，Transformer 有十几层、多个注意力头，每层都要维护自己的得分矩阵，还要算上反向传播时保存的中间梯度，实际开销会放大几十倍。当 n 涨到 32768，一张表就变成 2 GB，直接撑爆单卡显存：

这也是为什么部署长文本模型时，工程师宁可做量化、做算子融合，也要省出这口气——因为“登记簿”比“名片”本身更占地方。值得一提的是，推理阶段的 KV 缓存（缓存历史 K 和 V 矩阵）开销是 O(n·d) 而不是 O(n²)，这是因为推理时只需把新词元与缓存中的 K、V 做注意力，不需要重新计算历史词元两两之间的得分。但训练阶段没有这个便利，整张 n×n 表格躲无可躲。所以严格说，训练是 O(n²) 的重灾区，推理的瓶颈则更多落在 KV 缓存的线性增长上。

长序列挑战

当图书馆的藏书量突破某个临界点，管理员发现整个系统开始“卡壳”：名片交换的耗时变得不可接受，登记簿占满了所有柜子，读者来借书要等上半天。这就是长序列场景下标准自注意力遇到的真实困境。业界把这个问题概括为“上下文窗口受限”——不是模型不想看更长的历史，而是硬件内存和计算预算不允许。

具体到实践层面，三个问题最为突出。第一，上下文窗口受限：GPT-2 的窗口是 1024，GPT-3 是 2048，即便 GPT-4 达到了 128K，也远小于人类阅读一本书的体量——限制它的不是模型架构，而是 O(n²) 的硬件账单。第二，显存占用爆炸：训练时得分矩阵随 n² 增长，序列从 32K 翻到 64K，显存需求翻四倍，往往直接触发 OOM（显存溢出），让训练任务崩溃。第三，推理吞吐下降：虽然推理阶段 KV 缓存是 O(n·d)，但每生成一个新词元，都要和全部历史做注意力计算，随着生成变长，单次解码延迟线性上升，系统吞吐骤降。

用一个公式来概括序列长度对资源账单的影响。设序列长度增加 k 倍，则：

正因如此，处理一本书级别的长文本（约 10 万词元）时，标准 Transformer 的显存需求会达到 TB 级别，远超任何单卡容量。这也是为什么长文档问答、多轮长对话、代码仓库理解这类真实应用，必须依赖后面几节要讲的优化手段。图书馆管理员痛定思痛：与其让每本书都认识所有书，不如设计更聪明的信息流通方案。

稀疏化方向

管理员想出的第一套改革方案很务实：并不是所有书都需要互相认识。一个读者在查某一章的细节时，最相关的往往是前后相邻的几页；而真正需要“全馆广播”的，只有少数几本核心参考书。顺着这个直觉，研究者提出了稀疏注意力——把全连接的名片交换网，裁剪成一张稀疏但够用的“社交网络”。

最常见的稀疏化模式是局部窗口注意力（Window Attention）：规定每本书只与前后 w 本书交换名片，其余一律不联系。这样一来，每本书的名片交换次数从 n-1 次降到 2w 次，总交换量从 O(n²) 降到 O(n·w)。如果 w 是固定常数（如 256），复杂度就变成了线性的 O(n)。代价呢？远处的书失去了直接联系——但这恰恰是局部窗口想避免的：远处的信息可以通过一层层窗口的接力传递，就像消息通过朋友的朋友传到远方。实际应用中，为了弥补长距离信息的流失，还会放置少量“全局锚点”：

稀疏化的哲学是“用结构换速度”——主动放弃一部分理论上可能的连接，换取计算和内存的大幅缩减。它特别适合文本、图像这类局部性很强的数据：相邻像素、相邻词元之间的相关性本来就最强。但在某些需要深层长程推理的任务上，稀疏连接可能让信息流动变得迂回，这是工程师选型时必须权衡的。图书馆管理员点点头：与其让每本书都握着全馆名片，不如给每本书发一张“附近街区地图”外加几本核心手册——够用，且便宜得多。

线性注意力

图书馆管理员看着星型结构的“索引卡方案”陷入了沉思：如果所有书都把信息汇总到一张总索引卡上，新书来了只需和索引卡交换一次名片，总数就是 O(n)，而不是 O(n²)。这就是线性注意力（Linear Attention）的核心直觉——把“两两交互”的矩阵乘法，通过数学技巧改写成“先聚合、再分发”的形式。

技巧的关键在于核函数特征映射。标准注意力里的 softmax 分数本质上是一个核函数，它把 Q 和 K 映射到高维空间并求内积。线性注意力用一个新的映射函数 φ(x) 来近似这个核，使得注意力权重可以拆解成 φ(Q) 与 φ(K) 的乘积。于是矩阵乘法的结合律给了我们一个惊喜：原来必须“Q 先乘 K 再乘 V”的顺序，可以调整为“K 先乘 V，再让 Q 去乘结果”：

更妙的是，这个“先聚合后查询”的形式天然契合递推：由于每一时刻的输出只依赖一个固定大小的状态矩阵，线性注意力可以被写成 RNN 式的递推形式——每读入一个新词元，更新一次状态，状态大小恒定，从原理上支持无限长的上下文。当然，天下没有免费的午餐：核近似带来信息损失，实验表明线性注意力在需要精确局部匹配的任务上（如复制、计数）弱于标准注意力。但图书馆管理员的算盘打得响：如果我用 95% 的精度换来“无限藏书量”，对很多真实场景来说，这笔交易太划算了。

低秩与未来

管理员的最后一张底牌来自一个更深刻的观察：那张 n×n 的得分矩阵，真的是“满编”的吗？统计发现，经过 softmax 归一化后，注意力得分往往集中在一小部分关键词元上，大部分分数趋近于零——矩阵的有效信息其实集中在一个低维子空间里。就像图书馆里真正被反复借阅的书只占一小部分，厚厚的书单里大部分条目无人问津。这就是低秩近似的出发点。

数学上，若 n×n 的得分矩阵 P 的秩远小于 n，则存在两个瘦长矩阵的乘积来逼近它：P ≈ A·B，其中 A 是 n×r，B 是 r×n，r 是远小于 n 的“有效秩”（例如 64 或 128）。这样一来，计算量就从 O(n²·d) 降为构建 A 和 B 的代价加上一次 O(n·r·d) 的乘法：

把视野放回全景，今天的主流高效注意力方案几乎都是上述思路的组合拳：稀疏化（Sparse）、线性化（Linear）、低秩化（Low-rank），以及它们的混合体。比如 Longformer 走稀疏窗口路线，Linear Transformer 走核函数路线，Performer 用随机特征映射做低秩近似，最新的混合架构（如 Jamba、Mamba-2）则在同一模型里同时部署注意力层和线性层，各取所长。图书馆管理员的最终方案也逐渐成型：要快速查身边的书就用窗口注意力，要永久记住核心馆藏就用线性状态，要偶尔做深度检索就留几个全局锚点——没有万能药，只有按场景配比的工程智慧。

AI 应用场景

Python 演示

import numpy as np def elu_plus_one(x, alpha=1.0):    return np.where(x > 0, x, alpha * (np.exp(x) - 1)) + 1 def standard_attention(Q, K, V):    scores = Q @ K.T / np.sqrt(K.shape[-1])    weights = np.exp(scores - scores.max(axis=-1, keepdims=True))    weights /= weights.sum(axis=-1, keepdims=True)    return weights @ V def linear_attention(Q, K, V):    phi_q = elu_plus_one(Q)    phi_k = elu_plus_one(K)    kv = phi_k.T @ V    z = phi_k.sum(axis=0)    out = phi_q @ kv    norm = phi_q @ z    return out / norm[:, None] def main():    n, d = 256, 32    rng = np.random.default_rng(42)    Q = rng.normal(size=(n, d))    K = rng.normal(size=(n, d))    V = rng.normal(size=(n, d))     out_std = standard_attention(Q, K, V)    out_lin = linear_attention(Q, K, V)    rel_err = (np.linalg.norm(out_std - out_lin, 'fro')               / np.linalg.norm(out_std, 'fro'))    mem_std = n * n * 4 / 1024 / 1024    print(f"序列长度 n = {n}, 维度 d = {d}")    print(f"线性注意力相对误差: {rel_err:.4f}")    print(f"标准注意力得分矩阵内存: {mem_std:.2f} MB")    print("线性注意力不构建得分矩阵：内存为 0 MB") if __name__ == "__main__":    main()

1. scores = Q @ K.T / np.sqrt(K.shape[-1])生成 n×n 得分矩阵并缩放，这一步正是 O(n²) 的源头，内存大头从此处来。
2. weights /= weights.sum(axis=-1, keepdims=True)softmax 归一化，每一行的注意力权重之和为 1，模拟“预算有限、按亲疏分配”。
3. phi_q = elu_plus_one(Q)用 ELU+1 作为核函数特征映射，把 Q、K 映射到新空间，让乘法顺序可以交换。
4. kv = phi_k.T @ V先聚合 K 和 V，得到 d×d 小矩阵，复杂度与序列长度 n 彻底解耦，这是线性化的关键一步。
5. norm = phi_q @ z计算归一化常数，保证输出尺度与标准注意力可比，避免数值漂移。

常见误区

面试题