2026-04-15 · transformer
【Transformer 与注意力机制】10 RNN 的根本局限：为什么需要 Transformer
RNN 三难（长程依赖、梯度稳定、训练并行）的系统分析；attention 如何作为补丁逐步把 RNN 推向极限；Vaswani 2017 抛弃循环的范式革命
如果只让我挑一个 Transformer 时代最重要的工程问题，我会毫不犹豫地选「O(n²) 的注意力复杂度」。它不是某个论文的边角细节，而是过去七年（2017–2024）里几乎所有「Transformer 改进型工作」共同的攻击目标——FlashAttention、Sparse Transformer、Longformer、Performer、Linformer、Mamba、RWKV、RetNet、Hyena，名字层出不穷，但全都在回答同一个问题：能不能让 attention 不再是 n²？
这个问题之所以难，不是因为数学上找不到 O(n) 算法（早就有了），而是因为：在大规模语言模型这个具体场景下，所有「降复杂度」的方案在效果上都打不过 vanilla full attention。这就出现了一个奇怪的局面：理论上明显次优的设计，工程上压制了所有理论上更优的设计。直到今天，主流大模型（GPT-4、Claude、Gemini、LLaMA、Qwen）使用的仍是 O(n²) 的 full attention，只是用 FlashAttention 把内存压到了 O(n)。
本篇就来彻底说清楚：复杂度到底是怎么来的，O(n²) 的具体瓶颈在哪里，5 类降复杂度方案各自的代价是什么，以及为什么 FlashAttention 不是 O(n)（这是最常见的误解之一）。读完后你应该能在面试或讨论中清晰回答：「LLaMA 跑 128k 上下文的瓶颈到底在哪？」「为什么 Mamba 这么火但还没替代 Transformer？」
回到第 13 篇推导的 scaled dot-product attention：
Q, K, V ∈ R^{n×d}
S = Q K^T / √d # n×n
A = softmax(S) # n×n
O = A V # n×d逐步算计算量：
Q K^T：两个 n×d 矩阵相乘，得到 n×n
矩阵。FLOPs ≈ 2 · n · n · d = 2n²d。softmax：对 n×n 矩阵每行做 softmax。约
5n²
FLOPs（exp、sum、div），相对前者可忽略。A V：n×n 与 n×d 相乘。FLOPs ≈
2 · n · n · d = 2n²d。总计：4n²d FLOPs。计算复杂度
O(n²·d)。
显存方面：
S 和 A：n×n 矩阵，占
n² 个元素。fp16 每个 2 字节，所以一个 head 是
2n² 字节。显存复杂度 O(n²)——这才是真正卡死长上下文的瓶颈，比计算更要命。
我们做个具体数字：n=128k 时，单层单 head 的 attention 矩阵 = 128k×128k×2 = 32 GB。一个 80GB 的 H100 单卡都装不下一层。这就是为什么没有 FlashAttention 的话，128k 上下文在硬件上完全不可能。
很多人下意识以为 attention 是模型计算量最大的部分。实际上这取决于 n 和 d 的相对大小。
FFN 的计算量：每层 FFN 做两次矩阵乘
n·d → n·4d → n·d，FLOPs ≈
16n·d²。复杂度 O(n·d²)。
Attention 的计算量：FLOPs ≈
4n²·d。复杂度 O(n²·d)。
两者比值：4n²d / 16nd² = n/(4d)。
n < 4d 时，FFN 主导。n > 4d 时，attention 主导。LLaMA-7B 的 d=4096，所以 n=16k 才是 attention 与 FFN 各占一半的临界点。在 n=2k（典型短对话）时，attention 只占总 FLOPs 的约 1/8；FFN 才是大头。这是为什么短上下文场景里 attention 复杂度其实不那么紧迫，FFN 的优化（量化、MoE）反而更值钱。
但训练长上下文（n=32k 或 128k）的模型时，attention 立刻反超：n=128k 时 attention 是 FFN 的 8 倍。这就是为什么长上下文的训练成本爆炸式增长。
n²·d 的计算可以靠堆 GPU 时间解决，但 O(n²) 显存是硬墙。
让我们算几个典型配置：
| 序列长度 n | n² 矩阵 fp16 显存 | 单 H100（80 GB）能放几层？ |
|---|---|---|
| 2k | 8 MB | 一万层都行 |
| 8k | 128 MB | 几百层 |
| 32k | 2 GB | 30 层左右（接近极限） |
| 128k | 32 GB | 2 层（不能训练） |
| 1M | 2 TB | 不可能 |
注意以上是单 head 的数字。实际多 head 会再 ×8 或 ×32（取决于 GQA 配置）。再加上 batch 维度，128k 训练时即便 batch=1 都装不下。
这是为什么 FlashAttention 和 KV cache 是绝对必需品而不是优化项。没有它们，Transformer 在 n>16k 上就不能存在。
推理阶段，KV cache 增加了一项隐藏的复杂度：每生成一个新 token，KV cache 都要存下所有历史的 K、V。
KV cache 大小 =
2 · n · d · L（2 是 K 和 V，L 是层数）。
LLaMA-2-7B 配置：d=4096，L=32。每个 token 的 KV cache =
2 · 4096 · 32 · 2 字节 = 512 KB。
n=4k 时 KV cache = 2 GB；n=128k 时 KV cache = 64 GB。这甚至超过了模型权重本身的大小。
KV cache 还和 batch size 是乘性的：服务 32
个并发用户、每个 8k 上下文，仅 KV cache 就是
32 × 8k × 512KB = 128 GB，必须分多卡。
正因如此，工业界又发明了 GQA（Grouped Query Attention，第 16 篇讨论过）、MLA（Multi-head Latent Attention，DeepSeek-V2 的核心创新）、Paged Attention（vLLM 的核心创新）等专门针对 KV cache 的优化。
工业界和学术界对 O(n²) 的不满几乎是同步爆发的。从 2019 年开始，出现了一系列试图把 attention 复杂度降到 O(n·log n)、O(n·√n) 甚至 O(n) 的工作。我把它们归成五大类：
让每个位置只看前后各 w 个位置。复杂度 O(n·w)，w 是固定常数所以本质是 O(n)。
代表工作：Longformer（2020，Allen AI）、Mistral-7B 的 SWA（sliding window attention）。
优点：实现极简，把 mask 改成带状即可。GPU 友好（因为不需要复杂索引）。
缺点：彻底丢失全局信息。模型无法建立长距离依赖。
补救方案：加 global token（Longformer）、加 dilation（Sparse Transformer）。
按某种规则稀疏化 attention 矩阵。复杂度通常 O(n·√n)。
代表工作：Sparse Transformer（OpenAI 2019，GPT-3 之前）、BigBird（Google 2020）。
优点：理论好，可证明能近似 full attention。
缺点：实现复杂，GPU 上不易高效，效果与 full attention 仍有差距。BigBird 论文发表后实际工业部署很少。
把 softmax 换成可分解的核函数 φ(Q)φ(K)^T，从而 (φ(Q)φ(K)^T)V = φ(Q)(φ(K)^T V)，括号里先算就是 O(n)。
代表工作：Linformer（2020）、Performer（2020）、Linear Attention（2020）。
优点：真正的 O(n)，对长序列友好。
缺点：在大规模语言模型上效果显著弱于 full attention。多数工作只验证了几亿参数规模，没有在 LLaMA 级别复现成功。
不走 attention 路线，用状态空间模型或 gated RNN 实现 O(n) 序列建模。
代表工作：S4（2021）、Mamba（2023）、RWKV（2023）、RetNet（2023）。
优点：真正的 O(n) 推理，并行训练（通过 selective scan 等技巧），实测效果在中小规模上接近 Transformer。
缺点：长程检索任务（needle in haystack）仍弱于 attention；目前最大模型仍未到 70B+；混合架构（Jamba=Mamba+attention）成为务实方案。
这一类不降低渐进复杂度，而是利用 GPU 内存层次（HBM → SRAM）做 tiling 和 recomputation，让显存从 O(n²) 降到 O(n)，计算时间因为减少 HBM 访问而加速 2–4×。
代表工作：FlashAttention（Dao 2022）、FlashAttention-2（2023）、FlashAttention-3（2024）。
优点：精确而非近似，效果与 full attention 完全一致；显存 O(n)；训练加速显著；成为事实标准。
缺点：仍是 O(n²) 计算，长序列计算量本身没降。
我见过太多次这种错误描述：「FlashAttention 把注意力复杂度降到 O(n)」。这是错的。
正确说法是：FlashAttention 把 attention 的显存复杂度从 O(n²) 降到 O(n)，但计算复杂度仍然是 O(n²)。
它的核心想法是：
所以：
这种「不改变渐进复杂度但巨幅改善实测性能」的工作有个专有名词：IO-aware。其本质是认识到现代 GPU 的瓶颈早已不是计算而是带宽。FlashAttention 的成功告诉我们：在系统层面优化往往比改算法更值钱。
2020–2022 年是「线性注意力」的黄金期，论文层出不穷，但到了 2023 年大家发现：主流大模型基本都还在用 full attention。
原因有几个：
质量差距难以弥合。线性 attention 在小规模任务上能逼近 full attention，但放大到 7B+ 时差距被放大。Anthropic、OpenAI、Google 都内部尝试过，都没把线性方案 ship 出去。
FlashAttention 让 full attention 不再是瓶颈。一旦显存问题解决了，full attention 的「贵」就只剩计算贵——而计算贵可以用 GPU 数量解决。
长上下文的真正瓶颈是 KV cache 而不是 attention 计算。这让线性方案的「O(n) 计算」优势失去意义，因为推理时大头根本不在那。
稀疏方案 GPU 不友好。理论上的 O(n·√n) 在 GPU 上往往跑不出来，因为稀疏 indexing 的 memory access pattern 不规则。
混合架构兜底。如果线性方案确实有用，可以做成混合架构（如 Jamba），这种灵活性比纯线性更受工业界欢迎。
所以现在的实际格局是：
「O(n²) 复杂度」常被作为长上下文困难的简单解释。但在 2024 年的现实里，让模型支持 128k 上下文的难点不是复杂度，而是质量。
具体说，当 n 从 4k 扩到 128k 时，遇到的问题包括：
位置编码外推（第 19 篇主题相关，第 27、28 篇会展开 RoPE 与外推）：训练只到 4k，推理到 128k，位置嵌入是怎么处理的？需要 NTK-aware、YaRN 等技巧。
attention 注意力发散：序列变长后 softmax 越来越平，attention 几乎对所有位置都给一点点权重，「focus」消失。这与 attention sink（第 17 篇）现象相关。
训练数据不够：互联网上 128k+ 的连续高质量文本极少。需要拼接、人工合成。
显存爆炸：即便有 FlashAttention，128k 训练时激活值（activation）的显存仍极大，需要 gradient checkpointing、sequence parallelism 等。
质量评估难：传统困惑度在长上下文上不区分。需要 needle-in-haystack、RULER、∞-Bench 等专用评测。
复杂度只是表面问题。真正的长上下文工程涉及训练、数据、评测、推理优化的全链条。
讲到这里必须说清楚一个事实：现代 GPU 的瓶颈早已不是 FLOPs 而是 memory bandwidth。
H100 的 fp16 算力约 989 TFLOPs，HBM 带宽约 3.35 TB/s。FLOPs/byte 的「算术强度」需要达到 ~295 才能让算力跑满。
attention 的算术强度本身不高：每读一字节 K/V，做大约 d 次 FMA。d=4096 时算术强度 = 4096 ≈ 跑满。但是！如果你按朴素方式实现 attention，每次都要把 n×n 的 attention matrix 写回 HBM 再读出来 softmax 再读出来乘 V，这些「memory round trip」会让实际算术强度暴跌到 ~10。
这是 FlashAttention 巨幅加速的真正原因——它并没有让 GPU 算得更快，而是让 GPU 不用频繁来回搬数据。
理解这一点，对所有「为什么我的训练没有跑满 GPU」的疑问都有帮助。答案 99% 是带宽瓶颈，1% 是算力瓶颈。
研究界一个反复被验证的经验法则：任何打着「线性 attention」旗号的方案，质量上都会损失几个点。这种损失在小规模任务上不明显，但在 LLM 上会被放大。
具体表现：
为什么？直觉理解：线性 attention 把 n×n 矩阵分解成两个 n×k 的乘积，丢失了 rank=n 的表达力。这种压缩在「精确建模任意 token 对」的场景下损失大。Mamba 的 selective scan 等机制是在试图弥补这种损失，但完全弥补尚无定论。
所以「O(n)」从来不是免费的午餐。
attention 的复杂度在训练和推理阶段是不同的。
训练（一次前向 + 反向）：
推理（单 token）：
推理（生成 m 个 token）：
注意推理的总计算仍是 n²-级别，因为虽然每步是 O(n)，但要做 n 步。这意味着 chatbot 生成长回复时也是 O(n²)。
让我们做一个有趣的估算：把 LLaMA-7B 从 4k 上下文扩到 32k，计算成本是多少？
LLaMA-7B 训练数据约 2T tokens。如果都按 4k 上下文训练：
现在改成 32k：
整体训练 cost 涨 ~10×，主要由 attention 推动。这就是为什么长上下文模型训练费用高得离谱。
线性 attention 还有个非技术但同样重要的障碍：生态系统。
线性 attention 要替代 full attention，不仅要在质量上证明自己，还要建立一整套等价的工具链。这是 Mamba 直到 2024 年才有 vLLM 支持的原因。生态系统粘性是 full attention 的护城河。
把 PyTorch 朴素实现拿出来对照看：
def naive_attention(Q, K, V, mask=None):
# Q,K,V: (B, H, n, d)
n = Q.size(-2)
d = Q.size(-1)
# 1) 计算 n×n 矩阵：O(n²·d)
scores = (Q @ K.transpose(-2, -1)) / d**0.5 # (B,H,n,n) ← O(n²) 显存
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
# 2) softmax over n
attn = scores.softmax(dim=-1) # (B,H,n,n)
# 3) 加权和：再 O(n²·d)
out = attn @ V # (B,H,n,d)
return out显存大头在 scores、attn 两个
n×n 矩阵。把它们物化是问题所在。
FlashAttention 接口：
from flash_attn import flash_attn_func
out = flash_attn_func(Q, K, V, causal=True)
# 内部分块计算，永远不显式构造 n×n接口几乎一样，性能差几倍到几十倍。这是「IO-aware」工程的魅力。
下一篇（第 19 篇）会回到《Attention Is All You Need》论文本身的历史背景：2017 年的它是怎么从 Google Brain 诞生、怎么在 LSTM 时代杀出重围、为什么作者最初甚至没意识到 decoder-only 的潜力。理解了复杂度的工程现实，再回看那篇论文会有截然不同的感受——「all you need」这句话在当时是个挑衅，七年后却成了预言。
← 上一篇：17. Causal Mask 与自回归 | 下一篇：19.《Attention Is All You Need》论文背景 →
把当前热点继续串成多页阅读，而不是停在单篇消费。
2026-04-15 · transformer
RNN 三难（长程依赖、梯度稳定、训练并行）的系统分析；attention 如何作为补丁逐步把 RNN 推向极限；Vaswani 2017 抛弃循环的范式革命
2026-04-15 · transformer
很多人第一次读 Vaswani 2017 的公式时，都会卡在那一个 √dk 上。
2026-04-15 · transformer
这是【Transformer 与注意力机制】系列的第一篇，承担两件事：一是把这套五十多篇文章为谁写、解决什么问题、彼此之间是什么关系交代清楚；二是为完全没基础的读者画出一条从向量、点积、矩阵乘法走到自注意力、再走到大语言模型的爬升路径，让你在投入时间之前先知道终点在哪、路上要经过哪些坎、读完之后你会、还不会做什么事。
2026-04-15 · transformer
把矩阵乘法掰开成两种等价但风格不同的视角——『行 × 列』的点积视角和『列的线性组合』视角，最终落到 QK^T 的形状分析。
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。