大模型 | Self-Attention 的时间复杂度/空间复杂度是怎么计算的?

Self-Attention时间复杂度：𝑂(𝑛^2⋅𝑑) ，这里，n是序列的长度，d是embedding的维度。

一、回顾矩阵乘法

C = np.zeros((m, l))

for i in range(m):
  for j in range(l):
    for k in range(n):
      C[i][j] += A[i][k] * B[k][j]

显而易见，3 个𝑓𝑜𝑟 循环，因此矩阵乘法时间复杂度𝑂(𝑚𝑛𝑙)。

二、具体步骤

Self-Attention包括三个步骤：相似度计算，softmax和加权平均

step1: 相似度计算可以看作大小为(n,d)和(d,n)的两个矩阵相乘：(𝑛,𝑑)∗(𝑑,𝑛)=𝑂(𝑛^2⋅𝑑) ，得到一个(n,n)的矩阵.

step2: softmax就是直接计算了，时间复杂度为 𝑂(𝑛^2)

step3: 加权平均可以看作大小为(n,n)和(n,d)的两个矩阵相乘：(𝑛,𝑛)∗(𝑛,𝑑)=𝑂(𝑛^2⋅𝑑) ，得到一个(n,d)的矩阵

因此，Self-Attention的时间复杂度是 𝑂(𝑛^2⋅𝑑) 。

再分析一下Multi-Head Attention，它的作用类似于CNN中的多核。

Transformer/BERT中把 d ，也就是hidden_size/embedding_size这个维度做了reshape拆分，可以去看Google的TF源码或者上面的pytorch源码：

hidden_size (d) = num_attention_heads (m) * attention_head_size (a)，也即 d=m*a

并将 num_attention_heads 维度transpose到前面，使得Q和K的维度都是(m,n,a)，这里不考虑batch维度。

这样点积可以看作大小为(m,n,a)和(m,a,n)的两个张量相乘，得到一个(m,n,n)的矩阵，其实就相当于(n,a)和(a,n)的两个矩阵相乘，做了m次，时间复杂度是 𝑂(𝑛^{2⋅𝑚⋅𝑎)=𝑂(𝑛}2⋅𝑑) 。

因此Multi-Head Attention时间复杂度也是 𝑂(𝑛^2⋅𝑑) ，复杂度相较单头并没有变化，主要还是transposes and reshapes 的操作，相当于把一个大矩阵相乘变成了多个小矩阵的相乘。

三、空间复杂度

存储𝑄𝐾𝑇, 是𝑂(𝑁2).

假设𝑆=𝑠𝑜𝑓𝑡𝑚𝑎𝑥(𝑄𝐾𝑇/𝑑^0.5)∈𝑅𝑁×𝑁,𝑉∈𝑅𝑁×𝑑

存储𝑆𝑉∈𝑅𝑁×𝑑的空间复杂度是𝑂(𝑁𝑑)。

这样，整个𝑠𝑜𝑓𝑡𝑚𝑎𝑥(𝑄𝐾𝑇/𝑑^0.5)𝑉的空间复杂度可以看作：𝑂(𝑁2+𝑁𝑑)

如果把向量维度𝑑看作常数，则可以说 self-attention 的空间复杂度是序列长度的平方。

---

最后分享

学习AI大模型是一个系统的过程，需要从基础开始，逐步深入到更高级的技术。

这里给大家精心整理了一份全面的AI大模型学习资源，包括：AI大模型全套学习路线图（从入门到实战）、精品AI大模型学习书籍手册、视频教程、实战学习、面试题等，资料免费分享！

1. 成长路线图&学习规划

要学习一门新的技术，作为新手一定要先学习成长路线图，方向不对，努力白费。

这里，我们为新手和想要进一步提升的专业人士准备了一份详细的学习成长路线图和规划。可以说是最科学最系统的学习成长路线。

2. 大模型经典PDF书籍

书籍和学习文档资料是学习大模型过程中必不可少的，我们精选了一系列深入探讨大模型技术的书籍和学习文档，它们由领域内的顶尖专家撰写，内容全面、深入、详尽，为你学习大模型提供坚实的理论基础。（书籍含电子版PDF）

3. 大模型视频教程

对于很多自学或者没有基础的同学来说，书籍这些纯文字类的学习教材会觉得比较晦涩难以理解，因此，我们提供了丰富的大模型视频教程，以动态、形象的方式展示技术概念，帮助你更快、更轻松地掌握核心知识。

4. 大模型项目实战

学以致用 ，当你的理论知识积累到一定程度，就需要通过项目实战，在实际操作中检验和巩固你所学到的知识，同时为你找工作和职业发展打下坚实的基础。

5. 大模型面试题

面试不仅是技术的较量，更需要充分的准备。

在你已经掌握了大模型技术之后，就需要开始准备面试，我们将提供精心整理的大模型面试题库，涵盖当前面试中可能遇到的各种技术问题，让你在面试中游刃有余。

全套的AI大模型学习资源已经整理打包，有需要的小伙伴可以微信扫描下方优快云官方认证二维码，免费领取【保证100%免费】

如有侵权，请联系删除