NLP 经典面试题————Self-Attention 的时间复杂度/空间复杂度是怎么计算的?

本人是某双一流大学硕士生，也最近刚好准备参加 2024年秋招，在找大模型算法岗实习中，遇到了很多有意思的面试，所以将这些面试题记录下来，并分享给那些和我一样在为一份满意的offer努力着的小伙伴们！！！

面试题

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

标准答案

阅读 Transformer 相关的论文，在讨论 self-attention 的时间和空间复杂度时，都会提到是 O(N^2)，其中 N 是序列长度。

关于时间复杂度(time complexity)或空间复杂度 O(·)，首先要知道这只是一种定性分析，而不是精确的定量分析。

我们来看下 scaled dot-product attention的时间和空间复杂度：

为了分析时间和时间复杂度，我们把上面的计算过程拆分：

只要我们分析出每个计算的复杂度，就可以得到整体计算的复杂度。

我们先来看第一个矩阵乘法

矩阵乘法的朴素算法时间复杂度是

至于空间复杂度，只看存储 计算结果，复杂度是 ，但是也不要觉得这个数字很大，如果 ，其实存储 Q 和 K 要比 更占内(显)存。除非是序列很长 ,空间复杂度 才会是瓶颈。

简单回顾下矩阵乘法

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 个 for 循环，因此矩阵乘法时间复杂度 。

我在网上查找矩阵乘法时间复杂度分析的资料时，发现很多人喜欢用numpy.dot+画图的方式来直观展示，很有趣：

m = 64
n = 64
l = 64

times = []
ms = []
for i in range(20):
ms.append(m)
begin = time.time()
m1 = np.random.random((m, n))
m2 = np.random.random((n, l))
times.append(time.time() - begin)
m *= 2# 改变 m 的大小, 同理可以改变 n 或 l 的大小

# 画图
fig, ax = plt.subplots()
ax.set_ylabel('Time') 
ax.set_xlabel('Array Dimension Size') 
ax.plot(ms, times) 
plt.show()

可以看到，矩阵乘法时间和其中 m 维度大小成正相关，斜率 ~1，如果改变 n 或 l 也会得到相同的结论，因此矩阵乘法时间复杂度是 。

由于

因此

的时间复杂度为

我们再来看 softmax 时间复杂度，假设 z 是一维向量:

def softmax(x):
m_val = max(x)
x = [i-m_val for i in x]
x = [math.exp(i) for i in x]
deno = sum(x)
return [item / deno for item in x]

softmax([1,2,3])# [0.0900, 0.2447, 0.6652]

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

• step1: 相似度计算可以看作大小为(n,d)和(d,n)的两个矩阵相乘：，得到一个 的矩阵.
• step2: softmax就是直接计算了，时间复杂度为
• step3: 加权平均可以看作大小为 和 的两个矩阵相乘： ，得到一个 的矩阵

因此，Self-Attention的时间复杂度是 。

这样，整个

的时间复杂度是：

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

再来看下空间复杂度，不论是存储

最后存储

的空间复杂度是

这样，整个空间复杂度可以看作：

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