Transformer/CNN/RNN 时间复杂度的对比

一.矩阵、张量乘法的时间复杂度

一个形状为  的矩阵，与另一个形状为  的矩阵相乘，其运算复杂度来源于乘法操作的次数，时间复杂度为 

二. 具体网络

Self-Attention

• 
• 相似度计算  ：  与  运算，得到  矩阵，复杂度为 
• softmax计算：对每行做softmax，复杂度为  ，则n行的复杂度为 
• 加权和：  与  运算，得到  矩阵，复杂度为 

故最后self-attention的时间复杂度为 

对于受限的self-attention，每个元素仅能和周围  个元素进行交互，即和  个  维向量做内积运算，复杂度为  ，则  个元素的总时间复杂度为 

Multi-Head Attention

对于multi-head attention，假设有  个head，这里  是一个常数，对于每个head，首先需要把三个矩阵分别映射到  维度。这里考虑一种简化情况：  。(对于dot-attention计算方式，  与  可以不同)。

• 输入线性映射的复杂度：  与  运算，忽略常系数，复杂度为  。
• Attention操作复杂度：主要在相似度计算及加权和的开销上，  与  运算，复杂度为 
• 输出线性映射的复杂度：concat操作拼起来形成  的矩阵，然后经过输出线性映射，保证输入输出相同，所以是  与  计算，复杂度为 

故最后的复杂度为： 

注意：多头的计算并不是通过循环完成的，而是通过 transposes and reshapes，用矩阵乘法来完成的。假设有  个head，则新的representation dimension：  。因为，我们将  的矩阵拆为  的张量，再利用转置操作转为  的张量。故  的计算为：  与  做计算，得到  的张量，复杂度为  ，即  。注意，此处  实际是一个常数，故  复杂度为  。

Recurrent

•  ：  与  运算，复杂度为  ，  为input size
•  ：  与  运算，复杂度为 

故一次操作的时间复杂度为  ，  次序列操作后的总时间复杂度为 

Convolution

注: 这里保证输入输出都是一样的，即均是 

• 为了保证输入和输出在第一个维度都相同，故需要对输入进行padding操作，因为这里kernel size为  ，（实际kernel的形状为  ）如果不padding的话，那么输出的第一个维度为  ，因为这里stride是为1的。为了保证输入输出相同，则需要对序列的前后分别padding长度为  。
• 大小为  的卷积核一次运算的复杂度为：  ，一共做了  次，故复杂度为 
• 为了保证第二个维度在第二个维度都相同，故需要  个卷积核，所以卷积操作总的时间复杂度为 

三.参考

• Transformer/CNN/RNN的对比（时间复杂度，序列操作数，最大路径长度）
• 矩阵、张量乘法（numpy.tensordot）的时间复杂度分析