Multi-Head-Attention的时间复杂度

Bert中的self-attention使用的是Multi-Head Attention.

记序列的长度为$n$, 每个token向量对应的hidden size为$d$. 整个Multi-Head Attention的计算过程如下.

生成Q, K, V

Self Attention的输入为一个$(n, d)$大小的矩阵. 这个矩阵需要分别经过三个不同的大小为$(d, d)$的Dense, 生成Q, K, V矩阵, 对应的大小都是$(n, d)$.

做三次$(n, d)$大小的矩阵与$(d, d)$大小的矩阵相乘, 这一步的时间复杂度为$O(nd^2)$.

计算相似度矩阵

普通的Attention中, Q与K两个矩阵相乘, 得到相似度矩阵. 两者的大小都是$(n, d)$, 矩阵相乘得到$(n, n)$大小的相似度矩阵, 对应的时间复杂度为$O(n^2d)$.

但Multi-Head Attention要先把Q, K, V分割为$m$个head, 每个head中的hidden size记为$a=d/m$. 然后每个head内部进行Attention的计算, 不同head之间互不干扰.

因此考虑上head, Q和K矩阵需要转换为$(n, m, a)$大小的tensor, 点积得到$(n, n, m)$大小的相似度矩阵, 因此这一步的时间复杂度为$O(n^2ma)=O(n^2d)$.

Softmax计算

对得到的按head划分的, 大小为$(n, n, m)$的相似度矩阵计算softmax, 这一步的时间复杂度为$O(n^2m)$.

计算加权和

将$(n, n, m)$大小的权值矩阵与$(n, m, a)$大小的O点乘, 得到$O(n, m, a)$大小的输入. 这一步沿着$n$所在的维度进行点乘, 所用的时间复杂度为$O(n^2ma)=O(n^2d)$.

再经过reshape变成最后的输出, 大小仍然为$(n, d)$.

总结

• 生成Q, K, V的时间复杂度为$O(nd^2)$
• 计算相似度矩阵的时间复杂度为$O(n^2md)$
• Softamax计算的时间复杂度为$O(n^2m)$
• 计算加权和的时间复杂度为$O(n^2d)$

整体的时间复杂度为$O(n^2d+nd^2)$, 即序列长度与hidden size更大的一个占主导.