gongel

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

一个形状为 $N\times M$ 的矩阵，与另一个形状为 $M\times P$ 的矩阵相乘，其运算复杂度来源于乘法操作的次数，时间复杂度为 $\mathcal{O}(NMP)$

二. 具体网络

Self-Attention

$A(Q,K,V)=\mathrm{Softmax}(QK^T)V \\$

$Q,K,V:n\times d$
相似度计算 $QK^T$ ： $n\times d$ 与 $d\times n$ 运算，得到 $n\times n$ 矩阵，复杂度为 $\mathcal{O}(n^2d)$
softmax计算：对每行做softmax，复杂度为 $\mathcal{O}(n)$ ，则n行的复杂度为 $\mathcal{O}(n^2)$
加权和： $n\times n$ 与 $n\times d$ 运算，得到 $n\times d$ 矩阵，复杂度为 $\mathcal{O}(n^2d)$

故最后self-attention的时间复杂度为 $\mathcal{O}(n^2d)$

对于受限的self-attention，每个元素仅能和周围 $r$ 个元素进行交互，即和 $r$ 个 $d$ 维向量做内积运算，复杂度为 $\mathcal{O}(rd)$ ，则 $n$ 个元素的总时间复杂度为 $\mathcal{O(rnd)}$

Multi-Head Attention

$\mathrm{MultiHead}(Q,K,V)=\mathrm{Concat(head_1,...,head_h)}W^O \\ \mathrm{where\quad head_i}=A(QW_i^Q,KW_i^K,VW_i^V)\\$

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

输入线性映射的复杂度： $n \times d$ 与 $d\times \frac{d}{h}$ 运算，忽略常系数，复杂度为 $\mathcal{O}(nd^2)$ 。
Attention操作复杂度：主要在相似度计算及加权和的开销上， $n\times \frac{d}{h}$ 与 $\frac{d}{h}\times {n}$ 运算，复杂度为 $\mathcal{O}(n^2d)$ $\mathcal{}$
输出线性映射的复杂度：concat操作拼起来形成 $n\times d$ 的矩阵，然后经过输出线性映射，保证输入输出相同，所以是 $n\times d$ 与 $d\times d$ 计算，复杂度为 $\mathcal{O}(nd^2)$

故最后的复杂度为： $\mathcal{O}(n^2d+nd^2)$

注意：多头的计算并不是通过循环完成的，而是通过 transposes and reshapes，用矩阵乘法来完成的。假设有 $h$ 个head，则新的representation dimension： $m=\frac{d}{h}$ 。因为，我们将 $n\times d$ 的矩阵拆为 $n\times h\times m$ 的张量，再利用转置操作转为 $h\times n \times m$ 的张量。故 $QK^T$ 的计算为： $h\times n\times m$ 与 $h\times m \times n$ 做计算，得到 $h\times n\times n$ 的张量，复杂度为 $\mathcal{O}(h^2n^2m)$ ，即 $\mathcal{O}(n^2dh)$ 。注意，此处 $h$ 实际是一个常数，故 $QK^T$ 复杂度为 $\mathcal{O}(n^2d)$ 。

Recurrent

$h_t=f(Ux_t+Wh_{t-1}) \\$

$Ux_t$ ： $d\times m$ 与 $m\times 1$ 运算，复杂度为 $\mathcal{O}(md)$ ， $m$ 为input size
$Wh_{t-1}$ ： $d\times d$ 与 $d\times 1$ 运算，复杂度为 $\mathcal{O}(d^2)$

故一次操作的时间复杂度为 $\mathcal{O}(d^2)$ ， $n$ 次序列操作后的总时间复杂度为 $\mathcal{O}(nd^2)$

Convolution

注: 这里保证输入输出都是一样的，即均是 $n\times d$

为了保证输入和输出在第一个维度都相同，故需要对输入进行padding操作，因为这里kernel size为 $k$ ，（实际kernel的形状为 $k\times d$ ）如果不padding的话，那么输出的第一个维度为 $n-k+1$ ，因为这里stride是为1的。为了保证输入输出相同，则需要对序列的前后分别padding长度为 $(k-1)/2$ 。
大小为 $k\times d$ 的卷积核一次运算的复杂度为： $\mathcal{O}(kd)$ ，一共做了 $n$ 次，故复杂度为 $\mathcal{O}(nkd)$
为了保证第二个维度在第二个维度都相同，故需要 $d$ 个卷积核，所以卷积操作总的时间复杂度为 $\mathcal{O}(nkd^2)$