阿里面试官问：Self-Attention 的时间复杂度/空间复杂度是怎么计算的?

来源：互联网时间：2026-06-08 14:29:06

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

一位正在备战 2024 年秋招的双一流大学硕士生，在寻找大模型算法岗实习的过程中遇到了不少有意思的面试题。这里记录下其中一道，分享给同样在为 offer 努力的小伙伴们。

面试题

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

标准答案

但凡看过 Transformer 论文的同学都知道，Self-Attention 的复杂度常常被直接标记为 O(N²)，N 就是序列长度。但为什么是这个结果？背后其实是几个矩阵运算的叠加。

先明确一点：复杂度分析是定性而非定量，它告诉我们的是趋势，不是精确数值。接下来把 scaled dot-product attention 的计算过程拆开看看。

为了分析复杂度，把上面的计算拆成这几步：

只要分析清楚每一步的复杂度，总和就出来了。

先看第一个矩阵乘法：

矩阵乘法的朴素算法时间复杂度是 O(m·n·l)。至于空间复杂度，只看存储计算结果，复杂度是 O(m·l)。不过别被这个数字吓到，如果 d=64，存储 Q 和 K 本身比结果更占显存。除非序列很长（N 很大），空间复杂度 O(N²) 才会成为瓶颈。

简单回顾下矩阵乘法：

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]

三个 for 循环，时间复杂度自然是 O(m·n·l)。网上有人用 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 也会得到相同结论，所以时间复杂度确实是 O(m·n·l)。

由于

因此

这一步的时间复杂度为 O(N²·d)。

再来看 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 和加权平均。