✨ 复制成功!

以文会友,打造好学人设!

📋 已复制到剪贴板:

Self-Attention 的计算复杂度是多少?为什么长序列下会成为瓶颈?有哪些降低复杂度的方法?

创建时间:2026-09-15 更新时间:2026-09-16 阅读次数:1493 次 分类:模型架构与Transformer

1、标准 Self-Attention 的计算复杂度

先回顾 Self-Attention 的公式:

Attention(Q, K, V) = softmax(QK^T / √d_k) V

其中:

  • 输入序列长度为 n
  • 每个 token 的维度为 d
  • Q、K、V 的形状都是 (n, d)

逐步算复杂度:

第一步:QK^T

  • Q 是 (n, d),K^T 是 (d, n)。备注:K^T 表示 K 的转置。
  • 矩阵乘法结果 (n, n)
  • 计算量:O(n² · d)

第二步:softmax

  • 对 (n, n) 矩阵做 softmax
  • 计算量:O(n²)

第三步:softmax(...) V

  • (n, n) × (n, d) → (n, d)
  • 计算量:O(n² · d)

总计算复杂度:O(n² · d)

空间复杂度:

  • 注意力矩阵 (n, n),即 O(n²)
  • 这是显存瓶颈的主要来源

结论:

  • 时间复杂度 O(n² · d)
  • 空间复杂度 O(n²)(不考虑多头时)
  • 多头情况下,每个头维度是 d/h,总计算量仍是 O(n² · d)

2、为什么长序列下会成为瓶颈

(1)平方级增长

序列长度从 n 增加到 2n:

  • 计算量变成 4 倍
  • 显存占用变成 4 倍

序列长度到 8K、16K、32K 时:

  • n² 增长极快
  • 计算和显存都爆炸

举例:

  • n = 512,注意力矩阵约 26 万元素
  • n = 8192,注意力矩阵约 6700 万元素
  • n = 32768,注意力矩阵约 10.7 亿元素

单层、单头就这么大,多层多头更夸张。

(2)显存瓶颈

注意力矩阵 (n, n) 需要存储:

  • 前向传播要存
  • 反向传播还要存中间激活
  • 即使 FlashAttention 不显式存,计算过程中也有大量中间量

(3)计算瓶颈

QK^T 和 softmax·V 都是 O(n²·d):

  • GPU 算力有限
  • 长序列下 attention 成为主要耗时
  • KV Cache 在推理时也随 n 线性增长,长上下文显存压力大

(4)和 FFN 的对比

FFN 复杂度是 O(n · d²),和 n 线性相关。

  • 短序列:FFN 占主导
  • 长序列:Attention 的 O(n²·d) 反超,成为瓶颈

所以长上下文场景,优化重点在 Attention。

3、降低复杂度的主要方法

(1)稀疏注意力(Sparse Attention)

思路:不让每个 token 和所有 token 交互,只和部分 token 交互。

  • Local / Sliding Window Attention:每个 token 只看邻近窗口,复杂度 O(n · w)。
  • Strided Attention:隔几个 token 看一个,类似空洞卷积。
  • Block Sparse:把序列分块,只计算部分块对。
  • Longformer / BigBird:局部 + 全局 + 随机组合,近似全注意力。

优点:复杂度降到线性或 O(n√n)。

缺点:可能丢失长距离依赖,实现复杂。

(2)低秩 / 线性注意力(Linear Attention)

思路:用核函数或低秩近似,避免显式构造 (n, n) 矩阵。

  • Linformer:把 K、V 投影到低维,复杂度 O(n · d · k)。
  • Performer:用随机特征近似 softmax,复杂度 O(n · d²)。
  • Linear Transformer:把 softmax 换成可分解核,先算 K^T V,再乘 Q,复杂度 O(n · d²)。

优点:线性复杂度。

缺点:近似有损,效果可能下降。

(3)分块 / IO 优化(FlashAttention)

思路:不改变数学结果,只优化显存访问和计算顺序。

  • FlashAttention:分块计算,利用 SRAM,减少 HBM 读写。
  • FlashAttention-2/3:进一步优化并行和流水线。
  • PagedAttention:优化 KV Cache 显存管理,vLLM 用它。

优点:精确计算,速度快、显存省。

缺点:不改变 O(n²) 复杂度,只是常数优化;但工程收益巨大。

(4)KV Cache 优化

推理阶段,KV Cache 随序列线性增长。

  • MQA(Multi-Query Attention):所有头共享 K、V,大幅减少 KV Cache。
  • GQA(Grouped-Query Attention):分组共享,折中方案,LLaMA 2/3 用。
  • 量化 KV Cache:INT8/INT4 压缩。

优点:推理显存和带宽大幅下降。

缺点:可能损失效果,需要训练适配。

(5)状态空间模型(SSM)

思路:用递归或卷积结构替代 Attention,天然线性复杂度。

  • Mamba:选择性状态空间,复杂度 O(n)。
  • RWKV:线性注意力 + RNN 思想。
  • RetNet:保留并行训练,推理用递归。

优点:长序列线性复杂度,推理快。

缺点:表达能力和生态还在验证中。

(6)其他

  • 滑动窗口 + 全局 token:如 Longformer。
  • 检索增强:长文档先检索再进 Attention。
  • 压缩记忆:把历史压缩成少量 token。
  • 混合架构:Attention + SSM 混合,如 Jamba。

4、方法对比总结

方法复杂度是否精确代表
标准 AttentionO(n²·d)精确Transformer
稀疏 AttentionO(n·w) 或 O(n√n)近似Longformer、BigBird
线性 AttentionO(n·d²)近似Performer、Linformer
FlashAttentionO(n²·d)精确FlashAttention
MQA/GQAO(n²·d)精确LLaMA
SSMO(n)不同范式Mamba、RWKV

一句话:降低复杂度有两条路——改数学(稀疏、线性、SSM) 和 改工程(FlashAttention、KV Cache 优化)。前者降复杂度,后者降常数和显存。

5、这道面试题的质量点评

优点

  • 经典核心题:Self-Attention 复杂度是 Transformer 最基础也最重要的知识点。
  • 有层次:从复杂度计算,到瓶颈分析,再到优化方法,能考察深度和广度。
  • 区分度好:

    • 初级:能说出 O(n²·d)
    • 中级:能解释为什么长序列瓶颈、知道 FlashAttention
    • 高级:能对比稀疏、线性、SSM、MQA/GQA,说清取舍
  • 贴近前沿:长上下文是当前热点,优化方法更新快,能看出候选人是否跟进。

缺点

  • 偏八股:复杂度计算是标准答案,容易背。
  • 优化方法太开放:可以答很多,面试官需要控制范围,否则容易泛泛而谈。
  • 缺少场景:没问“你实际用过哪种?效果如何?”,容易停留在理论。
  • 容易变成名词堆砌:候选人可能列一堆方法但说不清原理和取舍。

总体评价

质量:上。 作为 Transformer 核心题,它有很好的区分度,尤其适合中高级岗位。但单靠这一题不够,最好追问:

  • “FlashAttention 为什么快?它改变了复杂度吗?”
  • “MQA 和 GQA 区别是什么?为什么 LLaMA 用 GQA?”
  • “线性 Attention 为什么效果可能下降?”
  • “Mamba 和 Transformer 的本质区别是什么?”
  • “你实际项目中长序列怎么处理的?”

6、怎么回答才是最好的

回答原则

  • 先给复杂度结论:O(n²·d),空间 O(n²)。
  • 解释瓶颈:平方增长、显存、计算、和 FFN 对比。
  • 分类说优化:改数学 vs 改工程,不要只列名词。
  • 说清取舍:每种方法的代价。
  • 结合实践:说“我们长序列用 FlashAttention + GQA”。

推荐回答结构

“标准 Self-Attention 的时间复杂度是 O(n²·d),空间复杂度 O(n²),主要来自 QK^T 和 softmax·V 两个矩阵乘法。”

“长序列下瓶颈主要是平方级增长:序列翻倍,计算和显存都翻 4 倍。8K 以上注意力矩阵就非常大。短序列时 FFN 占主导,长序列时 Attention 反超。”

“优化方法我分两类。一类改数学:稀疏注意力、线性注意力、SSM,把复杂度降到线性或 O(n√n),但都是近似或有范式差异。另一类改工程:FlashAttention 分块计算,不改变复杂度但大幅省显存、提速度;MQA/GQA 优化 KV Cache,推理显存降很多。”

“我们实际用 FlashAttention-2 + GQA,长上下文推理显存和延迟都明显改善。如果是超长序列,会考虑稀疏或检索增强。”

“需要注意的是,FlashAttention 不改变 O(n²),只是常数优化;线性 Attention 虽然快,但效果可能下降,要评估。”

加分点

  • 提到 FlashAttention 不改变复杂度、MQA/GQA、SSM、Longformer、Performer。
  • 提到 KV Cache 随序列线性增长。
  • 提到 Attention 和 FFN 的复杂度对比。
  • 结合真实长序列处理经验。

减分点

  • 只答 O(n²),不说 d。
  • 把 FlashAttention 说成降低复杂度。
  • 只列方法,不说取舍。
  • 不知道 MQA/GQA 或 SSM。

📌 面试天下网:一款服务于上班族的口袋书,让大家在地铁里/公交上可以学习大模型技术!
📌 网站公告:【悬赏 200 元/篇】寻找“大模型面试战场”的一手回忆录
📌 网站公告:【大模型/Agent面试陪练1V1指导】正式启动......

评论专区

评论加载中...