我们通过一个具体的例子来深入探讨在注意力机制中使用核函数替代 softmax的计算复杂度变化。
我们将以 N = 8000（序列长度）、d = 1024（嵌入维度）为基础，假设使用核函数将嵌入向量从 1024 维 映射到 2048 维，并计算相应的计算复杂度。

1. 标准自注意力机制的计算复杂度

首先，回顾一下标准 Transformer 的自注意力机制在计算复杂度上的表现。

步骤与计算复杂度

1. 计算查询和键的点积 $Q{K}^{\top }$

   • 操作：矩阵乘法
   • 维度：
     • $Q$（查询矩阵）：$N\times d=8000\times 1024$
     • $K$（键矩阵）：$N\times d=8000\times 1024$
     • $Q{K}^{\top }$：$8000\times 8000$
   • 计算复杂度：$O(N^2\cdot d)=O(8000^2\times 1024)\approx 6.554\times 10^{10}$ 次浮点运算（FLOPs）

2. 应用 softmax 函数

   • 操作：对每一行的 $Q{K}^{\top }$ 应用 softmax
   • 计算复杂度：$O(N^2)=O(8000^2)\approx 6.4\times 10^7$ 次运算（相对较小，可以忽略不计）

3. 计算注意力输出 $\text{softmax}(Q{K}^{\top })V$

   • 操作：矩阵乘法
   • 维度：
     • $\text{softmax}(Q{K}^{\top })$：$8000\times 8000$
     • $V$（值矩阵）：$8000\times 1024$
     • 输出：$8000\times 1024$
   • 计算复杂度：$O(N^2\cdot d)=O(8000^2\times 1024)\approx 6.554\times 10^{10}$ FLOPs

总计算复杂度

• 总复杂度：约 $2\times O(N^2\cdot d)=2\times 6.554\times 10^{10}\approx 1.311\times 10^{11}$ FLOPs

2. 使用核函数替代 softmax 的线性注意力机制

接下来，我们探讨线性注意力（Linear Attention）如何通过核函数降低计算复杂度。

线性注意力的基本步骤

1. 应用核函数 $\varphi (\cdot )$ 对查询 $Q$ 和键 $K$ 进行映射

   • 操作：特征映射（如通过非线性激活函数）
   • 维度：
     • $Q$：$8000\times 1024$ → $\varphi (Q)$：$8000\times 2048$
     • $K$：$8000\times 1024$ → $\varphi (K)$：$8000\times 2048$
   • 计算复杂度：每个映射 $O(N\cdot d\cdot d^{\prime })=O(8000\times 1024\times 2048)\approx 1.679\times 10^{10}$ FLOPs（对 $Q$ 和 $K$ 总共约 $3.358\times 10^{10}$ FLOPs）

2. 计算 $\varphi (K{)}^{\top }V$

   • 操作：矩阵乘法
   • 维度：
     • $\varphi (K{)}^{\top }$：$2048\times 8000$
     • $V$：$8000\times 1024$
     • 结果：$2048\times 1024$
   • 计算复杂度：$O(N\cdot d^{\prime }\cdot d_v)=O(8000\times 2048\times 1024)\approx 1.679\times 10^{10}$ FLOPs

3. 计算 $\varphi (Q)\times (\varphi (K{)}^{\top }V)$

   • 操作：矩阵乘法
   • 维度：
     • $\varphi (Q)$：$8000\times 2048$
     • $\varphi (K{)}^{\top }V$：$2048\times 1024$
     • 结果：$8000\times 1024$
   • 计算复杂度：$O(N\cdot d^{\prime }\cdot d_v)=O(8000\times 2048\times 1024)\approx 1.679\times 10^{10}$ FLOPs

4. 计算归一化项 $\varphi (Q)\times (\varphi (K{)}^{\top }1)$

   • 操作：矩阵乘法
   • 维度：
     • $\varphi (Q)$：$8000\times 2048$
     • $\varphi (K{)}^{\top }1$：$2048\times 1$
     • 结果：$8000\times 1$
   • 计算复杂度：$O(N\cdot d^{\prime })=O(8000\times 2048)\approx 1.638\times 10^7$ FLOPs（可忽略）

5. 最终计算注意力输出

   • 操作：逐元素除法（ broadcasting division ）
   • 维度：$8000\times 1024$ / $8000\times 1$ → $8000\times 1024$
   • 计算复杂度：$O(N\cdot d_v)=O(8000\times 1024)\approx 8.192\times 10^6$ FLOPs（可忽略）

总计算复杂度

• 总复杂度：
  • 特征映射：约 $3.358\times 10^{10}$ FLOPs
  • 矩阵乘法：约 $3.358\times 10^{10}$ FLOPs
  • 总计：约 $6.716\times 10^{10}$ FLOPs

3. 比较与分析

标准自注意力

• 计算复杂度：约 $1.311\times 10^{11}$ FLOPs
• 特点：
  • 计算复杂度与序列长度 $N$ 的平方成正比 $O(N^2\cdot d)$
  • 高计算和内存消耗，难以处理超长序列

线性注意力（使用核函数）

• 计算复杂度：约 $6.716\times 10^{10}$ FLOPs
• 特点：
  • 计算复杂度与序列长度 $N$ 成线性关系 $O(N\cdot d\cdot d^{\prime })$
  • 由于 $d^{\prime }=2d$（升维），相对于降低了二次增长为线性增长，但实际计算量仍然增加
  • 但与标准自注意力相比，整体复杂度 降低了约80%（$6.716\times 10^{10}$ FLOPs vs $1.311\times 10^{11}$ FLOPs）

具体数值差异

• 序列长度 $N=8000$
• 原始嵌入维度 $d=1024$
• 映射后维度 $d^{\prime }=2048$

标准自注意力的计算量约为 $6.554\times 10^{10}$ 次乘加操作（乘法与加法合并计算），需要执行两次这样的操作，总计算负担非常高。

线性注意力通过核函数映射后，总计算复杂度约为 $6.716\times 10^{10}$ FLOPs，相比标准自注意力的 $1.311\times 10^{11}$ FLOPs，计算量减少了约一半（具体减少比例取决于 $d^{\prime }$ 相对于 $d$ 的增长）。

4. 为什么线性注意力仍然可行？

即使核函数映射后维度增加至2048，线性注意力仍然能显著降低计算复杂度，主要原因有：

1. 分解计算：

   • 标准自注意力需要提前计算 $Q{K}^{\top }$（$O(N^2\cdot d)$），而线性注意力通过分解将二次复杂度转换为线性复杂度（$O(N\cdot d\cdot d^{\prime })$）。

2. 特征映射的选择：

   • 核函数 $\varphi (\cdot )$ 的设计确保了特征映射后的空间能够近似原始自注意力的效果，同时分解后的计算仍然保留了重要的依赖关系。

3. 优化与硬件加速：

   • 矩阵乘法和特征映射操作可以高度优化和并行化，利用现代硬件（如GPU）的优势，进一步提升计算效率。

5. 实际应用中的权衡

虽然线性注意力显著降低了计算复杂度，但仍有一些权衡需要考虑：

1. 性能与精度：

   • 线性注意力通过近似替代 softmax，可能会在某些任务上略微牺牲一定的精度。但在许多实际应用中，这种牺牲是可以接受的，尤其是当需要处理超长序列时。

2. 内存消耗：

   • 线性注意力避免了存储 $N\times N$ 的注意力矩阵，显著减少了内存消耗，适合处理更长的序列。

3. 映射维度 $d^{\prime }$：

   • 虽然映射后维度 $d^{\prime }=2048$ 较高，但相对于标准自注意力的 $O(N^2\cdot d)$ 复杂度，线性注意力的 $O(N\cdot d\cdot d^{\prime })$ 仍然具有显著优势。

6. 直观类比

为了更直观地理解，让我们用一个日常生活中的类比：

• 标准自注意力：

  • 类比：想象你在一个会议室，里面有8000个人。你要向每个人发表讲话，并与每个人一对一地交流，了解他们的观点。这样，你需要进行大量的交流和互动，每个人之间的交流数量是平方级别的。

• 线性注意力：

  • 类比：你使用一种智能筛选器（核函数）来快速总结每个人的观点，将他们的意见简化为更紧凑的形式（从1024维升维到2048维），然后仅基于这些总结信息进行汇总和决策。这样，你不需要与每个人逐一详细交流，而是通过简化后的信息快速做出决策，大大减少了交流的次数和复杂度。

7. 总结

通过上述例子和解释，我们可以清楚地看到 线性注意力（通过核函数替代 softmax）的计算复杂度如何从 $O(N^2\cdot d)$ 降低到 $O(N\cdot d\cdot d^{\prime })$，**即使在嵌入维度升高的情况下，整体计算复杂度仍然大幅降低。**例如：

• 标准自注意力：

  • 序列长度 $N=8000$
  • 嵌入维度 $d=1024$
  • 总计算复杂度：$1.311\times 10^{11}$ FLOPs

• 线性注意力（核函数映射至2048维）：

  • 序列长度 $N=8000$
  • 映射后维度 $d^{\prime }=2048$
  • 总计算复杂度：$6.716\times 10^{10}$ FLOPs

这表明线性注意力在处理长序列时，尽管映射后维度增加，但整体计算复杂度仍然大幅低于标准自注意力，从而显著提升了模型的可扩展性和处理效率。