当前位置：网站首页>注意力机制的一种卷积替代方式

注意力机制的一种卷积替代方式

2022-07-06 08:51:00 【cyz0202】

参考自：fairseq

背景

常见的注意力机制如下

$Attention(Q,K,V) = Softmax(\frac{QK^T}{\sqrt{d}})V$

而且一般使用multi heads的方式进行上述计算；

耗时耗力；

卷积很早就用在NLP中，只是常见的使用方法计算量也不小，而且效果不如Attention；

有没有可能改进卷积在NLP的应用方式，达到既快又好的效果？

一个思路就是 depthwise convolution（减少计算量），并且模仿注意力的softmax机制（加权平均）；

改进方案

常见的depthwise convolution如下

$\\O_{i,c} \\= DepthwiseConv(X, W_{c,:}, i, c) \\= \sum^k_{j=1}X_{(i+j-\left \lfloor \frac{k+1}{2} \right \rfloor), c} \cdot W_{c,j} \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ (1)$

W为kernel， $W \in R^{d \times k}$ ，参数量大小为：当d=1024，k=7，d*k=7168

-------

为了进一步减少计算量，可以令 $W \in R^{H \times k}$ ，其中一般 $H \ll d$ ，如H=16，d=1024；

H所在维度变小，为了保持对x做完整的卷积计算，则需要在原来d这个维度重复利用W，也就是不同channel可能用到W中同一行参数；

有两种重复利用方式，一种是W在d这个维度的平移重复，即每H行使用一次W；

另一种可以考虑W的行在d维度上的重复，即卷积计算过程中，x在d这个维度上每d/H行对应到 W中某一行，如x的d维度上的[0-d/H)行 depthwise卷积计算使用 W的第0行参数；（稍微有点绕，可参考下面公式3，会更直观）

这里使用第二种方式，即行的重复；考虑的是模仿Attention的multi heads，每个头（大小为d/H）能独立计算；第二方式相当于x的每d/H行（d维度上的每个头）使用各自独立的kernel参数；

细心的读者可以发现，这里的H其实就是模仿Attention中muliti heads的H；

通过以上设计，W的参数量缩减到 H*k（如H=16，k=7，则参数量只有112）

-------

除了减少计算量，还可以让卷积模仿Attention的softmax过程；

观察公式(1)， $W_{c,j}$ 的作用和Attention softmax输出很像，即加权平均；

所以为什么我们不让 $W_{c,:}$ 也是一个distribution呢？

因此令

$\\W_{c,:} = Softmax(W_{c,:}) \\softmax(W)_{c,j} = \frac{e^{W_{c,j}}}{\sum_{j'} e^{W_{c,j'}}} \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ (2)$

-------

最终卷积计算如下

$\\O_{i,c} \\= DepthwiseConv(X, W_{\left \lceil \frac{cH}{d} \right \rceil,:}, i, c) \\= \sum^k_{j=1}X_{(i+j-\left \lfloor \frac{k+1}{2} \right \rfloor), c} \cdot W_{\left \lceil \frac{cH}{d} \right \rceil,j} \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ (3)$

注：W下角标上取整默认 $c \in [1, d]$ ，也可以换成下取整并令 $c \in [0, d-1]$

具体实现

上述行重复的计算方式并不能直接利用已有的卷积算子计算；

为了利用矩阵计算，可以想一个折中的方式

令 $W => W' \in R^{BH \times n \times n}$ ，令 $x => x' \in R^{BH \times n \times \frac{d}{H}}$ ，执行BMM（Batch MatMul），可得

$\\Out = W' \cdot x' \\ Out \in R^{BH\times n \times \frac{d}{H}} \\ or\\ Out \in R^{B \times n \times d} \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ (4)$

上述设计可以保证结果Out有我们需要的正确的shape，那怎么保证Out值也是正确，即怎么设计W'，x的具体值？

不难，只要根据上一节的令W行重复以对x进行multi-head计算的思想出发即可；

考虑 $x => x' \in R^{BH \times n \times \frac{d}{H}}$ ，就是简单的一个reshape，维度1的BH代表batch*H，假设batch=1，那就是H个计算头；维度2的n代表序列长度n，即n个位置；最后一个维度（即channel）代表的是一个head，即一个计算头大小；

卷积计算时，

上述x'的每个头（最后一个维度大小d/H，共H个），对应的W参数是W中按顺序的一行参数（大小k，H个头就对应W的H行）；同时x'第2个维度表示有n个位置要进行卷积计算，并且每个位置计算方式是与窗口大小为k且固定值的序列加权求和；

所以 W'维度1的BH 代表与x'的BH个头一一对应的卷积参数（B=1时就是H个头对应H行kernel参数）；

W'维度2的n，代表n个位置要进行卷积计算；

W'维度3的n，是对当前头对应的kernel内一行的 k个参数扩展到n个参数，扩展方式是0填充；

之所以要对当前kernel大小为k的行参数填充到n，主要是上述卷积计算是一个固定大小为k的窗口在长度为n的序列上滑动；

因为卷积参数个数恒定，但是位置发生变化，看起来是不规则计算；但是想一下，虽然位置在滑动，但是是在长度为n的序列上滑动；能不能构造一个长度为n的fake kernel，除了窗口所在位置放置真正kerenl参数值，其他非窗口所在的位置都填充0呢？这样我们得到了一个长度固定为n的fake kernel，就能以统一形式计算了；

举例说明，假设此时计算位置为 10，则窗口中心也在10这个位置，而窗口具体位置可能是[7,13]，那么我把除了窗口所在位置外的其他位置都填充0，就得到一个长度为n的参数序列，就能用统一的 n*n 点积计算方式去和 x长度为n的输入进行计算了；

当然，上述填充方式每次得到的填充序列（fake kernel）都是不一样的，因为当前卷积行参数值虽然不变，但是位置在变；

了解过深度学习框架是如何实现卷积计算的同学可能对这种填充方式会更熟悉；

上述说法可能还是有点绕，读者想一下 $R^{n \times n} \times R^{n \times \frac{d}{H}}$ 矩阵相乘，第1个n表示n个位置，第2/3个n表示当前某个位置发生的卷积计算，之所以是n*n，是为了达到统一计算形式进行了 0填充；