根本原因是标准自注意力必须显式构建O(n²)复杂度的Q@K.T注意力分数矩阵，导致内存随序列长度平方增长；线性注意力通过可分解核函数ϕ(Q)@ϕ(K).T替代，将关键中间张量降为与n无关的形状，从而规避该瓶颈。
Self-attention 在 Python 中耗内存，根本原因不是代码写得不好，而是它必须显式构建 Q @ K.T 这个形状为 (batch_size, num_heads, seq_len, seq_len) 的注意力分数矩阵——这个矩阵的内存占用随序列长度平方增长，躲不掉。
为什么 Q @ K.T 矩阵是内存瓶颈
不管用 PyTorch、JAX 还是自定义 NumPy 实现，只要走标准缩放点积注意力（softmax(Q @ K.T / sqrt(d_k)) @ V），就必须在 GPU 或 CPU 上临时存下完整的 seq_len × seq_len 矩阵。哪怕你只关心 top-k 个最相关的位置，框架默认仍会算全。
- 序列长 1024、batch=8、head=12、dtype=float16 → 单层单头就要约
8 × 12 × 1024 × 1024 × 2字节 ≈ 200MB；12 头叠加就是 2.4GB - 模型有 12 层？那仅这一项就占满一张 24GB A100 的一半以上显存
- 反向传播时还要缓存这个矩阵用于梯度计算，实际峰值显存通常是前向的 1.5–2 倍
线性注意力（Linear Attention）怎么绕过这个矩阵
线性注意力不直接算 Q @ K.T，而是把 softmax 拆开，用一个可分解的核函数近似：把 softmax(QK.T) 替换成 ϕ(Q) @ ϕ(K).T，其中 ϕ 是一个特征映射（比如随机傅里叶特征、ELU+1 或 softmax 本身）。这样矩阵乘法顺序就变成：
ϕ(Q) ∈ (b, h, n, d') → 先和 V 合并：ϕ(Q) @ (ϕ(K).T @ V)
关键点在于：ϕ(K).T @ V 形状是 (b, h, d', d_v)，与序列长度 n 无关；整个过程避免了 n×n 中间张量。
- 典型实现如
linformer、performer、flash-attn（部分模式）都依赖这个 trick - 但要注意：
ϕ的选择直接影响近似质量；ELU+1 在长尾分布上比 softmax 更稳，但对短序列可能略降精度 - PyTorch 中若用
torch.einsum手写，务必确认ϕ输出维度d'不设太高（一般 64–128），否则反而更耗显存
你在 Python 里调 nn.MultiheadAttention 却没省内存？检查这三点
PyTorch 官方 nn.MultiheadAttention 默认仍是 full attention，所谓“支持线性 attention”只是接口预留，不等于自动启用。
- 确认你没误以为设置了
batch_first=True就触发了优化——它只改输入格式，不改计算逻辑 - 检查是否手动替换了
attn_mask或用了need_weights=False：后者能省一点输出缓存，但Q @ K.T矩阵照常分配 - 如果你用的是 Hugging Face 的
transformers库，注意model.config.attention_implementation = "eager"（默认）≠ 线性；要切到"sdpa"或第三方后端（如flash_attn）才可能启用 kernel 优化
别被“O(n) 复杂度”误导：线性 attention 的真实代价
理论复杂度降为 O(n)，但实际显存节省受三个现实因素压制：
- 特征映射
ϕ引入额外中间张量，比如ϕ(Q)和ϕ(K)各占b×h×n×d'，当d'接近d_k时，省下的n²可能被新张量吃掉大半 - CUDA kernel 未针对稀疏或低秩结构优化时，
ϕ(Q) @ (ϕ(K).T @ V)的访存带宽压力可能比原版更高 - 多数线性变体在
seq_len 时反而更慢、显存略增——它们专为 2k+ 长度设计，短序列别硬套
真正卡住你的，往往不是公式推导，而是没意识到：线性 attention 不是开关，而是一组需协同调整的组件——特征维度、数值稳定性处理、kernel 实现路径，漏掉任一环，都可能白忙活一场。
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
